sorghum_pipeline/pipeline.py

"""
Main pipeline class for the Sorghum Plant Phenotyping Pipeline.

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

import os
import logging
from pathlib import Path
from typing import Dict, Any, Optional
import numpy as np
import cv2
from sklearn.decomposition import PCA

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 plant phenotyping."""
    
    def __init__(self, config: Config):
        """Initialize the minimal pipeline."""
        self._setup_logging()
        self.config = config
        self.config.validate()
        self._initialize_components()
        logger.info("Sorghum Pipeline initialized")
    
    def _setup_logging(self):
        """Setup logging configuration."""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[logging.StreamHandler()]
        )
    
    def _initialize_components(self):
        """Initialize pipeline components."""
        self.preprocessor = ImagePreprocessor(target_size=None)
        self.mask_handler = MaskHandler(min_area=1000, kernel_size=7)
        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(),
            threshold=0.5,
            trust_remote_code=True
        )
        self.output_manager = OutputManager(
            output_folder=self.config.paths.output_folder,
            settings=self.config.output
        )

    def run(self, single_image_path: str) -> Dict[str, Any]:
        """
        Run minimal pipeline on single image.
        
        Args:
            single_image_path: Path to input image
            
        Returns:
            Dictionary containing results
        """
        logger.info("Starting minimal single-image pipeline...")
        
        try:
            import time
            from PIL import Image as _Image
            
            total_start = time.perf_counter()
            
            # Load single image
            _p = Path(single_image_path)
            _img = _Image.open(str(_p))
            plants = {
                "demo_demo_frame1": {
                    "raw_image": (_img, _p.name),
                    "plant_name": "demo",
                    "file_path": str(_p)
                }
            }
            
            # Create composite
            plants = self.preprocessor.create_composites(plants)
            
            # Segment
            plants = self._segment_plants(plants)
            
            # Extract features
            plants = self._extract_features(plants)
            
            # Generate outputs
            self._generate_outputs(plants)
            
            # Summary
            summary = self._create_summary(plants)
            
            total_time = time.perf_counter() - total_start
            logger.info(f"Pipeline completed in {total_time:.2f}s")
            
            return {
                "plants": plants,
                "summary": summary,
                "config": self.config,
                "timing_seconds": total_time
            }
            
        except Exception as e:
            logger.error(f"Pipeline failed: {e}")
            raise

    def _segment_plants(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Segment plants using BRIA model (full image)."""
        for key, pdata in plants.items():
            try:
                composite = pdata['composite']
                soft_mask = self.segmentation_manager.segment_image_soft(composite)
                pdata['soft_mask'] = soft_mask
                pdata['mask'] = (soft_mask * 255.0).astype(np.uint8)
                logger.info(f"Segmented {key}")
            except Exception as e:
                logger.error(f"Segmentation failed for {key}: {e}")
                pdata['soft_mask'] = np.zeros(composite.shape[:2], dtype=np.float32)
                pdata['mask'] = np.zeros(composite.shape[:2], dtype=np.uint8)
        return plants
    
    def _extract_features(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Extract features from plants."""
        for key, pdata in plants.items():
            try:
                pdata['texture_features'] = self._extract_texture_features(pdata)
                pdata['vegetation_indices'] = self._extract_vegetation_indices(pdata)
                pdata['morphology_features'] = self._extract_morphology_features(pdata)
                logger.info(f"Features extracted for {key}")
            except Exception as e:
                logger.error(f"Feature extraction failed for {key}: {e}")
                pdata['texture_features'] = {}
                pdata['vegetation_indices'] = {}
                pdata['morphology_features'] = {}
        return plants
    
    def _extract_texture_features(self, pdata: Dict[str, Any]) -> Dict[str, Any]:
        """Extract texture features from pseudo-color image only."""
        features = {}
        try:
            # Only process pseudo-color composite
            composite = pdata['composite']
            mask = pdata.get('mask')
            if mask is not None:
                masked = self.mask_handler.apply_mask_to_image(composite, mask)
                gray_image = cv2.cvtColor(masked, cv2.COLOR_BGR2GRAY)
            else:
                gray_image = cv2.cvtColor(composite, cv2.COLOR_BGR2GRAY)
            
            band_features = self.texture_extractor.extract_all_texture_features(gray_image)
            stats = self.texture_extractor.compute_texture_statistics(band_features, mask)
            
            features['color'] = {
                'features': band_features,
                'statistics': stats
            }
        except Exception as e:
            logger.error(f"Texture extraction failed: {e}")
            features['color'] = {'features': {}, 'statistics': {}}
        
        return features
    
    def _extract_vegetation_indices(self, pdata: Dict[str, Any]) -> Dict[str, Any]:
        """Extract vegetation indices (NDVI, ARI, GNDVI only)."""
        try:
            spectral_stack = pdata.get('spectral_stack', {})
            mask = pdata.get('mask')
            if not spectral_stack or mask is None:
                return {}
            
            out: Dict[str, Any] = {}
            for name in ("NDVI", "ARI", "GNDVI"):
                bands = self.vegetation_extractor.index_bands.get(name, [])
                if not all(b in spectral_stack for b in bands):
                    continue
                arrays = []
                for b in bands:
                    arr = spectral_stack[b]
                    if isinstance(arr, np.ndarray):
                        arr = arr.squeeze(-1)
                    arrays.append(np.asarray(arr, dtype=np.float64))
                
                values = self.vegetation_extractor.index_formulas[name](*arrays).astype(np.float64)
                binary_mask = (np.asarray(mask).astype(np.int32) > 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,
                    'min': float(np.min(valid)) if valid.size else 0.0,
                    'max': float(np.max(valid)) if valid.size else 0.0,
                    'median': float(np.median(valid)) if valid.size else 0.0,
                }
                out[name] = {'values': masked_values, 'statistics': stats}
            return out
        except Exception as e:
            logger.error(f"Vegetation index extraction failed: {e}")
            return {}
    
    def _extract_morphology_features(self, pdata: Dict[str, Any]) -> Dict[str, Any]:
        """Extract morphological features."""
        try:
            composite = pdata.get('composite')
            mask = pdata.get('mask')
            if composite is None or mask is None:
                return {}
            return self.morphology_extractor.extract_morphology_features(composite, mask)
        except Exception as e:
            logger.error(f"Morphology extraction failed: {e}")
            return {}
    
    def _generate_outputs(self, plants: Dict[str, Any]) -> None:
        """Generate output files."""
        self.output_manager.create_output_directories()
        for key, pdata in plants.items():
            try:
                self.output_manager.save_plant_results(key, pdata)
            except Exception as e:
                logger.error(f"Output generation failed for {key}: {e}")
    
    def _create_summary(self, plants: Dict[str, Any]) -> Dict[str, Any]:
        """Create summary of results."""
        return {
            "total_plants": len(plants),
            "successful_plants": sum(1 for p in plants.values() if p.get('texture_features')),
            "features_extracted": {
                "texture": sum(1 for p in plants.values() if p.get('texture_features')),
                "vegetation": sum(1 for p in plants.values() if p.get('vegetation_indices')),
                "morphology": sum(1 for p in plants.values() if p.get('morphology_features'))
            }
        }