feature: separate annotation and cutout

.gitignore

@@ -1,2 +1,3 @@
 __pycache__/
-model/
+model/
+flagged/

app.py

@@ -33,31 +33,31 @@ with block:
         return []
     def point_labels_empty():
         return []
+    raw_image = gr.Image(type='pil', visible=False)
     point_coords = gr.State(point_coords_empty)
     point_labels = gr.State(point_labels_empty)
-    raw_image = gr.Image(type='pil', visible=False)
+    masks = gr.State()
+    cutout_idx = gr.State(set())
 
     # UI
-    with gr.Row():
-        with gr.Column():
+    with gr.Column():
+        with gr.Row():
             input_image = gr.Image(label='Input', height=512, type='pil')
-            with gr.Row():
-                point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
-                reset_btn = gr.Button('Reset')
-                run_btn = gr.Button('Run', variant = 'primary')
-            gr.Examples(examples=[['examples/cat-256.png','examples/cat-256.png']],inputs=[input_image, raw_image])
-        with gr.Column():
-            with gr.Tab('Cutout'):
-                cutout_gallery = gr.Gallery()
-            with gr.Tab('Annotation'):
-                masks_annotated_image = gr.AnnotatedImage(label='Segments')
-    
+            masks_annotated_image = gr.AnnotatedImage(label='Segments')
+        with gr.Row():
+            point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
+            reset_btn = gr.Button('Reset')
+            run_btn = gr.Button('Run', variant = 'primary')
+        cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain')
+
     # components
-    components = {point_coords, point_labels, raw_image, input_image, point_label_radio, reset_btn, run_btn, cutout_gallery, masks_annotated_image}
+    components = {
+        point_coords, point_labels, raw_image, masks, cutout_idx,
+        input_image, point_label_radio, reset_btn, run_btn, masks_annotated_image}
 
     # event - init coords
     def on_reset_btn_click(raw_image):
-        return raw_image, point_coords_empty(), point_labels_empty(), None
+        return raw_image, point_coords_empty(), point_labels_empty(), None, []
     reset_btn.click(on_reset_btn_click, [raw_image], [input_image, point_coords, point_labels], queue=False)
 
     def on_input_image_upload(input_image):
@@ -91,9 +91,16 @@ with block:
                                                    point_coords=np.array(inputs[point_coords]), 
                                                    point_labels=np.array(inputs[point_labels]))
         annotated = (image, [(masks[i], f'Mask {i}') for i in range(len(masks))])
-        cutouts = [service.cutout(image, mask) for mask in masks]
-        return cutouts, annotated
-    run_btn.click(on_run_btn_click, components, [cutout_gallery, masks_annotated_image], queue=True)
+        return annotated, masks, set()
+    run_btn.click(on_run_btn_click, components, [masks_annotated_image, masks, cutout_idx], queue=True)
+
+    # event - get cutout 
+    def on_masks_annotated_image_select(inputs, evt:gr.SelectData):
+        inputs[cutout_idx].add(evt.index)
+        cutouts = [service.cutout(inputs[raw_image], inputs[masks][idx]) for idx in list(inputs[cutout_idx])]
+        tight_cutouts = [service.crop_empty(cutout) for cutout in cutouts]
+        return inputs[cutout_idx], tight_cutouts
+    masks_annotated_image.select(on_masks_annotated_image_select, components, [cutout_idx, cutout_galary])
 
 if __name__ == '__main__':
     block.queue()

app_configs.py

@@ -2,4 +2,5 @@ model_type = r'vit_b'
 # model_ckpt_path = None
 model_ckpt_path = "checkpoints/sam_vit_b_01ec64.pth"
 device = 'cpu'
-enable_segment_all = False
+enable_segment_all = False
+flagging_dir = r'.\flagged'

service.py

@@ -1,4 +1,5 @@
 from typing import IO, List
+import cv2
 import torch
 from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
@@ -87,4 +88,25 @@ def box_pts_to_xyxy(pt1, pt2):
     """
     x1, y1 = pt1
     x2, y2 = pt2
-    return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))
+    return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))
+
+def crop_empty(image:Image.Image):
+    # Convert image to numpy array
+    np_image = np.array(image)
+
+    # Find non-transparent pixels
+    non_transparent_pixels = np_image[:, :, 3] > 0
+
+    # Calculate bounding box coordinates
+    rows = np.any(non_transparent_pixels, axis=1)
+    cols = np.any(non_transparent_pixels, axis=0)
+    ymin, ymax = np.where(rows)[0][[0, -1]]
+    xmin, xmax = np.where(cols)[0][[0, -1]]
+
+    # Crop the image
+    cropped_image = np_image[ymin:ymax+1, xmin:xmax+1, :]
+
+    # Convert cropped image back to PIL image
+    pil_image = Image.fromarray(cropped_image)
+
+    return pil_image