Upload src/pipeline/visualization.py with huggingface_hub

src/pipeline/visualization.py

@@ -0,0 +1,270 @@
+"""
+Visualization module for EL defect analysis.
+
+Creates overlay images with color-coded defect masks:
+- Crack → Blue (visible against bright cell regions)
+- Dark → Red (contrast with bright areas)
+- Cross → Cyan (distinguishable from regular cracks)
+- Busbar → Green (feature, not defect)
+
+All overlays use alpha blending so original image detail remains visible.
+Handles resize alignment to prevent mask/image size mismatches.
+"""
+
+import cv2
+import numpy as np
+from typing import Dict, List, Tuple, Optional
+
+
+# Color scheme (BGR for OpenCV)
+DEFECT_COLORS_BGR = {
+    "background": (0, 0, 0),       # Black (not drawn)
+    "dark": (0, 0, 255),           # Red
+    "crack": (255, 0, 0),          # Blue
+    "cross": (255, 255, 0),        # Cyan
+    "busbar": (0, 255, 0),         # Green
+}
+
+# RGB for matplotlib/PIL/Streamlit
+DEFECT_COLORS_RGB = {
+    "background": (0, 0, 0),
+    "dark": (255, 0, 0),           # Red
+    "crack": (0, 0, 255),          # Blue
+    "cross": (0, 255, 255),        # Cyan
+    "busbar": (0, 255, 0),         # Green
+}
+
+CLASS_NAMES = ["background", "dark", "crack", "cross", "busbar"]
+
+
+def create_overlay(
+    image: np.ndarray,
+    mask: np.ndarray,
+    alpha: float = 0.4,
+    show_background: bool = False,
+) -> np.ndarray:
+    """
+    Create colored overlay of segmentation mask on image.
+    
+    Args:
+        image: Grayscale or BGR image (any size)
+        mask: Class index mask (any size, will be resized to match image)
+        alpha: Overlay transparency (0 = fully transparent, 1 = fully opaque)
+        show_background: If True, also color background class
+        
+    Returns:
+        BGR image with colored overlay
+    """
+    # Ensure image is BGR
+    if image.ndim == 2:
+        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    else:
+        vis = image.copy()
+    
+    # Ensure uint8
+    if vis.dtype != np.uint8:
+        if vis.max() <= 1.0:
+            vis = (vis * 255).astype(np.uint8)
+        else:
+            vis = vis.astype(np.uint8)
+    
+    h, w = vis.shape[:2]
+    
+    # Resize mask to match image (CRITICAL: use NEAREST to preserve labels)
+    if mask.shape[:2] != (h, w):
+        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    
+    # Create color overlay
+    overlay = vis.copy()
+    
+    for class_idx, class_name in enumerate(CLASS_NAMES):
+        if class_idx == 0 and not show_background:
+            continue
+        
+        color = DEFECT_COLORS_BGR[class_name]
+        class_mask = mask == class_idx
+        
+        if class_mask.any():
+            overlay[class_mask] = color
+    
+    # Alpha blend
+    result = cv2.addWeighted(vis, 1 - alpha, overlay, alpha, 0)
+    
+    return result
+
+
+def create_class_overlay(
+    image: np.ndarray,
+    mask: np.ndarray,
+    class_name: str,
+    alpha: float = 0.5,
+    color: Optional[Tuple[int, int, int]] = None,
+) -> np.ndarray:
+    """
+    Create overlay for a single class.
+    
+    Args:
+        image: Grayscale or BGR image
+        mask: Binary mask for one class
+        class_name: For color lookup
+        alpha: Overlay transparency
+        color: Override color (BGR)
+    """
+    if image.ndim == 2:
+        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    else:
+        vis = image.copy()
+    
+    if vis.dtype != np.uint8:
+        vis = (vis * 255).astype(np.uint8) if vis.max() <= 1 else vis.astype(np.uint8)
+    
+    h, w = vis.shape[:2]
+    if mask.shape[:2] != (h, w):
+        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    
+    if color is None:
+        color = DEFECT_COLORS_BGR.get(class_name, (255, 255, 255))
+    
+    overlay = vis.copy()
+    overlay[mask > 0] = color
+    
+    return cv2.addWeighted(vis, 1 - alpha, overlay, alpha, 0)
+
+
+def create_color_mask(
+    mask: np.ndarray,
+    include_background: bool = False,
+) -> np.ndarray:
+    """
+    Convert class index mask to RGB color visualization.
+    
+    Returns:
+        (H, W, 3) uint8 RGB image
+    """
+    h, w = mask.shape[:2]
+    color_img = np.zeros((h, w, 3), dtype=np.uint8)
+    
+    for class_idx, class_name in enumerate(CLASS_NAMES):
+        if class_idx == 0 and not include_background:
+            continue
+        
+        color = DEFECT_COLORS_RGB[class_name]
+        color_img[mask == class_idx] = color
+    
+    return color_img
+
+
+def draw_cell_results(
+    image: np.ndarray,
+    cell_results: List[dict],
+    cells: list,
+) -> np.ndarray:
+    """
+    Draw cell analysis results on module image.
+    
+    Shows per-cell:
+    - Bounding box (green = PASS, red = FAIL)
+    - Cell ID
+    - Defect score
+    """
+    if image.ndim == 2:
+        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    else:
+        vis = image.copy()
+    
+    if vis.dtype != np.uint8:
+        vis = (vis * 255).astype(np.uint8) if vis.max() <= 1 else vis.astype(np.uint8)
+    
+    for cell_info, result in zip(cells, cell_results):
+        y1, x1, y2, x2 = cell_info.bbox
+        score = result.get("defect_score", 0)
+        
+        # Color: green for good, yellow for moderate, red for bad
+        if score < 25:
+            color = (0, 255, 0)  # Green
+        elif score < 50:
+            color = (0, 255, 255)  # Yellow
+        else:
+            color = (0, 0, 255)  # Red
+        
+        cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)
+        
+        # Label
+        label = f"C{cell_info.cell_id}: {score:.0f}"
+        cv2.putText(
+            vis, label, (x1 + 2, y1 + 15),
+            cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1
+        )
+    
+    return vis
+
+
+def create_summary_image(
+    original: np.ndarray,
+    overlay: np.ndarray,
+    mask_color: np.ndarray,
+    decision: str,
+    score: float,
+) -> np.ndarray:
+    """
+    Create a summary image with original, overlay, and color mask side by side.
+    
+    Returns:
+        (H, W*3, 3) BGR image with all three panels
+    """
+    # Ensure all are BGR
+    panels = []
+    for img in [original, overlay, mask_color]:
+        if img.ndim == 2:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        if img.dtype != np.uint8:
+            img = (img * 255).astype(np.uint8) if img.max() <= 1 else img.astype(np.uint8)
+        panels.append(img)
+    
+    # Resize to same height
+    target_h = 400
+    resized = []
+    for p in panels:
+        scale = target_h / p.shape[0]
+        new_w = int(p.shape[1] * scale)
+        resized.append(cv2.resize(p, (new_w, target_h)))
+    
+    # Concatenate horizontally
+    summary = np.hstack(resized)
+    
+    # Add decision text
+    color = (0, 255, 0) if decision == "PASS" else (0, 0, 255)
+    text = f"{decision} (Score: {score:.1f})"
+    cv2.putText(
+        summary, text, (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2
+    )
+    
+    return summary
+
+
+def create_legend(height: int = 400, width: int = 200) -> np.ndarray:
+    """Create a color legend for defect classes."""
+    legend = np.ones((height, width, 3), dtype=np.uint8) * 255
+    
+    y_offset = 30
+    for class_name in CLASS_NAMES[1:]:  # Skip background
+        color = DEFECT_COLORS_BGR[class_name]
+        
+        # Color swatch
+        cv2.rectangle(
+            legend, (10, y_offset), (40, y_offset + 20), color, -1
+        )
+        cv2.rectangle(
+            legend, (10, y_offset), (40, y_offset + 20), (0, 0, 0), 1
+        )
+        
+        # Label
+        cv2.putText(
+            legend, class_name.capitalize(), (50, y_offset + 15),
+            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1
+        )
+        
+        y_offset += 35
+    
+    return legend