Upload src/pipeline/full_pipeline.py with huggingface_hub

src/pipeline/full_pipeline.py

@@ -0,0 +1,358 @@
+"""
+Full Pipeline Orchestrator.
+
+Connects all pipeline stages into a single coherent system:
+
+input image
+   ↓
+preprocessing (normalize brightness)
+   ↓
+module segmentation (grid detection)
+   ↓
+cell extraction
+   ↓
+deep learning inference (U-Net)
+   ↓
+post-processing (clean masks)
+   ↓
+defect analysis
+   ↓
+decision engine
+   ↓
+visualization + results
+
+This is the main entry point for the entire system.
+All error handling and fallback logic is centralized here.
+"""
+
+import cv2
+import numpy as np
+import time
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+
+from .preprocessing import ELPreprocessor
+from .module_segmentation import ModuleSegmenter, CellInfo, estimate_pixel_to_mm
+from .inference import DefectInferenceEngine, MockInferenceEngine
+from .mask_cleaning import MaskCleaner
+from .crack_analysis import DefectAnalyzer, CellAnalysisResult
+from .decision_engine import DecisionEngine, ModuleDecision
+from .visualization import (
+    create_overlay, create_color_mask, draw_cell_results,
+    create_summary_image, DEFECT_COLORS_RGB, CLASS_NAMES
+)
+
+
+@dataclass
+class PipelineResult:
+    """Complete pipeline output."""
+    # Input
+    original_image: np.ndarray
+    preprocessed_image: np.ndarray
+    
+    # Segmentation
+    cells: List[CellInfo]
+    num_cells: int
+    grid_image: Optional[np.ndarray] = None
+    
+    # Per-cell results
+    cell_masks: List[np.ndarray] = field(default_factory=list)
+    cell_overlays: List[np.ndarray] = field(default_factory=list)
+    cell_analyses: List[CellAnalysisResult] = field(default_factory=list)
+    
+    # Module-level
+    module_mask: Optional[np.ndarray] = None
+    module_overlay: Optional[np.ndarray] = None
+    decision: Optional[ModuleDecision] = None
+    
+    # Timing
+    timings: Dict[str, float] = field(default_factory=dict)
+    
+    # Errors/warnings
+    warnings: List[str] = field(default_factory=list)
+    
+    def to_dict(self) -> dict:
+        return {
+            "num_cells": self.num_cells,
+            "decision": self.decision.to_dict() if self.decision else None,
+            "cell_analyses": [c.to_dict() for c in self.cell_analyses],
+            "timings": {k: round(v, 3) for k, v in self.timings.items()},
+            "warnings": self.warnings,
+        }
+
+
+class ELInspectionPipeline:
+    """
+    Production-grade EL defect inspection pipeline.
+    
+    Usage:
+        pipeline = ELInspectionPipeline(model_path="model.pth")
+        result = pipeline.run(image)
+        
+        # Access results
+        print(result.decision.decision)  # "PASS" or "FAIL"
+        print(result.decision.overall_score)  # 0-100
+    """
+    
+    def __init__(
+        self,
+        model_path: Optional[str] = None,
+        encoder_name: str = "resnet34",
+        input_size: int = 512,
+        device: Optional[str] = None,
+        cell_type: str = "standard",
+        quality_standard: str = "default",
+        use_mock: bool = False,
+    ):
+        """
+        Args:
+            model_path: Path to trained model weights
+            encoder_name: Model encoder name
+            input_size: Model input size
+            device: 'cuda', 'cpu', or None for auto
+            cell_type: Solar cell type for mm conversion
+            quality_standard: 'default', 'strict', or 'lenient'
+            use_mock: Use mock inference (for testing without model)
+        """
+        self.input_size = input_size
+        self.cell_type = cell_type
+        
+        # Initialize components
+        self.preprocessor = ELPreprocessor(
+            target_size=(input_size, input_size)
+        )
+        self.segmenter = ModuleSegmenter()
+        
+        if use_mock or model_path is None:
+            self.engine = MockInferenceEngine(input_size=input_size)
+            self._using_mock = True
+        else:
+            try:
+                self.engine = DefectInferenceEngine(
+                    model_path=model_path,
+                    encoder_name=encoder_name,
+                    device=device,
+                    input_size=input_size,
+                )
+                self._using_mock = False
+            except Exception as e:
+                print(f"WARNING: Failed to load model: {e}. Using mock inference.")
+                self.engine = MockInferenceEngine(input_size=input_size)
+                self._using_mock = True
+        
+        self.mask_cleaner = MaskCleaner()
+        self.decision_engine = DecisionEngine(standard=quality_standard)
+    
+    def run(
+        self,
+        image: np.ndarray,
+        custom_thresholds: Optional[Dict] = None,
+    ) -> PipelineResult:
+        """
+        Run the full inspection pipeline.
+        
+        Args:
+            image: Input EL image (any format/size)
+            custom_thresholds: Override decision thresholds
+            
+        Returns:
+            PipelineResult with all outputs
+        """
+        timings = {}
+        warnings = []
+        
+        # ── Step 1: Preprocessing ─────────────────────────────
+        t0 = time.time()
+        try:
+            preprocessed = self.preprocessor.process(image)
+            stats = self.preprocessor.get_image_stats(image)
+            
+            if stats["is_dark"]:
+                warnings.append("Input image is very dark — results may be less reliable")
+            if stats["is_low_contrast"]:
+                warnings.append("Input image has low contrast — enhanced with CLAHE")
+        except Exception as e:
+            warnings.append(f"Preprocessing error: {e}. Using raw image.")
+            preprocessed = self._safe_grayscale(image)
+        
+        timings["preprocessing"] = time.time() - t0
+        
+        # ── Step 2: Module Segmentation ───────────────────────
+        t0 = time.time()
+        try:
+            # Use preprocessed image for grid detection (better contrast)
+            gray_uint8 = (preprocessed * 255).astype(np.uint8)
+            cells = self.segmenter.segment(gray_uint8)
+            grid_image = self.segmenter.get_grid_visualization(gray_uint8, cells)
+        except Exception as e:
+            warnings.append(f"Grid detection failed: {e}. Treating as single cell.")
+            h, w = preprocessed.shape[:2]
+            cells = [CellInfo(
+                cell_id=1, row=0, col=0,
+                image=gray_uint8 if 'gray_uint8' in dir() else (preprocessed * 255).astype(np.uint8),
+                bbox=(0, 0, h, w), area_pixels=h * w
+            )]
+            grid_image = None
+        
+        timings["segmentation"] = time.time() - t0
+        
+        # ── Step 3: Estimate pixel-to-mm scale ────────────────
+        if len(cells) > 1:
+            # Use median cell size for calibration
+            cell_widths = [c.bbox[3] - c.bbox[1] for c in cells]
+            cell_heights = [c.bbox[2] - c.bbox[0] for c in cells]
+            median_w = int(np.median(cell_widths))
+            median_h = int(np.median(cell_heights))
+            px_per_mm_factor = estimate_pixel_to_mm(median_w, median_h, self.cell_type)
+            px_per_mm = 1.0 / px_per_mm_factor if px_per_mm_factor > 0 else 5.0
+        else:
+            # Single cell: estimate from image size
+            h, w = preprocessed.shape[:2]
+            px_per_mm_factor = estimate_pixel_to_mm(w, h, self.cell_type)
+            px_per_mm = 1.0 / px_per_mm_factor if px_per_mm_factor > 0 else 5.0
+        
+        # ── Step 4: Per-cell inference + analysis ─────────────
+        t0 = time.time()
+        cell_masks = []
+        cell_overlays = []
+        cell_analyses = []
+        
+        defect_analyzer = DefectAnalyzer(px_per_mm=px_per_mm)
+        
+        for cell in cells:
+            try:
+                # Run inference on cell image
+                cell_img = cell.image
+                
+                # Get mask
+                mask = self.engine.predict(cell_img)
+                
+                # Clean mask
+                cleaned_mask = self.mask_cleaner.clean(mask)
+                cell_masks.append(cleaned_mask)
+                
+                # Create overlay
+                overlay = create_overlay(cell_img, cleaned_mask, alpha=0.4)
+                cell_overlays.append(overlay)
+                
+                # Analyze defects
+                # Resize cell image to match mask size for analysis
+                if cell_img.shape[:2] != cleaned_mask.shape[:2]:
+                    cell_img_resized = cv2.resize(
+                        cell_img, 
+                        (cleaned_mask.shape[1], cleaned_mask.shape[0]),
+                        interpolation=cv2.INTER_LINEAR
+                    )
+                else:
+                    cell_img_resized = cell_img
+                
+                analysis = defect_analyzer.analyze_cell(
+                    cell_img_resized, cleaned_mask, cell.cell_id
+                )
+                cell_analyses.append(analysis)
+                
+            except Exception as e:
+                warnings.append(f"Cell {cell.cell_id} analysis failed: {e}")
+                # Create empty result
+                from .crack_analysis import DarkResult
+                cell_masks.append(np.zeros_like(mask if 'mask' in dir() else np.zeros((self.input_size, self.input_size), dtype=np.uint8)))
+                cell_overlays.append(cell.image)
+                cell_analyses.append(CellAnalysisResult(
+                    cell_id=cell.cell_id, cracks=[], dark=DarkResult(0,0,0,0,0,"none"),
+                    cross_cracks=[], total_crack_length_mm=0, num_cracks=0,
+                    num_cross_cracks=0, max_crack_severity="none", defect_score=0
+                ))
+        
+        timings["inference_analysis"] = time.time() - t0
+        
+        # ── Step 5: Compose module-level mask ─────────────────
+        t0 = time.time()
+        try:
+            module_mask = self._compose_module_mask(
+                preprocessed.shape, cells, cell_masks
+            )
+            module_overlay = create_overlay(
+                (preprocessed * 255).astype(np.uint8),
+                module_mask, alpha=0.4
+            )
+        except Exception as e:
+            warnings.append(f"Module mask composition failed: {e}")
+            module_mask = None
+            module_overlay = None
+        
+        timings["visualization"] = time.time() - t0
+        
+        # ── Step 6: Module-level decision ─────────────────────
+        t0 = time.time()
+        
+        if custom_thresholds:
+            self.decision_engine.update_thresholds(**custom_thresholds)
+        
+        decision = self.decision_engine.decide(cell_analyses)
+        timings["decision"] = time.time() - t0
+        
+        if self._using_mock:
+            warnings.append("Using mock inference — results are approximate. Load a trained model for production use.")
+        
+        timings["total"] = sum(timings.values())
+        
+        return PipelineResult(
+            original_image=image,
+            preprocessed_image=preprocessed,
+            cells=cells,
+            num_cells=len(cells),
+            grid_image=grid_image,
+            cell_masks=cell_masks,
+            cell_overlays=cell_overlays,
+            cell_analyses=cell_analyses,
+            module_mask=module_mask,
+            module_overlay=module_overlay,
+            decision=decision,
+            timings=timings,
+            warnings=warnings,
+        )
+    
+    def _compose_module_mask(
+        self, 
+        shape: Tuple, 
+        cells: List[CellInfo], 
+        cell_masks: List[np.ndarray]
+    ) -> np.ndarray:
+        """
+        Compose per-cell masks back into a module-level mask.
+        
+        Handles size mismatches between cell extraction and model output.
+        """
+        h, w = shape[:2]
+        module_mask = np.zeros((h, w), dtype=np.uint8)
+        
+        for cell, mask in zip(cells, cell_masks):
+            y1, x1, y2, x2 = cell.bbox
+            cell_h, cell_w = y2 - y1, x2 - x1
+            
+            # Resize mask to cell bbox size
+            if mask.shape[:2] != (cell_h, cell_w):
+                resized = cv2.resize(
+                    mask, (cell_w, cell_h),
+                    interpolation=cv2.INTER_NEAREST
+                )
+            else:
+                resized = mask
+            
+            # Place in module mask
+            module_mask[y1:y2, x1:x2] = resized
+        
+        return module_mask
+    
+    def _safe_grayscale(self, image: np.ndarray) -> np.ndarray:
+        """Safely convert image to grayscale float32 [0, 1]."""
+        if image.ndim == 3:
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        else:
+            gray = image
+        
+        if gray.dtype == np.uint8:
+            return gray.astype(np.float32) / 255.0
+        elif gray.max() > 1.0:
+            return gray.astype(np.float32) / 255.0
+        return gray.astype(np.float32)