Update app.py

app.py

@@ -1,10 +1,7 @@
 import gradio as gr
 from transformers import pipeline
-from PIL import Image
-import numpy as np
 
-# Use a pipeline as a high-level helper
-pipe = pipeline("mask-generation", model="lightmedsam")
+pipeline = pipeline(task="image-segmentation", model="Saim8250/litemedsam")
 
 def predict(image_path):
     # Perform image segmentation
@@ -18,22 +15,11 @@ def predict(image_path):
 
     return segmented_image
 
-def colorize_mask(mask):
-    # Assuming `mask` is a single-channel segmentation mask (grayscale)
-    # You may need to adjust this function based on the specifics of your model's output
-
-    # Convert single-channel mask to 3-channel (RGB) mask
-    colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
-    colored_mask[:, :, 0] = mask
-    colored_mask[:, :, 1] = mask
-    colored_mask[:, :, 2] = mask
-
-    return colored_mask
-
+    
 gr.Interface(
     predict,
-    inputs=gr.Image(label="Upload medical image", type="filepath"),
-    outputs=gr.Image(label="Segmented image"),
-    title="Segmented medical image",
+    inputs=gr.inputs.Image(label="Upload image", type="filepath"),
+    outputs=gr.outputs.Label(num_top_classes=2),
+    title="segmentation",
     allow_flagging="manual"
 ).launch()