added post processing

app.py

@@ -1,6 +1,6 @@
 import os, torch
 import gradio as gr
-from utils.imageHandling import hfImageToTensor, preprocessing
+from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing
 from model.modelLoading import loadModel
 
 # %% prediction on an image
@@ -24,8 +24,20 @@ def predict(inputImage: torch.Tensor, model) -> torch.Tensor:
 
 # %% Gradio interface
 def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
+    if image is None:
+        raise ValueError("No image provided for prediction.")
+    
+    if selected_model is None:
+        raise ValueError("No model selected for prediction.")
+    
+    try:
+        model = loadModel(selected_model, device) 
+    except NotImplementedError as e:
+        raise ValueError(f"Model loading failed: {e}")
+    
     image = hfImageToTensor(image, width=1024, height=512)
-    return predict(image, loadModel(selected_model, device))
+    prediction = predict(image, model)
+    return postprocessing(prediction)
 
 # Gradio UI
 with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
@@ -44,7 +56,10 @@ with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
             submit_btn = gr.Button("Run prediction")
         with gr.Column():
             result_display = gr.Image(label="Model prediction")
-
+    
+    result_display = gr.Image(label="Output")
+    error_text = gr.Textbox(label="Messaggio", interactive=False)
+    
     submit_btn.click(
         fn=run_prediction,
         inputs=[image_input, model_selector],

utils/imageHandling.py

@@ -31,3 +31,16 @@ def preprocessing(image_tensor: torch.Tensor) -> torch.Tensor:
     return torchvision.transforms.functional.normalize(
         image_tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
     ).unsqueeze(0)
+
+# %% postprocessing
+def postprocessing(pred: torch.Tensor) -> torch.Tensor:
+    """
+    Convert the model's output tensor to a format suitable for visualization.
+
+    Args:
+        pred (torch.Tensor): Model output tensor of shape (1, H, W).
+
+    Returns:
+        torch.Tensor: Processed tensor of shape (3, H, W) for visualization.
+    """ 
+    return torchvision.transforms.functional.to_pil_image(pred.squeeze(0).cpu().clamp(0, 1))