Merge branch 'master'

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+model/

app.py

@@ -0,0 +1,98 @@
+import os
+import app_configs as configs
+import service
+import gradio as gr
+import numpy as np
+import cv2 
+from PIL import Image
+import logging 
+from huggingface_hub import hf_hub_download
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+sam = None #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device)
+red = (255,0,0)
+blue = (0,0,255)
+
+def load_sam_instance():
+    global sam
+    if sam is None:
+        gr.Info('Initialising SAM, hang in there...')
+        if not os.path.exists(configs.model_ckpt_path):
+            gr.Info('Downloading weights from  hugging face hub')
+            chkpt_path = hf_hub_download("ybelkada/segment-anything", configs.model_ckpt_path)
+        else:
+            chkpt_path = configs.model_ckpt_path
+        sam = service.get_sam(configs.model_type, chkpt_path, configs.device)
+    return sam
+
+block = gr.Blocks()
+with block:
+    # states
+    def point_coords_empty():
+        return []
+    def point_labels_empty():
+        return []
+    point_coords = gr.State(point_coords_empty)
+    point_labels = gr.State(point_labels_empty)
+    raw_image = gr.Image(type='pil', visible=False)
+
+    # UI
+    with gr.Row():
+        with gr.Column():
+            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')
+        with gr.Column():
+            with gr.Tab('Cutout'):
+                cutout_gallery = gr.Gallery()
+            with gr.Tab('Annotation'):
+                masks_annotated_image = gr.AnnotatedImage(label='Segments')
+            gr.Examples(examples=[['examples/cat-256.png','examples/cat-256.png']],inputs=[input_image, raw_image])
+    
+    # components
+    components = {point_coords, point_labels, raw_image, input_image, point_label_radio, reset_btn, run_btn, cutout_gallery, masks_annotated_image}
+
+    # event - init coords
+    def on_reset_btn_click(raw_image):
+        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):
+        return input_image, point_coords_empty(), point_labels_empty(), None
+    input_image.upload(on_input_image_upload, [input_image], [raw_image, point_coords, point_labels], queue=False)
+
+    # event - set coords
+    def on_input_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
+        x, y = evt.index
+        color = red if point_label_radio == 0 else blue
+        img = np.array(input_image)
+        cv2.circle(img, (x, y), 10, color, -1)
+        img = Image.fromarray(img)
+        point_coords.append([x,y])
+        point_labels.append(point_label_radio)
+        return img, point_coords, point_labels
+    input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False)
+
+    # event - inference
+    def on_run_btn_click(data):
+        sam = load_sam_instance()
+        image = data[raw_image]
+        if len(data[point_coords]) == 0:
+            masks, _ = service.predict_all(sam, image)
+        else:
+            masks, _ = service.predict_conditioned(sam,
+                                                   image, 
+                                                   point_coords=np.array(data[point_coords]), 
+                                                   point_labels=np.array(data[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, masks
+    run_btn.click(on_run_btn_click, components, [cutout_gallery, masks_annotated_image], queue=True)
+
+if __name__ == '__main__':
+    block.queue()
+    block.launch()

app_configs.py

@@ -0,0 +1,4 @@
+model_type = r'vit_b'
+# model_ckpt_path = None
+model_ckpt_path = "checkpoints/sam_vit_b_01ec64.pth"
+device = 'cpu'

requirements.txt

@@ -0,0 +1,5 @@
+huggingface_hub
+gradio
+torch
+torchvision
+git+https://github.com/facebookresearch/segment-anything.git

service.py

@@ -0,0 +1,90 @@
+from typing import IO, List
+import torch
+from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
+from PIL import Image
+import numpy as np
+import io
+
+def to_file(item) -> IO[bytes]:
+    # Create a BytesIO object
+    file_obj = io.BytesIO()
+    if isinstance(item, Image.Image):
+        item.save(file_obj, format='PNG')
+    if isinstance(item, np.ndarray):
+        np.save(file_obj, item)
+    # Reset the file object's position to the beginning
+    file_obj.seek(0)
+    # Return the file object
+    return file_obj
+
+def get_sam(model_type, checkpoint_path, device=None):
+    if device is None:
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
+    sam.to(device=device)
+    return sam 
+
+def draw_mask(img: Image.Image, boolean_mask: np.ndarray, color: tuple, mask_alpha: float) -> Image.Image:
+    int_alpha = int(mask_alpha*255)
+    color_mask = Image.new('RGBA', img.size, color=color)
+    color_mask.putalpha(Image.fromarray(boolean_mask.astype(np.uint8)*int_alpha, mode='L'))
+    result = Image.alpha_composite(img, color_mask)
+    
+    return result
+def random_color():
+    return tuple(np.random.randint(0,255, 3))
+
+def draw_masks(img: Image.Image, boolean_masks: np.ndarray) -> Image.Image:
+    img = img.copy()
+    for boolean_mask in boolean_masks:
+        img = draw_mask(img, boolean_mask, random_color(), 0.2)
+    return img
+
+def cutout(img: Image.Image, boolean_mask: np.ndarray):
+    rgba_img = img.convert('RGBA')
+    mask = Image.fromarray(boolean_mask).convert("L")
+    rgba_img.putalpha(mask)
+    return rgba_img
+
+
+def predict_conditioned(sam, pil_img, **kwargs):
+    rgb_arr = pil_image_to_rgb_array(pil_img)
+    predictor = SamPredictor(sam)
+    predictor.set_image(rgb_arr)
+    masks, quality, _ = predictor.predict(**kwargs)
+    return masks, quality
+
+def predict_all(sam, pil_img):
+    rgb_arr = pil_image_to_rgb_array(pil_img)
+    mask_generator = SamAutomaticMaskGenerator(sam)
+    results = mask_generator.generate(rgb_arr)
+    masks = []
+    quality = []
+    for result in results:
+        masks.append(result['segmentation'])
+        quality.append(result['stability_score'])
+    masks = np.array(masks)
+    quality = np.array(quality)
+    return masks, quality
+    
+def pil_image_to_rgb_array(image):
+    if image.mode == "RGBA":
+        rgb_image = Image.new("RGB", image.size, (255, 255, 255))
+        rgb_image.paste(image, mask=image.split()[3])  # Apply alpha channel as the mask
+        rgb_array = np.array(rgb_image)        
+    else:
+        rgb_array = np.array(image.convert("RGB"))
+    return rgb_array
+
+def box_pts_to_xyxy(pt1, pt2):
+    """convert box from pts format to XYXY
+    Args:
+        pt1 : (x1, y1) first corner of a box
+        pt2 : (x2, y2) second corner, diagonal to pt1
+
+    Returns:
+        xyxy: (x_min, y_min, x_max, y_max)
+    """
+    x1, y1 = pt1
+    x2, y2 = pt2
+    return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))