updated app

app.py

@@ -1,10 +1,10 @@
+import random
 import gradio as gr
 from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
 from torchvision.transforms import ColorJitter, functional as F
 from PIL import Image, ImageDraw, ImageFont
 import numpy as np
 import torch
-import torch.nn as nn
 from datasets import load_dataset
 import evaluate
 
@@ -18,9 +18,10 @@ lora_model_id = "guimCC/segformer-v0-gta-cityscapes"
 original_model = SegformerForSemanticSegmentation.from_pretrained(original_model_id).to(device)
 lora_model = SegformerForSemanticSegmentation.from_pretrained(lora_model_id).to(device)
 
-# Load the dataset and slice it
+# Load the dataset and select 10 random images
 dataset = load_dataset("Chris1/cityscapes", split="validation")
-sampled_dataset = [dataset[i] for i in range(10)]  # Select the first 10 examples
+#sampled_dataset = random.sample(list(dataset), 10)  # Select 10 random examples
+sampled_dataset = dataset[:10]  # Select the first 10 examples
 
 # Define your custom image processor
 jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
@@ -66,7 +67,7 @@ def compute_miou(logits, labels):
     with torch.no_grad():
         logits_tensor = torch.from_numpy(logits)
         # Scale the logits to the size of the label
-        logits_tensor = nn.functional.interpolate(
+        logits_tensor = F.interpolate(
             logits_tensor,
             size=labels.shape[-2:],
             mode="bilinear",
@@ -90,9 +91,13 @@ def compute_miou(logits, labels):
             reduce_labels=processor.do_reduce_labels,
         )
         
-        return metrics['mean_iou']
-
-
+        mean_iou = metrics.get('mean_iou', 0.0)
+        
+        if np.isnan(mean_iou):
+            mean_iou = 0.0  # Handle NaN values gracefully
+        
+        return mean_iou
+    
 def apply_color_palette(segmentation):
     colored_segmentation = palette[segmentation]
     return Image.fromarray(colored_segmentation.astype(np.uint8))
@@ -123,9 +128,7 @@ def create_legend():
 
     return legend
 
-
-
-def inference(index, a):
+def inference(index, legend):
     """Run inference on the input image with both models."""
     image = sampled_dataset[index]['image']  # Fetch image from the sampled dataset
     pixel_values = preprocess_image(image)
@@ -140,28 +143,15 @@ def inference(index, a):
         lora_outputs = lora_model(pixel_values=pixel_values)
         lora_segmentation = postprocess_predictions(lora_outputs.logits)
 
-    # Compute mIoU
-    true_labels = np.array(sampled_dataset[index]['semantic_segmentation'])
-    original_miou = compute_miou(original_outputs.logits.detach().cpu().numpy(), true_labels)
-    lora_miou = compute_miou(lora_outputs.logits.detach().cpu().numpy(), true_labels)
-    # original_miou = 0
-    # lora_miou = 0
-
     # Apply color palette
     original_segmentation_image = apply_color_palette(original_segmentation)
     lora_segmentation_image = apply_color_palette(lora_segmentation)
 
-   # Create legend
-    legend = create_legend()
-
     # Return the original image, the segmentations, and mIoU
     return (
         image,
         original_segmentation_image,
         lora_segmentation_image,
-        legend,
-        f"Original Model mIoU: {original_miou:.2f}",
-        f"LoRA Model mIoU: {lora_miou:.2f}"
     )
 
 # Create a list of image options for the user to select from
@@ -175,15 +165,13 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Legend", value=create_legend)
     ],
     outputs=[
-        gr.Image(type="pil", label="Selected Image"),
-        gr.Image(type="pil", label="Original Model Output"),
-        gr.Image(type="pil", label="LoRA Model Output"),
-        gr.Textbox(label="Original Model mIoU"),
-        gr.Textbox(label="LoRA Model mIoU")
+        gr.Image(type="pil", label="Input Image"),
+        gr.Image(type="pil", label="Original Model Prediction"),
+        gr.Image(type="pil", label="LoRA Model Prediction"),
+
     ],
-    title="Segformer Cityscapes Inference",
-    description="Select an image from the Cityscapes dataset to see the segmentation results from both the original and fine-tuned Segformer models.",
+    live=True
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()