Add mask utils

modules/mask_utils.py

@@ -0,0 +1,102 @@
+import cv2
+import numpy as np
+from typing import Dict, List
+from pycocotools import mask as coco_mask
+from pytoshop import layers
+import pytoshop
+from pytoshop.enums import BlendMode
+
+
+def generate_random_color():
+    return np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)
+
+
+def create_base_layer(image):
+    rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+    return [rgba_image]
+
+
+def create_mask_layers(image, masks):
+    layer_list = []
+
+    for result in masks:
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
+
+        layer_list.append(rgba_image)
+
+    return layer_list
+
+
+def create_mask_gallery(image, masks):
+    mask_array_list = []
+    label_list = []
+
+    for index, result in enumerate(masks):
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+
+        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
+
+        mask_array_list.append(rgba_image)
+        label_list.append(f'Part {index}')
+
+    return [[img, label] for img, label in zip(mask_array_list, label_list)]
+
+
+def create_mask_combined_images(image, masks):
+    final_result = np.zeros_like(image)
+
+    for result in masks:
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+
+        color = generate_random_color()
+        colored_mask = np.zeros_like(image)
+        colored_mask[mask == 1] = color
+
+        final_result = cv2.addWeighted(final_result, 1, colored_mask, 0.5, 0)
+
+    combined_image = cv2.addWeighted(image, 1, final_result, 0.5, 0)
+    return [combined_image, "masked"]
+
+
+def insert_psd_layer(psd, image_data, layer_name, blending_mode):
+    channel_data = [layers.ChannelImageData(image=image_data[:, :, i], compression=1) for i in range(4)]
+
+    layer_record = layers.LayerRecord(
+        channels={-1: channel_data[3], 0: channel_data[0], 1: channel_data[1], 2: channel_data[2]},
+        top=0, bottom=image_data.shape[0], left=0, right=image_data.shape[1],
+        blend_mode=blending_mode,
+        name=layer_name,
+        opacity=255,
+    )
+    psd.layer_and_mask_info.layer_info.layer_records.append(layer_record)
+    return psd
+
+
+def save_psd(input_image_data, layer_data, layer_names, blending_modes, output_path):
+    psd_file = pytoshop.core.PsdFile(num_channels=3, height=input_image_data.shape[0], width=input_image_data.shape[1])
+    psd_file.layer_and_mask_info.layer_info.layer_records.clear()
+
+    for index, layer in enumerate(layer_data):
+        psd_file = insert_psd_layer(psd_file, layer, layer_names[index], blending_modes[index])
+
+    with open(output_path, 'wb') as output_file:
+        psd_file.write(output_file)
+
+
+def save_psd_with_masks(
+    image: np.ndarray,
+    masks: Dict,
+    output_path: str
+):
+    original_layer = create_base_layer(image)
+    mask_layers = create_mask_layers(image, masks)
+    names = [f'Part {i}' for i in range(len(mask_layers))]
+    modes = [BlendMode.normal] * (len(mask_layers)+1)
+    save_psd(image, original_layer+mask_layers, ['Original_Image']+names, modes, output_path)
+