Update app.py

app.py

@@ -1,39 +1,43 @@
 import gradio as gr
 import torch
-from PIL import Image
-from torchvision import transforms
+import cv2
+import numpy as np
+from fastsam import FastSAM, FastSAMPrompt
 
-# Load pre-trained U-Net model
-model = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'unet', pretrained=True)
+# Load the FastSAM model
+model = FastSAM('FastSAM-x.pt')
 
-# Define a function to segment an image
-def segment_image(image):
-    # Preprocess image
-    image = Image.fromarray(image)
-    image = transforms.Compose([
-        transforms.Resize((256, 256)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])(image)
+def segment_image(input_image, points):
+    # Prepare the image
+    input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
     
-    # Run segmentation model
-    output = model(image.unsqueeze(0))
-    output = torch.argmax(output, dim=1)
+    # Run the model
+    everything_results = model(input_image, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
     
-    # Postprocess output
-    output = output.squeeze(0).cpu().numpy()
-    output = Image.fromarray(output.astype('uint8'))
+    # Prepare prompts
+    prompt_process = FastSAMPrompt(input_image, everything_results, device='cpu')
     
-    return output
+    # Generate mask based on points
+    ann = prompt_process.point_prompt(points=points, pointlabel=[1] * len(points))
+    
+    # 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
+    
+    return result_image
 
-# Create Gradio app
-demo = gr.Interface(
+# Create Gradio interface
+iface = gr.Interface(
     fn=segment_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Segment Anything",
-    description="Segment any image using a pre-trained U-Net model"
+    inputs=[
+        gr.Image(type="numpy"),
+        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."
 )
 
-# Launch Gradio app
-demo.launch()
+# Launch the interface
+iface.launch()