Update app.py

app.py

@@ -1,25 +1,45 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from torchvision import transforms
+from PIL import Image
 
-pipeline = pipeline(task="image-segmentation", model="Saim8250/litemedsam")
+# Load the segmentation model (replace 'path/to/lightmed_model' with the actual path)
+model_path = 'path/to/lightmed_model'
+segmentation_model = torch.load(model_path, map_location=torch.device('cpu'))
+segmentation_model.eval()
 
-def predict(image_path):
-    # Perform image segmentation
-    predictions = pipe(image_path)
+# Define the preprocessing function for the input image
+def preprocess(image):
+    # Resize the image to match the model's expected input size
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+    ])
+    img = Image.fromarray(image)
+    img = transform(img).unsqueeze(0)
+    return img
 
-    # Access the segmented mask or any other relevant information in predictions
-    segmented_mask = predictions["segmentation_mask"]
+# Define the segmentation function
+def segment_image(input_image):
+    # Preprocess the input image
+    input_tensor = preprocess(input_image)
 
-    # Convert the segmentation mask to an RGB image
-    segmented_image = colorize_mask(segmented_mask)
+    # Perform segmentation using the model
+    with torch.no_grad():
+        output = segmentation_model(input_tensor)
 
+    # Convert the output tensor to a segmented image
+    segmented_image = torch.argmax(output, dim=1).squeeze().numpy()
+
+    # Return the segmented image
     return segmented_image
 
-    
-gr.Interface(
-    predict,
-    inputs=gr.inputs.Image(label="Upload image", type="filepath"),
-    outputs=gr.outputs.Label(num_top_classes=2),
-    title="segmentation",
-    allow_flagging="manual"
-).launch()
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=segment_image,
+    inputs=gr.Image(type="pil", preprocess=preprocess),
+    outputs=gr.Image(type="numpy")
+)
+
+# Launch the Gradio app
+iface.launch()