fix

app.py

@@ -17,22 +17,21 @@ def predict(inputImage: torch.Tensor, model) -> torch.Tensor:
         prediction (torch.Tensor): The predicted segmentation mask.
     """
     with torch.no_grad():
-        output = model(preprocessing(inputImage.clone()).to(model.device))
+        output = model(preprocessing(inputImage.clone()).to(device))
     output = output[0] if isinstance(output, (tuple, list)) else output
     return output[0].argmax(dim=0, keepdim=True).cpu()  
 
 
 # %% Gradio interface
 def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     image = hfImageToTensor(image, width=1024, height=512)
-    return image, predict(image, loadModel(selected_model, device))
+    return predict(image, loadModel(selected_model, device))
 
 # Gradio UI
 with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
-    gr.Markdown("## 🧠 Image Segmentation with BiSeNet and BiSeNetV2")
-    gr.Markdown("Upload an image and choose your preferred model for segmentation.")
+    gr.Markdown("## Semantic Segmentation with Real-Time Networks")
     gr.Markdown('A small user interface created to run semantic segmentation on images using city scapes like predictions and real time segmentation networks.')
+    gr.Markdown("Upload an image and choose your preferred model for segmentation.")
 
     with gr.Row():
         with gr.Column():
@@ -42,7 +41,7 @@ with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
                 label="Select model"
             )
             image_input = gr.Image(type="pil", label="Upload image")
-            submit_btn = gr.Button("🧪 Run prediction")
+            submit_btn = gr.Button("Run prediction")
         with gr.Column():
             result_display = gr.Image(label="Model prediction")
 
@@ -54,4 +53,5 @@ with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
     
     gr.Markdown("Made by group 21 semantic segmentation project. ")
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 demo.launch()