Update app.py

app.py

@@ -1,18 +1,16 @@
 import gradio as gr
 import torch
-import cv2
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
 import io
-from ultralytics import FastSAM
-from ultralytics.models.fastsam import FastSAMPrompt
+from sam2.sam2_image_predictor import SAM2ImagePredictor
 
 # Set up device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load FastSAM model
-model = FastSAM("FastSAM-s.pt")  # or FastSAM-x.pt
+# Load SAM 2 model
+predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
 
 def fig2img(fig):
     buf = io.BytesIO()
@@ -21,16 +19,13 @@ def fig2img(fig):
     img = Image.open(buf)
     return img
 
-def plot_masks(annotations, output_shape):
+def plot_masks(image, masks):
     fig, ax = plt.subplots(figsize=(10, 10))
-    ax.imshow(annotations[0].orig_img)
+    ax.imshow(image)
     
-    for ann in annotations:
-        for mask in ann.masks.data:
-            mask = cv2.resize(mask.cpu().numpy().astype('uint8'), output_shape[::-1])
-            masked = np.ma.masked_where(mask == 0, mask)
-            ax.imshow(masked, alpha=0.5, cmap=plt.cm.get_cmap('jet'))
-
+    for mask in masks:
+        masked = np.ma.masked_where(mask == 0, mask)
+        ax.imshow(masked, alpha=0.5, cmap=plt.cm.get_cmap('jet'))
     ax.axis('off')
     plt.close()
     return fig2img(fig)
@@ -42,19 +37,15 @@ def segment_everything(input_image):
         
         input_image = Image.fromarray(input_image).convert("RGB")
         
-        # Run FastSAM model in "everything" mode
-        everything_results = model(input_image, device=device, retina_masks=True, imgsz=1024, conf=0.25, iou=0.9, agnostic_nms=True)
-        
-        # Prepare a Prompt Process object
-        prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
-        
-        # Get everything segmentation
-        ann = prompt_process.everything_prompt()
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+            predictor.set_image(input_image)
+            # Use 'everything' prompt
+            masks, _, _ = predictor.predict([])
         
         # Plot the results
-        result_image = plot_masks(ann, input_image.size)
+        result_image = plot_masks(input_image, masks)
         
-        return result_image, f"Segmented everything in the image. Found {len(ann[0].masks)} objects."
+        return result_image, f"Segmented everything in the image. Found {len(masks)} objects."
     
     except Exception as e:
         return None, f"An error occurred: {str(e)}"
@@ -69,8 +60,8 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Segmented Image"),
         gr.Textbox(label="Status")
     ],
-    title="FastSAM Everything Segmentation",
-    description="Upload an image to segment all objects using FastSAM."
+    title="SAM 2 Everything Segmentation",
+    description="Upload an image to segment all objects using SAM 2."
 )
 
 # Launch the interface