Update app.py

app.py

@@ -105,7 +105,7 @@ def load_model(model_name: str):
 # ---------------------------
 
 def generate_heatmap_layout(img_rgb: np.ndarray, conf: float = 0.25):
-    """Return a composite XAI heatmap layout comparing vanilla and fine-tuned models."""
+    """Return separate XAI heatmap layouts for vanilla and fine-tuned models."""
 
     # Convert RGB (Gradio default) ➜ BGR (OpenCV default)
     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
@@ -128,9 +128,6 @@ def generate_heatmap_layout(img_rgb: np.ndarray, conf: float = 0.25):
     finetune_has_detections = (finetune_results and len(finetune_results) > 0 and 
                               hasattr(finetune_results[0], "boxes") and len(finetune_results[0].boxes) > 0)
     
-    if not vanilla_has_detections and not finetune_has_detections:
-        return img_rgb  # nothing detected, return original
-
     # Create heatmap visualizations for both models
     vanilla_heatmaps = []
     finetune_heatmaps = []
@@ -145,8 +142,15 @@ def generate_heatmap_layout(img_rgb: np.ndarray, conf: float = 0.25):
         finetune_names = [finetune_model.model.names[int(cls)] for cls in finetune_result.boxes.cls]
         finetune_heatmaps = create_heatmap_snippets(img_bgr, finetune_result, finetune_names, "Fine-tuned")
 
-    # Create comparison layout
-    return create_comparison_layout(img_bgr, vanilla_heatmaps, finetune_heatmaps)
+    # Create separate layouts for each model
+    vanilla_layout = create_model_layout(vanilla_heatmaps, "Vanilla Model", (0, 100, 0))
+    finetune_layout = create_model_layout(finetune_heatmaps, "Fine-tuned Model", (0, 0, 200))
+    
+    # Convert BGR to RGB for display
+    vanilla_output = cv2.cvtColor(vanilla_layout, cv2.COLOR_BGR2RGB) if vanilla_layout is not None else None
+    finetune_output = cv2.cvtColor(finetune_layout, cv2.COLOR_BGR2RGB) if finetune_layout is not None else None
+    
+    return vanilla_output, finetune_output
 
 
 def create_heatmap_snippets(img_bgr, result, names, model_type):
@@ -181,38 +185,10 @@ def create_heatmap_snippets(img_bgr, result, names, model_type):
     return snippets
 
 
-def create_comparison_layout(img_bgr, vanilla_heatmaps, finetune_heatmaps):
-    """Create side-by-side comparison layout for both models."""
+def create_model_layout(heatmaps, title, color):
+    """Create a layout for one model's heatmaps."""
     pad = 20
     
-    # Create model sections
-    vanilla_section = create_model_section(vanilla_heatmaps, "Vanilla Model", (0, 100, 0))
-    finetune_section = create_model_section(finetune_heatmaps, "Fine-tuned Model", (0, 0, 200))
-    
-    # Calculate dimensions
-    max_height = max(vanilla_section.shape[0], finetune_section.shape[0]) if vanilla_section.size > 0 and finetune_section.size > 0 else max(vanilla_section.shape[0] if vanilla_section.size > 0 else 0, finetune_section.shape[0] if finetune_section.size > 0 else 0)
-    total_width = vanilla_section.shape[1] + finetune_section.shape[1] + pad if vanilla_section.size > 0 and finetune_section.size > 0 else max(vanilla_section.shape[1] if vanilla_section.size > 0 else 0, finetune_section.shape[1] if finetune_section.size > 0 else 0)
-    
-    # Create final canvas
-    canvas_h = max_height + 2 * pad
-    canvas_w = total_width + 2 * pad
-    final = np.full((canvas_h, canvas_w, 3), 255, np.uint8)
-    
-    # Place sections side by side
-    if vanilla_section.size > 0:
-        y_off = (canvas_h - vanilla_section.shape[0]) // 2
-        final[y_off : y_off + vanilla_section.shape[0], pad : pad + vanilla_section.shape[1]] = vanilla_section
-    
-    if finetune_section.size > 0:
-        y_off = (canvas_h - finetune_section.shape[0]) // 2
-        x_start = pad + vanilla_section.shape[1] + pad if vanilla_section.size > 0 else pad
-        final[y_off : y_off + finetune_section.shape[0], x_start : x_start + finetune_section.shape[1]] = finetune_section
-    
-    return cv2.cvtColor(final, cv2.COLOR_BGR2RGB)
-
-
-def create_model_section(heatmaps, title, color):
-    """Create a section for one model's heatmaps."""
     if not heatmaps:
         # Create empty section with title
         font = cv2.FONT_HERSHEY_SIMPLEX
@@ -231,19 +207,22 @@ def create_model_section(heatmaps, title, color):
     section_h = max_h + th + 40
     section_w = max(total_w, tw + 20)
     
-    section = np.full((section_h, section_w, 3), 255, np.uint8)
+    # Create canvas with padding
+    canvas_h = section_h + 2 * pad
+    canvas_w = section_w + 2 * pad
+    canvas = np.full((canvas_h, canvas_w, 3), 255, np.uint8)
     
     # Add title
-    cv2.putText(section, title, (10, th + 20), title_font, 1.0, color, 2, cv2.LINE_AA)
+    cv2.putText(canvas, title, (pad + 10, pad + th + 20), title_font, 1.0, color, 2, cv2.LINE_AA)
     
     # Arrange heatmaps
-    cur_x = 0
+    cur_x = pad
     for h in heatmaps:
-        y_off = th + 30 + (max_h - h.shape[0]) // 2
-        section[y_off : y_off + h.shape[0], cur_x : cur_x + h.shape[1]] = h
+        y_off = pad + th + 30 + (max_h - h.shape[0]) // 2
+        canvas[y_off : y_off + h.shape[0], cur_x : cur_x + h.shape[1]] = h
         cur_x += h.shape[1] + 10
     
-    return section
+    return canvas
 
 
 # ---------------------------
@@ -267,19 +246,20 @@ def build_demo():
                     label="Click to load example"
                 )
             
-            # Right side - Output visualization
+            # Right side - Output visualizations (separated vertically)
             with gr.Column(scale=2):
-                outputs = gr.Image(type="numpy", label="XAI Heatmap Comparison")
+                vanilla_output = gr.Image(type="numpy", label="Vanilla Model Heatmap")
+                finetune_output = gr.Image(type="numpy", label="Fine-tuned Model Heatmap")
         
         # Connect inputs to the function
         def update_heatmap(image, confidence):
             if image is None:
-                return None
+                return None, None
             return generate_heatmap_layout(image, confidence)
         
         # Set up the interface
-        image_input.change(fn=update_heatmap, inputs=[image_input, conf_input], outputs=outputs)
-        conf_input.change(fn=update_heatmap, inputs=[image_input, conf_input], outputs=outputs)
+        image_input.change(fn=update_heatmap, inputs=[image_input, conf_input], outputs=[vanilla_output, finetune_output])
+        conf_input.change(fn=update_heatmap, inputs=[image_input, conf_input], outputs=[vanilla_output, finetune_output])
     
     return demo