segment_image.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  1. #!/usr/bin/python3
  2. # Copyright (C) 2009-2016, Panagiotis Christopoulos Charitos and contributors.
  3. # All rights reserved.
  4. # Code licensed under the BSD License.
  5. # http://www.anki3d.org/LICENSE
  6. import optparse
  7. import sys
  8. import struct
  9. class Config:
  10. in_file = ""
  11. out_dir = ""
  12. def printi(s):
  13. print("[I] %s" % s)
  14. def parse_commandline():
  15. """ Parse the command line arguments """
  16. parser = optparse.OptionParser(usage = "usage: %prog [options]", \
  17. description = "This program takes a single 2D image and spits "\
  18. "a number of images. Input and output images should be TGA.")
  19. parser.add_option("-i", "--input", dest = "inp",
  20. type = "string", help = "specify the image to split.")
  21. parser.add_option("-o", "--output", dest = "out",
  22. type = "string", help = "the directory of the output images.")
  23. parser.add_option("-s", "--size", dest = "size",
  24. type = "string", metavar="WxH",
  25. help = "size of the splits.")
  26. # Add the default value on each option when printing help
  27. for option in parser.option_list:
  28. if option.default != ("NO", "DEFAULT"):
  29. option.help += (" " if option.help else "") + "[default: %default]"
  30. (options, args) = parser.parse_args()
  31. if not options.inp or not options.out or not options.size:
  32. parser.error("argument is missing")
  33. # Parse the --size
  34. size = [0, 0]
  35. try:
  36. size_strs = options.size.split("x")
  37. size[0] = int(size_strs[0])
  38. size[1] = int(size_strs[1])
  39. except:
  40. parser.error("incorrect --size: %s" % sys.exc_info()[0])
  41. config = Config()
  42. config.in_file = options.inp
  43. config.out_dir = options.out
  44. config.size = size
  45. return config
  46. def split(filename, split_size, out_dir):
  47. """ Load an image """
  48. # Read and check the header
  49. uncompressed_tga_header = struct.pack("BBBBBBBBBBBB", \
  50. 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0)
  51. in_file = open(filename, "rb")
  52. tga_header = in_file.read(12)
  53. if len(tga_header) != 12:
  54. raise Exception("Failed reading TGA header")
  55. if uncompressed_tga_header != tga_header:
  56. raise Exception("Incorrect TGA header")
  57. # Read the size and bpp
  58. header6_buff = in_file.read(6)
  59. if len(header6_buff) != 6:
  60. raise Exception("Failed reading TGA header #2")
  61. header6 = struct.unpack("BBBBBB", header6_buff)
  62. img_width = header6[1] * 256 + header6[0]
  63. img_height = header6[3] * 256 + header6[2]
  64. img_bpp = header6[4];
  65. if img_bpp != 24 and img_bpp != 32:
  66. raise Exception("Unexpected bpp")
  67. # Check split size against the image
  68. if (img_width % split_size[0]) != 0 or (img_height % split_size[1]) != 0:
  69. raise Exception("Sizes of the input image and the split are not " \
  70. "compatible: %d %d vs %d %d"
  71. % (img_width, img_height, split_size[0], split_size[1]))
  72. # Dump the data to an array
  73. pixels = []
  74. for y in range(0, img_height):
  75. pixels.append([])
  76. for x in range(0, img_width):
  77. pixels[y].append([])
  78. if img_bpp == 24:
  79. pixel = in_file.read(3)
  80. pixels[y][x].append(pixel[0])
  81. pixels[y][x].append(pixel[1])
  82. pixels[y][x].append(pixel[2])
  83. else:
  84. pixel = in_file.read(4)
  85. pixels[y][x].append(pixel[0])
  86. pixels[y][x].append(pixel[1])
  87. pixels[y][x].append(pixel[2])
  88. pixels[y][x].append(pixel[3])
  89. # Iterate splits and write them
  90. split_count_x = int(img_width / split_size[0])
  91. split_count_y = int(img_height / split_size[1])
  92. count = 0
  93. for y in range(0, split_count_y):
  94. for x in range(0, split_count_x):
  95. # Open file and write header
  96. out_file = open("%s/out_%02d.tga" % (out_dir, count), "wb")
  97. out_file.write(uncompressed_tga_header)
  98. header2 = struct.pack("BBBBBB", split_size[0] % 256,
  99. int(split_size[0] / 256), split_size[1] % 256,
  100. int(split_size[1] / 256), img_bpp, 0)
  101. out_file.write(header2)
  102. # Iterate split pixels
  103. for sy in range(0, split_size[1]):
  104. for sx in range(0, split_size[0]):
  105. in_x = x * split_size[0] + sx
  106. in_y = y * split_size[1] + sy
  107. if(img_bpp == 24):
  108. pixels_s = struct.pack("BBB", pixels[in_y][in_x][0],
  109. pixels[in_y][in_x][1], pixels[in_y][in_x][2])
  110. else:
  111. pixels_s = struct.pack("BBBB", pixels[in_y][in_x][0],
  112. pixels[in_x][in_y][1], pixels[in_y][in_x][2],
  113. pixels[in_y][in_x][3])
  114. out_file.write(pixels_s)
  115. out_file.close()
  116. count += 1
  117. def main():
  118. """ The main """
  119. config = parse_commandline();
  120. split(config.in_file, config.size, config.out_dir)
  121. printi("Done!")
  122. if __name__ == "__main__":
  123. main()