Update app.py

app.py

@@ -5,22 +5,25 @@ import torch
 import traceback
 import numpy as np
 from itertools import chain
-from transformers import SamModel, SamProcessor
+from segment_anything import SamPredictor, sam_model_registry
 
 
-model = SamModel.from_pretrained('facebook/sam-vit-huge')
-processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
+sam_checkpoint = "./checkpoints/sam_vit_h_4b8939.pth"
+model_type = "vit_h"
 
 
 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)
-    image_embedding = model.get_image_embeddings(inputs['pixel_values'])
 
-    return [image, image_embedding, 'Done']
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+    sam.to(device=device)
+
+    predictor = SamPredictor(sam)
+    predictor.set_image(image)
+
+    return [image, predictor, 'Done']
 
 
 def get_polygon(points, image, image_embedding):
@@ -30,24 +33,12 @@ def get_polygon(points, image, image_embedding):
     """
     points = list(chain.from_iterable(points))
 
-    device = 'cpu'
-    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)
-    inputs.update({"image_embeddings": image_embedding})
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    masks = processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), 
-        inputs["original_sizes"].cpu(), 
-        inputs["reshaped_input_sizes"].cpu()
+    masks, _, _ = predictor.predict(
+        box=input_box[None, :],
+        multimask_output=False,
     )
 
-    mask = masks[0].squeeze().numpy()
-    img = mask.astype(np.uint8)[0]
+    img = masks[0].astype(np.uint8)
     contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
 
     if len(contours) == 0: