sorghum_pipeline/pipeline.py

"""
Minimal single-image pipeline for Hugging Face demo.
"""

import logging
from pathlib import Path
from typing import Dict, Any, Callable, Optional, Generator
import numpy as np
import cv2

from .config import Config
from .data import ImagePreprocessor, MaskHandler
from .features import TextureExtractor, VegetationIndexExtractor, MorphologyExtractor
from .output import OutputManager
from .segmentation import SegmentationManager

logger = logging.getLogger(__name__)


class SorghumPipeline:
    """Minimal pipeline for single-image processing."""
    
    def __init__(self, config: Config):
        """Initialize pipeline."""
        logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
        self.config = config
        self.config.validate()
        
        # Initialize components with defaults
        self.preprocessor = ImagePreprocessor()
        self.mask_handler = MaskHandler()
        self.texture_extractor = TextureExtractor()
        self.vegetation_extractor = VegetationIndexExtractor()
        self.morphology_extractor = MorphologyExtractor()
        self.segmentation_manager = SegmentationManager(
            model_name="briaai/RMBG-2.0",
            device=self.config.get_device(),
            trust_remote_code=True
        )
        self.output_manager = OutputManager(
            output_folder=self.config.paths.output_folder,
            settings=self.config.output
        )
        logger.info("Pipeline initialized")

    def run(self, single_image_path: str) -> Dict[str, Any]:
        """Run pipeline on single image."""
        logger.info("Processing single image...")

        import time, imghdr, tifffile
        from PIL import Image

        start = time.perf_counter()

        # --- Load image with TIFF preference ---
        kind = imghdr.what(single_image_path)
        suffix = Path(single_image_path).suffix.lower()

        arr = None
        if kind == "tiff" or suffix in [".tif", ".tiff"]:
            try:
                arr = tifffile.imread(single_image_path)
                logger.info(f"Loaded TIFF: shape={arr.shape}, dtype={arr.dtype}")
            except Exception as e:
                logger.warning(f"tifffile failed ({e}), falling back to cv2")
                arr = cv2.imread(single_image_path, cv2.IMREAD_UNCHANGED)
                logger.info(f"Fallback read: shape={arr.shape}, dtype={arr.dtype}")
        else:
            arr = cv2.imread(single_image_path, cv2.IMREAD_UNCHANGED)
            logger.info(f"Loaded non-TIFF: shape={arr.shape}, dtype={arr.dtype}")

        # --- Normalize array shape ---
        if arr is None:
            raise ValueError(f"Could not read image: {single_image_path}")
        if arr.ndim > 3:
            arr = arr[..., 0]  # drop extra dimension
        if arr.ndim == 3 and arr.shape[-1] == 1:
            arr = arr[..., 0]  # squeeze singleton

        logger.info(f"DEBUG normalized input: shape={arr.shape}, dtype={arr.dtype}")

        # Wrap into PIL image for downstream pipeline
        img = Image.fromarray(arr)

        plants = {
            "demo": {
                "raw_image": (img, Path(single_image_path).name),
                "plant_name": "demo",
                # Keep original normalized array for visualization
                "normalized_input": arr,
            }
        }

        # Process: composite → segment → features → save
        plants = self.preprocessor.create_composites(plants)
        plants = self._segment(plants)
        plants = self._extract_features(plants)
        self.output_manager.create_output_directories()

        for key, pdata in plants.items():
            self.output_manager.save_plant_results(key, pdata)

        elapsed = time.perf_counter() - start
        logger.info(f"Completed in {elapsed:.2f}s")

        return {"plants": plants, "timing": elapsed}

    def run_with_progress(self, single_image_path: str, progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None) -> Generator[Dict[str, Any], None, None]:
        """Run pipeline on single image, yielding intermediate results progressively."""
        logger.info("Processing single image with progress...")

        import time, imghdr, tifffile
        from PIL import Image

        start = time.perf_counter()

        # --- Load image with TIFF preference ---
        kind = imghdr.what(single_image_path)
        suffix = Path(single_image_path).suffix.lower()

        arr = None
        if kind == "tiff" or suffix in [".tif", ".tiff"]:
            try:
                arr = tifffile.imread(single_image_path)
                logger.info(f"Loaded TIFF: shape={arr.shape}, dtype={arr.dtype}")
            except Exception as e:
                logger.warning(f"tifffile failed ({e}), falling back to cv2")
                arr = cv2.imread(single_image_path, cv2.IMREAD_UNCHANGED)
                logger.info(f"Fallback read: shape={arr.shape}, dtype={arr.dtype}")
        else:
            arr = cv2.imread(single_image_path, cv2.IMREAD_UNCHANGED)
            logger.info(f"Loaded non-TIFF: shape={arr.shape}, dtype={arr.dtype}")

        # --- Normalize array shape ---
        if arr is None:
            raise ValueError(f"Could not read image: {single_image_path}")
        if arr.ndim > 3:
            arr = arr[..., 0]  # drop extra dimension
        if arr.ndim == 3 and arr.shape[-1] == 1:
            arr = arr[..., 0]  # squeeze singleton

        logger.info(f"DEBUG normalized input: shape={arr.shape}, dtype={arr.dtype}")

        # Wrap into PIL image for downstream pipeline
        img = Image.fromarray(arr)

        plants = {
            "demo": {
                "raw_image": (img, Path(single_image_path).name),
                "plant_name": "demo",
                "normalized_input": arr,
            }
        }

        # Create output directories early
        self.output_manager.create_output_directories()

        # Stage 1: Create composite
        logger.info("Stage 1: Creating composite...")
        plants = self.preprocessor.create_composites(plants)
        if progress_callback:
            progress_callback("composite", plants)
        # Save composite immediately for display
        for key, pdata in plants.items():
            self.output_manager.save_plant_results(key, pdata)
        yield {"plants": plants, "stage": "composite"}

        # Stage 2: Segmentation
        logger.info("Stage 2: Segmentation...")
        plants = self._segment(plants)
        if progress_callback:
            progress_callback("segmentation", plants)
        # Save mask and overlay
        for key, pdata in plants.items():
            self.output_manager.save_plant_results(key, pdata)
        yield {"plants": plants, "stage": "segmentation"}

        # Stage 3: Extract features (texture, vegetation, morphology)
        logger.info("Stage 3: Extracting features...")
        plants = self._extract_features(plants)
        if progress_callback:
            progress_callback("features", plants)
        # Save all final outputs
        for key, pdata in plants.items():
            self.output_manager.save_plant_results(key, pdata)
        
        elapsed = time.perf_counter() - start
        logger.info(f"Completed in {elapsed:.2f}s")

        yield {"plants": plants, "timing": elapsed, "stage": "complete"}

    def _segment(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Segment using BRIA."""
        for key, pdata in plants.items():
            composite = pdata['composite']
            logger.info(f"Composite shape: {composite.shape}")
            soft_mask = self.segmentation_manager.segment_image_soft(composite)
            logger.info(f"Soft mask shape: {soft_mask.shape}")
            mask_uint8 = (soft_mask * 255.0).astype(np.uint8)
            logger.info(f"Mask uint8 shape: {mask_uint8.shape}")
            pdata['mask'] = mask_uint8
        return plants
    
    def _extract_features(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Extract features: texture + vegetation indices + morphology."""
        for key, pdata in plants.items():
            composite = pdata['composite']
            mask = pdata.get('mask')

            # --- Texture: LBP/HOG/Lacunarity on green band (downsample for speed) ---
            pdata['texture_features'] = {}
            spectral = pdata.get('spectral_stack', {})
            if 'green' in spectral:
                green_band = np.asarray(spectral['green'], dtype=np.float64)
                if green_band.ndim == 3 and green_band.shape[-1] == 1:
                    green_band = green_band[..., 0]

                # Downsample to 256x256 max for faster texture computation
                h, w = green_band.shape[:2]
                if h > 256 or w > 256:
                    scale = 256 / max(h, w)
                    green_band = cv2.resize(green_band, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
                    if mask is not None:
                        mask_resized = cv2.resize(mask, (green_band.shape[1], green_band.shape[0]), interpolation=cv2.INTER_NEAREST)
                    else:
                        mask_resized = None
                else:
                    mask_resized = mask

                if mask_resized is not None:
                    valid = np.where(mask_resized > 0, green_band, np.nan)
                else:
                    valid = green_band

                v = np.nan_to_num(valid, nan=np.nanmin(valid))
                m, M = np.min(v), np.max(v)
                denom = (M - m) if (M - m) > 1e-6 else 1.0
                gray8 = ((v - m) / denom * 255.0).astype(np.uint8)

                lbp_map = self.texture_extractor.extract_lbp(gray8)
                hog_map = self.texture_extractor.extract_hog(gray8)
                lac1_map = self.texture_extractor.compute_local_lacunarity(gray8)
                pdata['texture_features'] = {'green': {'features': {'lbp': lbp_map, 'hog': hog_map, 'lac1': lac1_map}}}

            # --- Vegetation indices ---
            if spectral and mask is not None:
                pdata['vegetation_indices'] = self._compute_vegetation(spectral, mask)
            else:
                pdata['vegetation_indices'] = {}

            # --- Morphology ---
            try:
                if mask is not None and isinstance(composite, np.ndarray):
                    morph = self.morphology_extractor.extract_morphology_features(composite, mask)
                    pdata['morphology_features'] = morph
                else:
                    pdata['morphology_features'] = {}
            except Exception:
                pdata['morphology_features'] = {}
        
        return plants
    
    def _compute_vegetation(self, spectral: Dict[str, np.ndarray], mask: np.ndarray) -> Dict[str, Any]:
        """Compute NDVI, GNDVI, SAVI."""
        out = {}
        for name in ("NDVI", "GNDVI", "SAVI"):
            bands = self.vegetation_extractor.index_bands.get(name, [])
            if not all(b in spectral for b in bands):
                continue

            arrays = []
            for b in bands:
                arr = np.asarray(spectral[b], dtype=np.float64)
                if arr.ndim == 3 and arr.shape[-1] == 1:
                    arr = arr[..., 0]
                arrays.append(arr)

            values = self.vegetation_extractor.index_formulas[name](*arrays).astype(np.float64)
            binary_mask = (mask > 0)
            masked_values = np.where(binary_mask, values, np.nan)
            valid = masked_values[~np.isnan(masked_values)]

            stats = {
                'mean': float(np.mean(valid)) if valid.size else 0.0,
                'std': float(np.std(valid)) if valid.size else 0.0,
            }
            out[name] = {'values': masked_values, 'statistics': stats}
        return out