Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
-import torch
+import ast
 import cv2
+import torch
 import traceback
 import numpy as np
+from itertools import chain
 from transformers import SamModel, SamProcessor
 
 
-model = SamModel.from_pretrained('facebook/sam-vit-huge').to('cpu')
+model = SamModel.from_pretrained('facebook/sam-vit-huge').to('cuda')
 processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
 
 
@@ -14,7 +16,8 @@ def set_predictor(image):
     """
     Creates a Sam predictor object based on a given image and model.
     """
-    inputs = processor(image, return_tensors='pt').to('cpu')
+    device = 'cuda'
+    inputs = processor(image, return_tensors='pt').to(device)
     image_embedding = model.get_image_embeddings(inputs['pixel_values'])
 
     return [image, image_embedding, 'Done']
@@ -23,11 +26,14 @@ def set_predictor(image):
 def get_polygon(points, image, image_embedding):
     """
     Returns the points of the polygon given a bounding box and a prediction
-    made by Sam, or if an exception was triggered, it returns such exception.
+    made by Sam.
     """
-    points = [int(w) for w in points.split(',')]
+    #points = [int(w) for w in points.split(',')]
+    points = list(chain.from_iterable(points))
+    print(points)
 
-    inputs = processor(image, input_boxes=[points], return_tensors="pt").to('cpu')
+    device = 'cuda'
+    inputs = processor(image, input_boxes=[points], return_tensors="pt").to(device)
 
     # pop the pixel_values as they are not neded
     inputs.pop("pixel_values", None)
@@ -43,39 +49,69 @@ def get_polygon(points, image, image_embedding):
     )
 
     mask = masks[0].squeeze().numpy()
-
     img = mask.astype(np.uint8)[0]
-
     contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
 
     if len(contours) == 0:
-        return [0], img
-        
+        return [], img
+
     points = contours[0]
+
     polygon = []
     for point in points:
         for x, y in point:
             polygon.append([int(x), int(y)])
 
-    return polygon
+    mask = np.zeros(image.shape, dtype='uint8')
+    poly = np.array(polygon)
+    cv2.fillPoly(mask, [poly], (0, 255, 0))
+
+    return polygon, mask
+
+
+def add_bbox(bbox, evt: gr.SelectData):
+    if bbox[0] == [0, 0]:
+        bbox[0] = [evt.index[0], evt.index[1]]
+        return bbox, bbox
+
+    bbox[1] = [evt.index[0], evt.index[1]]
+    return bbox, bbox
 
 
+def clear_bbox(bbox):
+    updated_bbox = [[0, 0], [0, 0]]
+    return updated_bbox, updated_bbox
 
-with gr.Blocks() as app:
+
+with gr.Blocks() as demo:
     image = gr.State()
     embedding = gr.State()
+    bbox = gr.State([[0, 0], [0, 0]])
 
-    with gr.Tab('Get embedding'):
+    with gr.Row():
         input_image = gr.Image(label='Image')
-        output_status = gr.Textbox(label='Status')
-        predictor_button = gr.Button('Send Image')
+        mask = gr.Image(label='Mask')
 
+    with gr.Row():
+        with gr.Column():
+            output_status = gr.Textbox(label='Status')
+            
+        with gr.Column():            
+            predictor_button = gr.Button('Send Image')
+ 
+    with gr.Row():
+        with gr.Column():
+            bbox_box = gr.Textbox(label="bbox")
 
-    with gr.Tab('Get points'):
-        bbox = gr.Textbox(label="bbox")
-        polygon = [gr.Textbox(label='Polygon')]
-        mask = gr.Image(label='Mask')
-        points_button = gr.Button('Send bounding box')
+        with gr.Column():
+            bbox_button = gr.Button('Clear bbox')
+
+    with gr.Row():
+        with gr.Column():
+            polygon = gr.Textbox(label='Polygon')
+
+        with gr.Column():
+            points_button = gr.Button('Send bounding box')
 
 
     predictor_button.click(
@@ -90,4 +126,17 @@ with gr.Blocks() as app:
         [polygon, mask],
     )
 
-app.launch(debug=True)
+    bbox_button.click(
+        clear_bbox, 
+        bbox,
+        [bbox, bbox_box],
+    )    
+
+    input_image.select(
+        add_bbox,
+        bbox,
+        [bbox, bbox_box]
+    )
+
+
+demo.launch(debug=True)