extract_parts.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. # main import
  2. import os
  3. import argparse
  4. import numpy as np
  5. # images imports
  6. from PIL import Image
  7. image_block = 800, 800
  8. zone_block = 200, 200
  9. def get_parts(img):
  10. parts = []
  11. h, w, c = img.shape
  12. h_block, w_block = image_block
  13. h_zone, w_zone = zone_block
  14. m_h = int(h/2)
  15. m_w = int(w/2)
  16. m_h_img_block = int(h_block/2)
  17. m_w_img_block = int(w_block/2)
  18. if h > h_block and w > w_block:
  19. # extract left up image
  20. h_start = int((h % h_zone) / 2)
  21. w_start = int((w % w_zone) / 2)
  22. left_up_image = img[h_start:h_block+h_start, w_start:w_start+w_block]
  23. parts.append(left_up_image)
  24. # extract middle up image
  25. h_start = int((h % h_zone) / 2)
  26. middle_up_image = img[h_start:h_block+h_start, m_w-m_w_img_block:m_w+m_w_img_block]
  27. parts.append(middle_up_image)
  28. # extract right up image
  29. h_start = int((h % h_zone) / 2)
  30. w_end = w - int((w % w_zone) / 2)
  31. right_up_image = img[h_start:h_block+h_start, w_end-w_block:w_end]
  32. parts.append(right_up_image)
  33. # extract left bottom image
  34. h_end = h - int((h % h_zone) / 2)
  35. h_start = h_end - h_block
  36. w_start = int((w % w_zone) / 2)
  37. left_bottom_image = img[h_start:h_end, w_start:w_start+w_block]
  38. parts.append(left_bottom_image)
  39. # extract middle bottom image
  40. h_end = h - int((h % h_zone) / 2)
  41. h_start = h_end - h_block
  42. middle_bottom_image = img[h_start:h_end, m_w-m_w_img_block:m_w+m_w_img_block]
  43. parts.append(middle_bottom_image)
  44. # extract left bottom image
  45. h_end = h - int((h % h_zone) / 2)
  46. h_start = h_end - h_block
  47. w_end = w - int((w % w_zone) / 2)
  48. right_bottom_image = img[h_start:h_end, w_end-w_block:w_end]
  49. parts.append(right_bottom_image)
  50. # extract middle center image
  51. middle_image = img[m_h-m_h_img_block:m_h+m_h_img_block, m_w-m_w_img_block:m_w+m_w_img_block]
  52. parts.append(middle_image)
  53. return parts
  54. def extract(folder, output):
  55. images_path = sorted(os.listdir(folder))
  56. check_exists = False
  57. # for each image get sub parts
  58. for img_index, img_name in enumerate(images_path):
  59. img_path = os.path.join(folder, img_name)
  60. img = np.array(Image.open(img_path))
  61. # get all expected parts from image
  62. parts = get_parts(img)
  63. if len(parts) > 1:
  64. for index, part in enumerate(parts):
  65. # get part output folder and create it if not exists
  66. prefix_image_folder_path = os.path.join(output, folder.split('/')[-1].replace('/', '') + '_part' + str(index))
  67. if not os.path.exists(prefix_image_folder_path):
  68. os.makedirs(prefix_image_folder_path)
  69. elif img_index == 0:
  70. check_exists = True
  71. if not check_exists:
  72. # build output image path and save it
  73. image_output_folder_path = os.path.join(prefix_image_folder_path, img_name)
  74. Image.fromarray(part).save(image_output_folder_path)
  75. else:
  76. # only middle of the scene
  77. part = parts[0]
  78. prefix_image_folder_path = os.path.join(output, folder.split('/')[-1].replace('/', '') + '_part6')
  79. if not os.path.exists(prefix_image_folder_path):
  80. os.makedirs(prefix_image_folder_path)
  81. elif img_index == 0:
  82. check_exists = True
  83. if not check_exists:
  84. # build output image path and save it
  85. image_output_folder_path = os.path.join(prefix_image_folder_path, img_name)
  86. Image.fromarray(part).save(image_output_folder_path)
  87. if check_exists:
  88. break
  89. def main():
  90. parser = argparse.ArgumentParser(description="extract and create 7 parts of image from HD one")
  91. parser.add_argument('--folder', type=str, help="folder with HD images", required=True)
  92. parser.add_argument('--output', type=str, help="output data folder", required=True)
  93. args = parser.parse_args()
  94. p_folder = args.folder
  95. p_output = args.output
  96. print('Start extraction for folder %s' % p_folder)
  97. extract(p_folder, p_output)
  98. if __name__ == "__main__":
  99. main()