Added type hinting and some clean up

.gitignore

@@ -1,2 +1,3 @@
 test.ipynb
 data
+__pycache__

app.py

@@ -5,6 +5,7 @@ from PIL import Image, ImageDraw
 import requests
 from transformers import SamModel, SamProcessor
 import cv2
+from typing import List
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -12,7 +13,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
-def mask_2_dots(mask):
+def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
     gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
     _, thresh = cv2.threshold(gray, 127, 255, 0)
     kernel = np.ones((5,5),np.uint8)
@@ -26,34 +27,32 @@ def mask_2_dots(mask):
         points.append([cx, cy])
     return [points]
 
-def main_func(inputs):
-    dots = inputs['mask']
-    points = mask_2_dots(dots)
-
-    image_input = inputs['image']
+def foward_pass(image_input: np.ndarray, points: List[List[int]]) -> np.ndarray:
     image_input = Image.fromarray(image_input)
 
     inputs = processor(image_input, input_points=points, return_tensors="pt").to(device)
-    # Forward pass
     outputs = model(**inputs)
-
-    # Postprocess outputs
-    draw = ImageDraw.Draw(image_input)
-    for point in points[0]:
-        draw.ellipse((point[0] - 10, point[1] - 10, point[0] + 10, point[1] + 10), fill="red")
-
-
     masks = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
     )
-    #scores = outputs.iou_scores
+    masks = masks[0].squeeze(0).numpy().transpose(1, 2, 0)
 
-    mask = masks[0].squeeze(0).numpy().transpose(1, 2, 0)
+    return masks
+
+def main_func(inputs) -> List[Image.Image]:
+    dots = inputs['mask']
+    points = mask_2_dots(dots)
+    image_input = inputs['image']
+    masks = foward_pass(image_input, points)
+
+    image_input = Image.fromarray(image_input)
+    draw = ImageDraw.Draw(image_input)
+    for point in points[0]:
+        draw.ellipse((point[0] - 10, point[1] - 10, point[0] + 10, point[1] + 10), fill="red")
 
     pred_masks = [image_input]
-    for i in range(mask.shape[2]):
-        #mask[:,:,i] = mask[:,:,i] * scores[0][i].item()
-        pred_masks.append(Image.fromarray((mask[:,:,i] * 255).astype(np.uint8)))
+    for i in range(masks.shape[2]):
+        pred_masks.append(Image.fromarray((masks[:,:,i] * 255).astype(np.uint8)))
 
     return pred_masks