Update app.py

app.py

@@ -1,3 +1,20 @@
+from transformers import pipeline
 import gradio as gr
 
-gr.Interface.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+# Load the image segmentation model using the pipeline
+segmentation = pipeline("image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512")
+
+# Function to process the image and return the segmentation mask
+def segment_image(image):
+    result = segmentation(image)
+    return result[0]['mask']  # The segmentation mask is returned as an image
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=segment_image,  # The function to call for segmentation
+    inputs=gr.Image(type="pil"),  # Accepts image as input
+    outputs="image"  # Outputs the segmentation mask
+)
+
+# Launch the interface
+interface.launch()