Refactor to add type hint

modules/mask_utils.py

@@ -5,23 +5,34 @@ from pycocotools import mask as coco_mask
 from pytoshop import layers
 import pytoshop
 from pytoshop.enums import BlendMode
+from pytoshop.core import PsdFile
+
+
+def decode_to_mask(seg: np.ndarray[np.bool_]) -> np.ndarray[np.uint8]:
+    return seg.astype(np.uint8) * 255
 
 
 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):
+def create_base_layer(image: np.ndarray):
     rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
     return [rgba_image]
 
 
-def create_mask_layers(image, masks):
+def create_mask_layers(
+    image: np.ndarray,
+    masks: List
+):
     layer_list = []
 
-    for result in masks:
-        rle = result['segmentation']
-        mask = coco_mask.decode(rle).astype(np.uint8)
+    sorted_masks = sorted(masks, key=lambda x: x['area'], reverse=True)
+
+    for info in sorted_masks:
+        rle = info['segmentation']
+        mask = decode_to_mask(rle)
+
         rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
         rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
 
@@ -30,13 +41,18 @@ def create_mask_layers(image, masks):
     return layer_list
 
 
-def create_mask_gallery(image, masks):
+def create_mask_gallery(
+    image: np.ndarray,
+    masks: List
+):
     mask_array_list = []
     label_list = []
 
-    for index, result in enumerate(masks):
-        rle = result['segmentation']
-        mask = coco_mask.decode(rle).astype(np.uint8)
+    sorted_masks = sorted(masks, key=lambda x: x['area'], reverse=True)
+
+    for index, info in enumerate(sorted_masks):
+        rle = info['segmentation']
+        mask = decode_to_mask(rle)
 
         rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
         rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
@@ -47,12 +63,15 @@ def create_mask_gallery(image, masks):
     return [[img, label] for img, label in zip(mask_array_list, label_list)]
 
 
-def create_mask_combined_images(image, masks):
+def create_mask_combined_images(
+    image: np.ndarray,
+    masks: List
+):
     final_result = np.zeros_like(image)
 
-    for result in masks:
-        rle = result['segmentation']
-        mask = coco_mask.decode(rle).astype(np.uint8)
+    for info in masks:
+        rle = info['segmentation']
+        mask = decode_to_mask(rle)
 
         color = generate_random_color()
         colored_mask = np.zeros_like(image)
@@ -64,7 +83,12 @@ def create_mask_combined_images(image, masks):
     return [combined_image, "masked"]
 
 
-def insert_psd_layer(psd, image_data, layer_name, blending_mode):
+def insert_psd_layer(
+    psd: PsdFile,
+    image_data: np.ndarray,
+    layer_name: str,
+    blending_mode: BlendMode
+):
     channel_data = [layers.ChannelImageData(image=image_data[:, :, i], compression=1) for i in range(4)]
 
     layer_record = layers.LayerRecord(
@@ -78,8 +102,14 @@ def insert_psd_layer(psd, image_data, layer_name, blending_mode):
     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])
+def save_psd(
+    input_image_data: np.ndarray,
+    layer_data: List,
+    layer_names: List,
+    blending_modes: List,
+    output_path: str
+):
+    psd_file = 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):
@@ -91,7 +121,7 @@ def save_psd(input_image_data, layer_data, layer_names, blending_modes, output_p
 
 def save_psd_with_masks(
     image: np.ndarray,
-    masks: Dict,
+    masks: List,
     output_path: str
 ):
     original_layer = create_base_layer(image)