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, single_plant_mode: bool = False):
        """Initialize pipeline."""
        logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
        self.config = config
        self.config.validate()
        self.single_plant_mode = single_plant_mode
        
        # Initialize components with defaults
        self.preprocessor = ImagePreprocessor()
        self.mask_handler = MaskHandler()
        self.texture_extractor = TextureExtractor()
        self.vegetation_extractor = VegetationIndexExtractor()
        self.morphology_extractor = MorphologyExtractor(single_plant_mode=single_plant_mode)
        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(f"Pipeline initialized (single_plant_mode={single_plant_mode})")

    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, single_plant_mode: bool = False) -> 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 0: Save and show input image immediately
        logger.info("Stage 0: Saving input image...")
        for key, pdata in plants.items():
            # Quick save of just the input image
            self.output_manager._save_input_image_only(key, pdata)
        yield {"plants": plants, "stage": "input"}

        # Stage 1: Create composite + Segmentation
        logger.info("Stage 1: Creating composite and segmenting...")
        plants = self.preprocessor.create_composites(plants)
        plants = self._segment(plants)
        if progress_callback:
            progress_callback("segmentation", plants)
        # Save composite, mask, overlay
        for key, pdata in plants.items():
            self.output_manager.save_plant_results(key, pdata)
        yield {"plants": plants, "stage": "segmentation"}

        # Stage 2: Extract features (texture, vegetation, morphology) - run in parallel where possible
        logger.info("Stage 2: Extracting features...")
        plants = self._extract_features_fast(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 128x128 max for faster texture computation (2x speedup)
                h, w = green_band.shape[:2]
                if h > 128 or w > 128:
                    scale = 128 / 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 _extract_features_fast(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Fast feature extraction - skip textures, minimal vegetation indices."""
        for key, pdata in plants.items():
            composite = pdata['composite']
            mask = pdata.get('mask')

            # Skip texture extraction for speed (can be added back if needed)
            pdata['texture_features'] = {}
            spectral = pdata.get('spectral_stack', {})
            
            # Compute texture on green band with high quality (512x512)
            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 512x512 for high texture quality
                h, w = green_band.shape[:2]
                if h > 512 or w > 512:
                    scale = 512 / 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