Update app.py

app.py

@@ -1,3 +1,38 @@
+# import gradio as gr
+
+# gr.Interface.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+
+
 import gradio as gr
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+from PIL import Image
+import torch
+
+# Load the processor and model
+processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+
+# Define the function for segmentation
+def segment_image(image):
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Perform inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Get the segmentation mask
+    logits = outputs.logits
+    segmentation_mask = logits.argmax(dim=1).squeeze().cpu().numpy()
+    
+    # Convert to a PIL image and return it
+    segmentation_image = Image.fromarray(segmentation_mask.astype("uint8"))
+    return segmentation_image
+
+# Create the Gradio interface
+interface = gr.Interface(fn=segment_image, 
+                         inputs=gr.Image(type="pil"),  # Using gr.Image for Gradio
+                         outputs=gr.Image(type="pil"))  # Using gr.Image for Gradio
 
-gr.Interface.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+# Launch the interface
+interface.launch()