Update app.py

app.py

@@ -2,28 +2,43 @@ import gradio as gr
 import torch
 import cv2
 import numpy as np
-from fastsam import FastSAM, FastSAMPrompt
+from transformers import SamModel, SamProcessor
+from PIL import Image
 
-# Load the FastSAM model
-model = FastSAM('FastSAM-x.pt')
+# Set up device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model and processor
+model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
+processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
 def segment_image(input_image, points):
-    # Prepare the image
-    input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
+    # Convert input_image to PIL Image
+    input_image = Image.fromarray(input_image)
+    
+    # Prepare inputs
+    inputs = processor(input_image, input_points=[points], return_tensors="pt").to(device)
     
-    # Run the model
-    everything_results = model(input_image, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
+    # Generate masks
+    with torch.no_grad():
+        outputs = model(**inputs)
     
-    # Prepare prompts
-    prompt_process = FastSAMPrompt(input_image, everything_results, device='cpu')
+    # Post-process masks
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )
+    scores = outputs.iou_scores
     
-    # Generate mask based on points
-    ann = prompt_process.point_prompt(points=points, pointlabel=[1] * len(points))
+    # Convert mask to numpy array
+    mask = masks[0][0].numpy()
     
     # Overlay the mask on the original image
-    result_image = input_image.copy()
-    mask = ann[0].astype(bool)
-    result_image[mask] = result_image[mask] * 0.5 + np.array([255, 0, 0]) * 0.5
+    result_image = np.array(input_image)
+    mask_rgb = np.zeros_like(result_image)
+    mask_rgb[mask > 0.5] = [255, 0, 0]  # Red color for the mask
+    result_image = cv2.addWeighted(result_image, 1, mask_rgb, 0.5, 0)
     
     return result_image
 
@@ -35,8 +50,8 @@ iface = gr.Interface(
         gr.Image(type="numpy", tool="sketch", brush_radius=5, label="Click on objects to segment")
     ],
     outputs=gr.Image(type="numpy"),
-    title="FastSAM Image Segmentation",
-    description="Click on objects in the image to segment them using FastSAM."
+    title="Segment Anything Model (SAM) Image Segmentation",
+    description="Click on objects in the image to segment them using SAM."
 )
 
 # Launch the interface