Update app.py

app.py

@@ -1,3 +1,21 @@
+from transformers import pipeline
 import gradio as gr
 
-gr.Interface.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+# Load the model for image segmentation
+segmentation = pipeline("image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512")
+
+# Define a function to segment the input image
+def segment_image(image):
+    result = segmentation(image)
+    # Return the segmentation mask
+    return result[0]['mask']  # Returns the mask as an image
+
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=segment_image,  # Function to be called for segmentation
+    inputs=gr.Image(type="pil"),  # Input: image
+    outputs="image"  # Output: segmented image (mask)
+)
+
+# Launch the Gradio app
+interface.launch()