Update app.py

app.py

@@ -9,8 +9,19 @@ from torch.utils.data import Dataset, DataLoader
 import streamlit as st
 import gradio as gr
 
-# Load the Segformer model using Gradio (Optional)
-gr.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+feature_extractor = SegformerFeatureExtractor.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')
+segformer_model = SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')
+
+# Inference function to use the Segformer model
+def segment_image(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = segformer_model(**inputs)
+    segmentation = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
+    return segmentation
+
+# Use Gradio interface to display the segmentation result
+iface = gr.Interface(fn=segment_image, inputs=gr.Image(type="pil"), outputs="image")
+iface.launch()  # Launch Gradio UI
 
 # Function to extract zip files
 def extract_zip(zip_file, extract_to):