Update app.py

app.py

@@ -8,20 +8,34 @@ from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
 import gradio as gr
+import os
+import zipfile
+import numpy as np
+import torch
+from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
+from transformers import ResNetForImageClassification, AdamW
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+import streamlit as st
+import gradio as gr
 
+# Load feature extractor and model
 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
+# Inference function for segmentation
 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
+# Gradio interface at the end of the script
 iface = gr.Interface(fn=segment_image, inputs=gr.Image(type="pil"), outputs="image")
-iface.launch()  # Launch Gradio UI
+
+# Specify a custom port if needed to avoid conflicts (optional)
+iface.launch(server_port=7861)  # Change port if 7860 is occupied
+
 
 # Function to extract zip files
 def extract_zip(zip_file, extract_to):