Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 from transformers import SamModel, SamProcessor
 
 
-model = SamModel.from_pretrained('facebook/sam-vit-huge')
+model = SamModel.from_pretrained('facebook/sam-vit-huge').to('cuda')
 processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
 
 
@@ -14,8 +14,7 @@ def set_predictor(image):
     """
     Creates a Sam predictor object based on a given image and model.
     """
-    device = 'cpu'
-    inputs = processor(image, return_tensors='pt').to(device)
+    inputs = processor(image, return_tensors='pt').to('cuda')
     image_embedding = model.get_image_embeddings(inputs['pixel_values'])
 
     return [image, image_embedding, 'Done']
@@ -28,8 +27,7 @@ def get_polygon(points, image, image_embedding):
     """
     points = [int(w) for w in points.split(',')]
 
-    device = 'cpu'
-    inputs = processor(image, input_boxes=[points], return_tensors="pt").to(device)
+    inputs = processor(image, input_boxes=[points], return_tensors="pt").to('cuda')
 
     # pop the pixel_values as they are not neded
     inputs.pop("pixel_values", None)
@@ -49,8 +47,11 @@ def get_polygon(points, image, image_embedding):
     img = mask.astype(np.uint8)[0]
 
     contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-    points = contours[0]
 
+    if len(contours) == 0:
+        return [0], img
+        
+    points = contours[0]
     polygon = []
     for point in points:
         for x, y in point:
@@ -73,6 +74,7 @@ with gr.Blocks() as app:
     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')
 
 
@@ -85,7 +87,7 @@ with gr.Blocks() as app:
     points_button.click(
         get_polygon, 
         [bbox, image, embedding],
-        polygon,
+        [polygon, mask],
     )
 
 app.launch(debug=True)