sorghum_pipeline/output/manager.py

"""
Minimal output manager for demo (saves only 7 required images).
"""

import os
import numpy as np
import cv2
import matplotlib
if os.environ.get('MPLBACKEND') is None:
    matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from pathlib import Path
from typing import Dict, Any
import logging

logger = logging.getLogger(__name__)


class OutputManager:
    """Minimal output manager for demo."""
    
    def __init__(self, output_folder: str, settings: Any):
        """Initialize output manager."""
        self.output_folder = Path(output_folder)
        self.settings = settings
        try:
            self.minimal_demo: bool = bool(int(os.environ.get('MINIMAL_DEMO', '0')))
        except Exception:
            self.minimal_demo = False
        self.output_folder.mkdir(parents=True, exist_ok=True)

    def create_output_directories(self) -> None:
        """Create output directories."""
        self.output_folder.mkdir(parents=True, exist_ok=True)
    
    def save_plant_results(self, plant_key: str, plant_data: Dict[str, Any]) -> None:
        """Save minimal demo outputs only."""
        if not self.minimal_demo:
            logger.warning("OutputManager configured for minimal demo only")
            return
        
        self._save_minimal_demo_outputs(plant_data)

    def _save_minimal_demo_outputs(self, plant_data: Dict[str, Any]) -> None:
        """Save only the 7 required images."""
        results_dir = self.output_folder / 'results'
        veg_dir = self.output_folder / 'Vegetation_indices_images'
        tex_dir = self.output_folder / 'texture_output'
        results_dir.mkdir(parents=True, exist_ok=True)
        veg_dir.mkdir(parents=True, exist_ok=True)
        tex_dir.mkdir(parents=True, exist_ok=True)

        # 1. Mask
        try:
            mask = plant_data.get('mask')
            if isinstance(mask, np.ndarray):
                cv2.imwrite(str(results_dir / 'mask.png'), mask)
        except Exception as e:
            logger.error(f"Failed to save mask: {e}")

        # 2. Overlay
        try:
            base_image = plant_data.get('composite')
            mask = plant_data.get('mask')
            if isinstance(base_image, np.ndarray) and isinstance(mask, np.ndarray):
                overlay = self._create_overlay(base_image, mask)
                cv2.imwrite(str(results_dir / 'overlay.png'), overlay)
        except Exception as e:
            logger.error(f"Failed to save overlay: {e}")

        # 3-5. Vegetation indices (NDVI, ARI, GNDVI)
        try:
            veg = plant_data.get('vegetation_indices', {})
            for name in ['NDVI', 'ARI', 'GNDVI']:
                data = veg.get(name, {})
                values = data.get('values') if isinstance(data, dict) else None
                if isinstance(values, np.ndarray) and values.size > 0:
                    try:
                        cmap = cm.RdYlGn if name in ['NDVI', 'GNDVI'] else cm.magma
                        vmin, vmax = (-1, 1) if name in ['NDVI', 'GNDVI'] else (0, 1)
                        
                        masked = np.ma.masked_invalid(values.astype(np.float64))
                        fig, ax = plt.subplots(figsize=(5, 5))
                        ax.set_axis_off()
                        ax.set_facecolor('white')
                        ax.imshow(masked, cmap=cmap, vmin=vmin, vmax=vmax)
                        plt.tight_layout()
                        plt.savefig(veg_dir / f"{name.lower()}.png", dpi=100, bbox_inches='tight')
                        plt.close(fig)
                    except Exception as e:
                        logger.error(f"Failed to save {name}: {e}")
        except Exception as e:
            logger.error(f"Failed to save vegetation indices: {e}")

        # 6-8. Texture features (LBP, HOG, Lacunarity)
        try:
            tex = plant_data.get('texture_features', {})
            color_band = tex.get('color', {})
            feats = color_band.get('features', {})
            
            if isinstance(feats.get('lbp'), np.ndarray) and feats['lbp'].size > 0:
                cv2.imwrite(str(tex_dir / 'lbp.png'), feats['lbp'].astype(np.uint8))
            
            if isinstance(feats.get('hog'), np.ndarray) and feats['hog'].size > 0:
                cv2.imwrite(str(tex_dir / 'hog.png'), feats['hog'].astype(np.uint8))
            
            lac = feats.get('lac2')
            if isinstance(lac, np.ndarray) and lac.size > 0:
                if lac.dtype != np.uint8:
                    lac = self._normalize_to_uint8(lac.astype(np.float64))
                cv2.imwrite(str(tex_dir / 'lacunarity.png'), lac)
        except Exception as e:
            logger.error(f"Failed to save texture: {e}")

        # 9. Morphology size analysis
        try:
            morph = plant_data.get('morphology_features', {})
            images = morph.get('images', {})
            size_img = images.get('size_analysis')
            if isinstance(size_img, np.ndarray) and size_img.size > 0:
                cv2.imwrite(str(results_dir / 'size.size_analysis.png'), size_img)
        except Exception as e:
            logger.error(f"Failed to save size analysis: {e}")
    
    def _create_overlay(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Create overlay (masked pixels only)."""
        if mask is None:
            return image
        if mask.shape[:2] != image.shape[:2]:
            mask = cv2.resize(mask.astype(np.uint8), (image.shape[1], image.shape[0]), 
                            interpolation=cv2.INTER_NEAREST)
        binary = (mask.astype(np.int32) > 0).astype(np.uint8) * 255
        return cv2.bitwise_and(image, image, mask=binary)
    
    def _normalize_to_uint8(self, arr: np.ndarray) -> np.ndarray:
        """Normalize to uint8."""
        arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
        ptp = np.ptp(arr)
        if ptp > 0:
            normalized = (arr - arr.min()) / (ptp + 1e-6) * 255
        else:
            normalized = np.zeros_like(arr)
        return np.clip(normalized, 0, 255).astype(np.uint8)