refactor: Simplify segmentation to single overlay panel

Changed from 3-panel debug view to clean single overlay:
- Removed: Original image panel (left)
- Removed: Pure mask panel (middle)
- Kept: Overlay panel with contours and legend

Benefits:
- Cleaner, more professional appearance
- Focuses on the actual result (overlay)
- Larger visualization size
- Better for production use

The 3-panel view was useful for debugging "no masks detected"
but now that segmentation works, we only need the overlay.

Co-Authored-By: Claude <[email protected]>

medrax/tools/segmentation/segmentation.py

@@ -173,9 +173,8 @@ class ChestXRaySegmentationTool(BaseTool):
 
     def _save_visualization(self, original_img: np.ndarray, pred_masks: torch.Tensor, organ_indices: List[int]) -> str:
         """Save visualization of original image with segmentation masks overlaid."""
-        # Create a 3-panel figure: original | mask only | overlay
-        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-        ax_orig, ax_mask, ax_overlay = axes
+        # Create single panel with overlay
+        fig, ax = plt.subplots(figsize=(10, 10))
 
         # Generate color palette for organs
         colors = plt.cm.tab10(np.linspace(0, 1, min(len(organ_indices), 10)))
@@ -220,27 +219,10 @@ class ChestXRaySegmentationTool(BaseTool):
 
         print(f"Total masks found and rendered: {masks_found}")
 
-        # Panel 1: Original image
-        ax_orig.imshow(original_img, cmap="gray")
-        ax_orig.set_title("Original X-Ray", fontsize=12, color='white')
-        ax_orig.axis("off")
+        # Display original image
+        ax.imshow(original_img, cmap="gray")
 
-        # Panel 2: Mask only (colorized)
-        if masks_found > 0:
-            # Create RGB image for mask
-            mask_rgb = np.zeros((*original_img.shape, 3))
-            for idx, (organ_name, color) in enumerate(legend_items):
-                mask_region = (combined_mask == idx + 1)
-                mask_rgb[mask_region] = color[:3]
-            ax_mask.imshow(mask_rgb)
-            ax_mask.set_title(f"Segmentation Mask ({masks_found} organ(s))", fontsize=12, color='white')
-        else:
-            ax_mask.imshow(np.zeros_like(original_img), cmap="gray")
-            ax_mask.set_title("No Masks Detected", fontsize=12, color='red')
-        ax_mask.axis("off")
-
-        # Panel 3: Overlay
-        ax_overlay.imshow(original_img, cmap="gray")
+        # Overlay masks with contours
         if masks_found > 0:
             from matplotlib.patches import Patch
             for idx, (organ_name, color) in enumerate(legend_items):
@@ -249,24 +231,22 @@ class ChestXRaySegmentationTool(BaseTool):
                 # Create colored overlay
                 overlay = np.zeros((*original_img.shape, 4))
                 overlay[mask_region] = [color[0], color[1], color[2], 0.5]
-                ax_overlay.imshow(overlay)
+                ax.imshow(overlay)
 
-                # Draw contours
+                # Draw contours for clear boundaries
                 contours = skimage.measure.find_contours(mask_region.astype(float), 0.5)
                 for contour in contours:
-                    ax_overlay.plot(contour[:, 1], contour[:, 0], color=color, linewidth=3, alpha=0.9)
+                    ax.plot(contour[:, 1], contour[:, 0], color=color, linewidth=3, alpha=0.9)
 
-            ax_overlay.set_title("Overlay", fontsize=12, color='white')
-            # Add legend to overlay panel
+            # Add legend
             patches = [Patch(facecolor=c, edgecolor=c, label=n, alpha=0.7) for n, c in legend_items]
-            ax_overlay.legend(handles=patches, loc="upper right", fontsize=8, framealpha=0.9)
+            ax.legend(handles=patches, loc="upper right", fontsize=10, framealpha=0.9)
+            ax.set_title("Segmentation Overlay", fontsize=14, color='white', pad=15)
         else:
-            ax_overlay.set_title("No Overlay (No Masks)", fontsize=12, color='red')
-        ax_overlay.axis("off")
+            ax.set_title("No Masks Detected", fontsize=14, color='red', pad=15)
 
-        # Set background
+        ax.axis("off")
         fig.patch.set_facecolor('black')
-        plt.tight_layout()
 
         save_path = self.temp_dir / f"segmentation_{uuid.uuid4().hex[:8]}.png"
         plt.savefig(save_path, bbox_inches="tight", dpi=150, facecolor='black')