sorghum_pipeline/pipeline.py

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

import logging
from pathlib import Path
from typing import Dict, Any
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",
            }
        }

        # 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 _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."""
        for key, pdata in plants.items():
            composite = pdata['composite']
            mask = pdata.get('mask')

            # --- Texture: LBP on green band ---
            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]

                if mask is not None:
                    valid = np.where(mask > 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)
                pdata['texture_features'] = {'green': {'features': {'lbp': lbp_map}}}

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

            # --- Morphology (currently empty) ---
            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