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)
                # Convert BGR→RGB for correct viewing in standard image viewers
                overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
                cv2.imwrite(str(results_dir / 'overlay.png'), overlay_rgb)
        except Exception as e:
            logger.error(f"Failed to save overlay: {e}")

        # 2b. Composite (input to segmentation)
        try:
            base_image = plant_data.get('composite')
            if isinstance(base_image, np.ndarray):
                # Ensure uint8
                if base_image.dtype != np.uint8:
                    base_image = self._normalize_to_uint8(base_image.astype(np.float64))
                # Convert BGR→RGB for human viewing
                comp_rgb = cv2.cvtColor(base_image, cv2.COLOR_BGR2RGB)
                cv2.imwrite(str(results_dir / 'composite.png'), comp_rgb)
        except Exception as e:
            logger.error(f"Failed to save composite: {e}")

        # 3-5. Vegetation indices (NDVI, GNDVI, SAVI)
        try:
            veg = plant_data.get('vegetation_indices', {})
            for name in ['NDVI', 'GNDVI', 'SAVI']:
                data = veg.get(name, {})
                values = data.get('values') if isinstance(data, dict) else None
                if isinstance(values, np.ndarray) and values.size > 0:
                    try:
                        # Colormap with colorbar similar to src: use matplotlib savefig
                        cmap = cm.RdYlGn
                        # Value ranges
                        if name in ['NDVI', 'GNDVI']:
                            vmin, vmax = (-1, 1)
                        else:
                            vmin, vmax = (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')
                        im = ax.imshow(masked, cmap=cmap, vmin=vmin, vmax=vmax)
                        # add colorbar
                        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
                        cbar.ax.tick_params(labelsize=8)
                        plt.tight_layout()
                        plt.savefig(veg_dir / f"{name.lower()}.png", dpi=120, 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. Texture features: ONLY LBP on green band
        try:
            tex = plant_data.get('texture_features', {})
            green_band = tex.get('green', {})
            feats = green_band.get('features', {})

            lbp = feats.get('lbp')
            if isinstance(lbp, np.ndarray) and lbp.size > 0:
                try:
                    img = lbp.astype(np.float64)
                    fig, ax = plt.subplots(figsize=(5, 5))
                    ax.set_axis_off()
                    ax.set_facecolor('white')
                    im = ax.imshow(img, cmap='gray', vmin=0, vmax=255)
                    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
                    cbar.ax.tick_params(labelsize=8)
                    plt.tight_layout()
                    plt.savefig(tex_dir / 'lbp_green.png', dpi=120, bbox_inches='tight')
                    plt.close(fig)
                except Exception as e:
                    logger.error(f"Failed to save LBP with colorbar: {e}")
        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 green overlay on brightened composite, following src pipeline style."""
        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
        base = image
        if base.dtype != np.uint8:
            base = self._normalize_to_uint8(base.astype(np.float64))
        bright = cv2.convertScaleAbs(base, alpha=1.2, beta=15)
        green_overlay = bright.copy()
        green_overlay[binary == 255] = (0, 255, 0)
        blended = cv2.addWeighted(bright, 1.0, green_overlay, 0.5, 0)
        return blended
    
    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)