def split_parts(pool, data_dir, image_files):
parts_dir = get_parts_dir(data_dir)
create_dir(parts_dir)
futures = []
for img_file in image_files:
parts_dir = get_parts_dir_name(img_file)
futures.append(
pool.apply_async(split_parts_for_image, (img_file, parts_dir)))
for future in futures:
name, success = future.get()
print("parts split: %s" %
name if success else "failed to split parts: %s" % name)
def segment_images(pool, data_dir, image_files):
seg_dir = get_segmentation_dir(data_dir)
create_dir(seg_dir)
futures = []
for img_file in image_files:
seg_file = get_seg_file_name(img_file)
if not os.path.exists(seg_file):
futures.append(
pool.apply_async(create_segmentation, (img_file, seg_file)))
for future in futures:
name, success = future.get()
print("segmented: %s" % name if success else "failed to segment: %s" %
name)
def generate_default_masks(pool,
data_dir,
image_files,
write_debug_images=False):
masks_dir = get_masks_dir(data_dir)
create_dir(masks_dir)
futures = []
for img_file in image_files:
mask_file = get_mask_file_name(img_file)
if not os.path.exists(mask_file):
futures.append(
pool.apply_async(create_default_mask,
(img_file, mask_file, write_debug_images)))
for future in futures:
name = future.get()
print("default mask: %s" % name)
def downsample_images(pool, data_dir, image_files):
downsampled_dir = get_downsampled_dir(data_dir)
create_dir(downsampled_dir)
futures = []
for img_file in image_files:
downsampled_file = get_downsampled_img_name(img_file)
if not os.path.exists(downsampled_file):
futures.append(
pool.apply_async(downsample_image,
(img_file, downsampled_file)))
for future in futures:
name, success = future.get()
if success:
print("downsampled: %s" % name)
else:
print("fail: %s" % name)
def split_parts_for_image(img_file, out_dir):
try:
if os.path.exists(out_dir):
clear_directory(out_dir)
else:
create_dir(out_dir)
segmented_img = cv2.imread(get_seg_file_name(img_file))
width = segmented_img.shape[1]
height = segmented_img.shape[0]
min_area = 10
max_area = width * height / 2
mask = cv2.cvtColor((segmented_img != 0).astype(np.uint8),
cv2.COLOR_BGR2GRAY)
part_index = 0
while True:
nz = np.nonzero(mask.flatten())[0].flatten()
if len(nz) == 0:
break
nz_i = 0
found_mask = None
found_image = None
while True:
index = nz[nz_i]
seed_x = index % width
seed_y = index // width
ff_mask = np.zeros((height + 2, width + 2), dtype=np.uint8)
area, _, __, rect = cv2.floodFill(mask,
ff_mask, (seed_x, seed_y),
255,
flags=cv2.FLOODFILL_MASK_ONLY
| cv2.FLOODFILL_FIXED_RANGE)
x = rect[0]
y = rect[1]
w = rect[2]
h = rect[3]
# slicing into found rect
roi_mask = ff_mask[y + 1:y + 1 + h, x + 1:x + 1 + w]
found = False
if min_area < area < max_area:
found_mask = roi_mask
found_image = segmented_img[y:y + h, x:x + w].copy()
found_image[roi_mask == 0] = 0 # removing background
found = True
# clearing found component in the mask
mask[y:y + h, x:x + w][roi_mask != 0] = 0
if found:
break
nz_i += 1
if nz_i >= len(nz):
break
if found_mask is not None:
# we found some part
title = os.path.splitext(os.path.split(img_file)[1])[0]
part_file = os.path.join(out_dir,
"%s_%02d.png" % (title, part_index))
cv2.imwrite(part_file, found_image)
part_index += 1
return img_file, True
except Exception as _:
return img_file, False