Update app.py

app.py

@@ -1,14 +1,18 @@
-# Code credit: [EdgeSAM Demo](https://huggingface.co/spaces/chongzhou/EdgeSAM).
+# Code credit: [EdgeSAM Demo](https://huggingface.co/spaces/yunyangx/EfficientSAM).
+
+import copy
+import os  # noqa
 
-import torch
 import gradio as gr
-from huggingface_hub import snapshot_download
 import numpy as np
-from tinysam import sam_model_registry, SamPredictor
+import torch
 from PIL import ImageDraw
+from torchvision.transforms import ToTensor
+
+from utils.tools import format_results, point_prompt
 from utils.tools_gradio import fast_process
-import copy
-import argparse
+from tinysam import sam_model_registry, SamPredictor
+
 
 snapshot_download("merve/tinysam", local_dir="tinysam")
 
@@ -19,6 +23,28 @@ sam.to(device=device)
 sam.eval()
 predictor = SamPredictor(sam)
 
+# Description
+title = "<center><strong><font size='8'>TinySAM<font></strong> <a href='https://github.com/xinghaochen/TinySAM'><font size='6'>[GitHub]</font></a> </center>"
+
+description_e = """This is a demo of TinySAM Model](https://github.com/xinghaochen/TinySAM).
+              """
+
+description_p = """# Interactive Instance Segmentation
+                - Point-prompt instruction
+                <ol>
+                <li> Click on the left image (point input), visualizing the point on the right image </li>
+                <li> Click the button of Segment with Point Prompt </li>
+                </ol>
+                - Box-prompt instruction
+                <ol>
+                <li> Click on the left image (one point input), visualizing the point on the right image </li>
+                <li> Click on the left image (another point input), visualizing the point and the box on the right image</li>
+                <li> Click the button of Segment with Box Prompt </li>
+                </ol>
+                - Github [link](https://github.com/xinghaochen/TinySAM)
+              """
+
+# examples
 examples = [
     ["assets/1.jpg"],
     ["assets/2.jpg"],
@@ -28,67 +54,126 @@ examples = [
     ["assets/6.jpeg"]
 ]
 
-# Description
-title = "<center><strong><font size='8'>TinySAM<font></strong> <a href='https://github.com/xinghaochen/TinySAM'><font size='6'>[GitHub]</font></a> </center>"
-
-description_p = """ # Instructions for point mode
+default_example = examples[0]
 
-                1. Upload an image or click one of the provided examples.
-                2. Select the point type.
-                3. Click once or multiple times on the image to indicate the object of interest.
-                4. The Clear button clears all the points.
-                5. The Reset button resets both points and the image.
-
-              """
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
 
-description_b = """ # Instructions for box mode
 
-                1. Upload an image or click one of the provided examples.
-                2. Click twice on the image (diagonal points of the box).
-                3. The Clear button clears the box.
-                4. The Reset button resets both the box and the image.
+def segment_with_boxs(
+    image,
+    seg_image,
+    global_points,
+    global_point_label,
+    input_size=1024,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
+):
+    if len(global_points) < 2:
+        return seg_image, global_points, global_point_label
+    print("Original Image : ", image.size)
 
-              """
+    input_size = int(input_size)
+    w, h = image.size
+    scale = input_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image = image.resize((new_w, new_h))
 
-css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+    print("Scaled Image : ", image.size)
+    print("Scale : ", scale)
 
+    scaled_points = np.array(
+        [[int(x * scale) for x in point] for point in global_points]
+    )
+    scaled_points = scaled_points[:2]
+    scaled_point_label = np.array(global_point_label)[:2]
 
-def reset(session_state):
-    session_state['coord_list'] = []
-    session_state['label_list'] = []
-    session_state['box_list'] = []
-    session_state['ori_image'] = None
-    session_state['image_with_prompt'] = None
-    session_state['feature'] = None
-    return None, session_state
+    print(scaled_points, scaled_points is not None)
+    print(scaled_point_label, scaled_point_label is not None)
 
+    if scaled_points.size == 0 and scaled_point_label.size == 0:
+        print("No points selected")
+        return image, global_points, global_point_label
 
-def reset_all(session_state):
-    session_state['coord_list'] = []
-    session_state['label_list'] = []
-    session_state['box_list'] = []
-    session_state['ori_image'] = None
-    session_state['image_with_prompt'] = None
-    session_state['feature'] = None
-    return None, None, session_state
+    nd_image = np.array(image)
+    img_tensor = ToTensor()(nd_image)
 
+    coord_np = np.array(session_state['coord_list'])
+    label_np = np.array(session_state['label_list'])
+    masks, scores, logits = predictor.predict(
+        point_coords=coord_np,
+        point_labels=label_np,
+    )
+    print(f'scores: {scores}')
+    area = masks.sum(axis=(1, 2))
+    print(f'area: {area}')
+    annotations = np.expand_dims(masks[scores.argmax()], axis=0)
+    
+    print(img_tensor.shape)
+    pts_sampled = torch.reshape(torch.tensor(scaled_points), [1, 1, -1, 2])
+    pts_sampled = pts_sampled[:, :, :2, :]
+    pts_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2])
+
+    predicted_logits, predicted_iou = model(
+        img_tensor[None, ...].to(device),
+        pts_sampled.to(device),
+        pts_labels.to(device),
+    )
+    predicted_logits = predicted_logits.cpu()
+    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
+    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
+
+
+    max_predicted_iou = -1
+    selected_mask_using_predicted_iou = None
+    selected_predicted_iou = None
+
+    for m in range(all_masks.shape[0]):
+        curr_predicted_iou = predicted_iou[m]
+        if (
+            curr_predicted_iou > max_predicted_iou
+            or selected_mask_using_predicted_iou is None
+        ):
+            max_predicted_iou = curr_predicted_iou
+            selected_mask_using_predicted_iou = all_masks[m:m+1]
+            selected_predicted_iou = predicted_iou[m:m+1]
+
+    results = format_results(selected_mask_using_predicted_iou, selected_predicted_iou, predicted_logits, 0)
+
+    annotations = results[0]["segmentation"]
+    annotations = np.array([annotations])
+    print(scaled_points.shape)
+    fig = fast_process(
+        annotations=annotations,
+        image=image,
+        device=device,
+        scale=(1024 // input_size),
+        better_quality=better_quality,
+        mask_random_color=mask_random_color,
+        use_retina=use_retina,
+        bbox = scaled_points.reshape([4]),
+        withContours=withContours,
+    )
 
-def clear(session_state):
-    session_state['coord_list'] = []
-    session_state['label_list'] = []
-    session_state['box_list'] = []
-    session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image'])
-    return session_state['ori_image'], session_state
+    global_points = []
+    global_point_label = []
+    # return fig, None
+    return fig, global_points, global_point_label
 
 
-def on_image_upload(
+def segment_with_points(
     image,
-    session_state,
-    input_size=1024
+    global_points,
+    global_point_label,
+    input_size=1024,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
 ):
-    session_state['coord_list'] = []
-    session_state['label_list'] = []
-    session_state['box_list'] = []
+    print("Original Image : ", image.size)
 
     input_size = int(input_size)
     w, h = image.size
@@ -96,232 +181,244 @@ def on_image_upload(
     new_w = int(w * scale)
     new_h = int(h * scale)
     image = image.resize((new_w, new_h))
-    session_state['ori_image'] = copy.deepcopy(image)
-    session_state['image_with_prompt'] = copy.deepcopy(image)
-    print("Image changed")
-    nd_image = np.array(image)
-    session_state['feature'] = predictor.set_image(nd_image)
 
-    return image, session_state
+    print("Scaled Image : ", image.size)
+    print("Scale : ", scale)
 
+    if global_points is None:
+        return image, global_points, global_point_label
+    if len(global_points) < 1:
+        return image, global_points, global_point_label
+    scaled_points = np.array(
+        [[int(x * scale) for x in point] for point in global_points]
+    )
+    scaled_point_label = np.array(global_point_label)
 
-def convert_box(xyxy):
-    min_x = min(xyxy[0][0], xyxy[1][0])
-    max_x = max(xyxy[0][0], xyxy[1][0])
-    min_y = min(xyxy[0][1], xyxy[1][1])
-    max_y = max(xyxy[0][1], xyxy[1][1])
-    xyxy[0][0] = min_x
-    xyxy[1][0] = max_x
-    xyxy[0][1] = min_y
-    xyxy[1][1] = max_y
-    return xyxy
+    print(scaled_points, scaled_points is not None)
+    print(scaled_point_label, scaled_point_label is not None)
 
+    if scaled_points.size == 0 and scaled_point_label.size == 0:
+        print("No points selected")
+        return image, global_points, global_point_label
 
-def segment_with_points(
-    label,
-    session_state,
-    evt: gr.SelectData,
-    input_size=1024,
-    better_quality=False,
-    withContours=True,
-    use_retina=True,
-    mask_random_color=False,
-):
-    x, y = evt.index[0], evt.index[1]
-    print(f"x y: {x,y}")
-    point_radius, point_color = 5, (97, 217, 54) if label == "Positive" else (237, 34, 13)
-    session_state['coord_list'].append([x, y])
-    session_state['label_list'].append(1 if label == "Positive" else 0)
-
-    print(f"coord_list: {session_state['coord_list']}")
-    print(f"label_list: {session_state['label_list']}")
-
-    draw = ImageDraw.Draw(session_state['image_with_prompt'])
-    draw.ellipse(
-        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
-        fill=point_color,
-    )
-    image = session_state['image_with_prompt']
-
-    coord_np = np.array(session_state['coord_list'])
-    label_np = np.array(session_state['label_list'])
+    nd_image = np.array(image)
+    img_tensor = ToTensor()(nd_image)
+    print(img_tensor.shape)    
+    
+    predictor.set_image(nd_image)
     masks, scores, logits = predictor.predict(
-        point_coords=coord_np,
-        point_labels=label_np,
+        point_coords=scaled_points,
+        point_labels=global_point_label,
     )
     print(f'scores: {scores}')
     area = masks.sum(axis=(1, 2))
     print(f'area: {area}')
-
     annotations = np.expand_dims(masks[scores.argmax()], axis=0)
 
-    seg = fast_process(
+    fig = fast_process(
         annotations=annotations,
         image=image,
         device=device,
         scale=(1024 // input_size),
         better_quality=better_quality,
         mask_random_color=mask_random_color,
+        points = scaled_points,
         bbox=None,
         use_retina=use_retina,
         withContours=withContours,
     )
 
-    return seg, session_state
+    global_points = []
+    global_point_label = []
+    # return fig, None
+    return fig, global_points, global_point_label
 
 
-def segment_with_box(
-        session_state,
-        evt: gr.SelectData,
-        input_size=1024,
-        better_quality=False,
-        withContours=True,
-        use_retina=True,
-        mask_random_color=False,
-):
+def get_points_with_draw(image, cond_image, global_points, global_point_label, evt: gr.SelectData):
+    print("Starting functioning")
+    if len(global_points) == 0:
+        image = copy.deepcopy(cond_image)
     x, y = evt.index[0], evt.index[1]
-    point_radius, point_color, box_outline = 5, (97, 217, 54), 5
-    box_color = (0, 255, 0)
-
-    if len(session_state['box_list']) == 0:
-        session_state['box_list'].append([x, y])
-    elif len(session_state['box_list']) == 1:
-        session_state['box_list'].append([x, y])
-    elif len(session_state['box_list']) == 2:
-        session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image'])
-        session_state['box_list'] = [[x, y]]
-
-    print(f"box_list: {session_state['box_list']}")
-
-    draw = ImageDraw.Draw(session_state['image_with_prompt'])
-    draw.ellipse(
-        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
-        fill=point_color,
+    label = "Add Mask"
+    point_radius, point_color = 15, (255, 255, 0) if label == "Add Mask" else (
+        255,
+        0,
+        255,
     )
-    image = session_state['image_with_prompt']
-
-    if len(session_state['box_list']) == 2:
-        box = convert_box(session_state['box_list'])
-        xy = (box[0][0], box[0][1], box[1][0], box[1][1])
-        draw.rectangle(
-            xy,
-            outline=box_color,
-            width=box_outline
-        )
+    global_points.append([x, y])
+    global_point_label.append(1 if label == "Add Mask" else 0)
 
-        box_np = np.array(xy)
-        masks, scores, _ = predictor.predict(
-            point_coords=None,
-            point_labels=None,
-            box=box_np[None, :],
-        )
-        annotations = np.expand_dims(masks[scores.argmax()], axis=0)
-
-        seg = fast_process(
-            annotations=annotations,
-            image=image,
-            device=device,
-            scale=(1024 // input_size),
-            better_quality=better_quality,
-            mask_random_color=mask_random_color,
-            bbox=None,
-            use_retina=use_retina,
-            withContours=withContours,
+    print(x, y, label == "Add Mask")
+
+    if image is not None:
+        draw = ImageDraw.Draw(image)
+
+        draw.ellipse(
+            [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
+            fill=point_color,
         )
-        return seg, session_state
-    return image, session_state
 
+    return image, global_points, global_point_label
 
-img_p = gr.Image(label="Input with points", type="pil", interactive=True)
-img_b = gr.Image(label="Input with box", type="pil", interactive=True)
+def get_points_with_draw_(image, cond_image, global_points, global_point_label, evt: gr.SelectData):
 
-with gr.Blocks(css=css, title="TinySAM") as demo:
-    session_state = gr.State({
-        'coord_list': [],
-        'label_list': [],
-        'box_list': [],
-        'ori_image': None,
-        'image_with_prompt': None,
-        'feature': None
-    })
+    if len(global_points) == 0:
+        image = copy.deepcopy(cond_image)
+    if len(global_points) > 2:
+        return image, global_points, global_point_label
+    x, y = evt.index[0], evt.index[1]
+    label = "Add Mask"
+    point_radius, point_color = 15, (255, 255, 0) if label == "Add Mask" else (
+        255,
+        0,
+        255,
+    )
+    global_points.append([x, y])
+    global_point_label.append(1 if label == "Add Mask" else 0)
+
+    print(x, y, label == "Add Mask")
+
+    if image is not None:
+        draw = ImageDraw.Draw(image)
 
+        draw.ellipse(
+            [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
+            fill=point_color,
+        )
+
+    if len(global_points) == 2:
+        x1, y1 = global_points[0]
+        x2, y2 = global_points[1]
+        if x1 < x2 and y1 < y2:
+            draw.rectangle([x1, y1, x2, y2], outline="red", width=5)
+        elif x1 < x2 and y1 >= y2:
+            draw.rectangle([x1, y2, x2, y1], outline="red", width=5)
+            global_points[0][0] = x1
+            global_points[0][1] = y2
+            global_points[1][0] = x2
+            global_points[1][1] = y1
+        elif x1 >= x2 and y1 < y2:
+            draw.rectangle([x2, y1, x1, y2], outline="red", width=5)
+            global_points[0][0] = x2
+            global_points[0][1] = y1
+            global_points[1][0] = x1
+            global_points[1][1] = y2
+        elif x1 >= x2 and y1 >= y2:
+            draw.rectangle([x2, y2, x1, y1], outline="red", width=5)
+            global_points[0][0] = x2
+            global_points[0][1] = y2
+            global_points[1][0] = x1
+            global_points[1][1] = y1
+
+    return image, global_points, global_point_label
+
+
+cond_img_p = gr.Image(label="Input with Point", value=default_example[0], type="pil")
+cond_img_b = gr.Image(label="Input with Box", value=default_example[0], type="pil")
+
+segm_img_p = gr.Image(
+    label="Segmented Image with Point-Prompt", interactive=False, type="pil"
+)
+segm_img_b = gr.Image(
+    label="Segmented Image with Box-Prompt", interactive=False, type="pil"
+)
+
+input_size_slider = gr.components.Slider(
+    minimum=512,
+    maximum=1024,
+    value=1024,
+    step=64,
+    label="Input_size",
+    info="Our model was trained on a size of 1024",
+)
+
+with gr.Blocks(css=css, title="TinySAM") as demo:
+    global_points = gr.State([])
+    global_point_label = gr.State([])
     with gr.Row():
         with gr.Column(scale=1):
             # Title
             gr.Markdown(title)
 
-    with gr.Tab("Point mode") as tab_p:
+    with gr.Tab("Point mode"):
         # Images
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
-                img_p.render()
-            with gr.Column(scale=1):
-                with gr.Row():
-                    add_or_remove = gr.Radio(
-                        ["Positive", "Negative"],
-                        value="Positive",
-                        label="Point Type"
-                    )
+                cond_img_p.render()
 
-                    with gr.Column():
-                        clear_btn_p = gr.Button("Clear", variant="secondary")
-                        reset_btn_p = gr.Button("Reset", variant="secondary")
-                with gr.Row():
-                    gr.Markdown(description_p)
+            with gr.Column(scale=1):
+                segm_img_p.render()
 
+        # Submit & Clear
+        # ###
         with gr.Row():
             with gr.Column():
+
+                with gr.Column():
+                    segment_btn_p = gr.Button(
+                        "Segment with Point Prompt", variant="primary"
+                    )
+                    clear_btn_p = gr.Button("Clear", variant="secondary")
+
                 gr.Markdown("Try some of the examples below ⬇️")
                 gr.Examples(
                     examples=examples,
-                    inputs=[img_p, session_state],
-                    outputs=[img_p, session_state],
-                    examples_per_page=8,
-                    fn=on_image_upload,
-                    run_on_click=True
+                    inputs=[cond_img_p],
+                    examples_per_page=4,
                 )
 
-    with gr.Tab("Box mode") as tab_b:
+            with gr.Column():
+                # Description
+                gr.Markdown(description_p)
+
+    with gr.Tab("Box mode"):
         # Images
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
-                img_b.render()
-            with gr.Row():
-                with gr.Column():
-                    clear_btn_b = gr.Button("Clear", variant="secondary")
-                    reset_btn_b = gr.Button("Reset", variant="secondary")
-                    gr.Markdown(description_b)
+                cond_img_b.render()
 
+            with gr.Column(scale=1):
+                segm_img_b.render()
+
+        # Submit & Clear
         with gr.Row():
             with gr.Column():
+
+                with gr.Column():
+                    segment_btn_b = gr.Button(
+                        "Segment with Box Prompt", variant="primary"
+                    )
+                    clear_btn_b = gr.Button("Clear", variant="secondary")
+
                 gr.Markdown("Try some of the examples below ⬇️")
                 gr.Examples(
                     examples=examples,
-                    inputs=[img_b, session_state],
-                    outputs=[img_b, session_state],
-                    examples_per_page=8,
-                    fn=on_image_upload,
-                    run_on_click=True
+                    inputs=[cond_img_b],
+
+                    examples_per_page=4,
                 )
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown(
-                "<center><img src='https://visitor-badge.laobi.icu/badge?page_id=xinghaochen.tinysam' alt='visitors'></center>")
+            with gr.Column():
+                # Description
+                gr.Markdown(description_p)
+
+    cond_img_p.select(get_points_with_draw, inputs = [segm_img_p, cond_img_p, global_points, global_point_label], outputs = [segm_img_p, global_points, global_point_label])
 
-    img_p.upload(on_image_upload, [img_p, session_state], [img_p, session_state])
-    img_p.select(segment_with_points, [add_or_remove, session_state], [img_p, session_state])
+    cond_img_b.select(get_points_with_draw_, [segm_img_b, cond_img_b, global_points, global_point_label], [segm_img_b, global_points, global_point_label])
 
-    clear_btn_p.click(clear, [session_state], [img_p, session_state])
-    reset_btn_p.click(reset, [session_state], [img_p, session_state])
-    tab_p.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])
+    segment_btn_p.click(
+        segment_with_points, inputs=[cond_img_p, global_points, global_point_label], outputs=[segm_img_p, global_points, global_point_label]
+    )
+
+    segment_btn_b.click(
+        segment_with_boxs, inputs=[cond_img_b, segm_img_b, global_points, global_point_label], outputs=[segm_img_b,global_points, global_point_label]
+    )
 
-    img_b.upload(on_image_upload, [img_b, session_state], [img_b, session_state])
-    img_b.select(segment_with_box, [session_state], [img_b, session_state])
+    def clear():
+        return None, None, [], []
 
-    clear_btn_b.click(clear, [session_state], [img_b, session_state])
-    reset_btn_b.click(reset, [session_state], [img_b, session_state])
-    tab_b.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])
+    clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p, global_points, global_point_label])
+    clear_btn_b.click(clear, outputs=[cond_img_b, segm_img_b, global_points, global_point_label])
 
 demo.queue()
 demo.launch()