{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Updated Documentation with New Preprocessing\n\n## Data Used\n* **Images**: Top-view images of 70cm×30cm quadrats (normalized to 2000×1000px)\n* **Targets**: 5 types of biomass (Dry_Green_g, Dry_Dead_g, Dry_Clover_g, GDM_g, Dry_Total_g)\n* **Auxiliary Information**: NDVI, height, species, date, state (training data only)\n\n## Preprocessing\n\n### Phase 1: Core Preprocessing\n* **White Balance Correction**: Normalize color differences between cameras using Gray World algorithm\n* **CLAHE Enhancement**: Improve local contrast and reduce shadow effects via Contrast Limited Adaptive Histogram Equalization\n* **Vegetation Indices**: Calculate RGB-based indices for vegetation detection\n  - ExG (Excess Green Index): Emphasizes green vegetation\n  - VARI (Visible Atmospherically Resistant Index): Resistant to lighting variations\n  - NDI (Normalized Difference Index): Detects dead vegetation\n* **HSV Masking**: Enhance green vegetation (probabilistic application during training)\n* **RGB Normalization**: ImageNet statistics-based normalization\n* **Data Augmentation**: Rotation, horizontal/vertical flip, color jitter\n\n### Phase 2: Advanced Preprocessing\n* **Shadow Removal**: Correct lighting variations in outdoor images using HSV-based correction\n* **LAB Color Space**: Transform to perceptually uniform color space for better species discrimination\n* **LBP Texture Features**: Extract Local Binary Pattern features to distinguish living vs dead vegetation\n\n### Target Processing\n* **Log Transformation**: Stabilize target variables using log(1 + x) transformation\n\n## Model Definition\n* **Backbone**: EfficientNetV2-S (pre-trained) with flexible multi-channel input support\n* **Architecture**: Multi-task output layer with 5 regression heads (one for each biomass component)\n* **Input Channels**: Variable (3-10 channels) depending on enabled preprocessing\n  - Base: RGB (3 channels)\n  - +Vegetation Indices: +3 channels (ExG, VARI, NDI)\n  - +LAB: +3 channels (L, A, B)\n  - +LBP: +1 channel (texture)\n\n## Training Configuration\n* **Loss Function**: Weighted MSE (weights aligned with evaluation metrics: Dry_Total_g=0.5, GDM_g=0.2, others=0.1)\n* **Optimizer**: AdamW with Cosine Annealing learning rate scheduler\n* **Cross-Validation**: 5-fold stratified CV evaluation\n* **Training Monitoring**: Real-time progress tracking with tqdm\n* **Model Checkpointing**: Automatic saving of best models per fold\n\n## Others\n* **Multi-Channel Processing**: Support for 3-10 input channels with dynamic configuration\n* **Image Feature Extraction**: Color histograms, texture features (LBP), color space transformations\n* **Model Persistence**: Save/load functionality with preprocessing metadata\n* **Ablation Study Support**: Toggle individual preprocessing steps via configuration flags\n* **Pretrained Weight Transfer**: Intelligent weight initialization for multi-channel inputs\n* **TTA (Test Time Augmentation)**: Available for inference (not yet implemented in this version)\n\n## Implementation Highlights\n* **Configurable Pipeline**: All preprocessing steps can be toggled on/off\n* **Preprocessing Summary**: Automatic tracking of enabled features in saved models\n* **Per-Target Metrics**: Individual R² scores reported for each biomass component\n* **Memory Efficient**: Proper cleanup between folds with garbage collection","metadata":{}},{"cell_type":"code","source":"\"\"\"\nPhase 2 Preprocessing Implementation\n- LAB Color Space Transformation\n- LBP Texture Features\n- Shadow Removal\n- Model Saving Functionality\nAll preprocessing can be toggled on/off via Config\n\"\"\"\n\nimport os\nimport gc\nfrom pathlib import Path\nfrom typing import Optional\nfrom dataclasses import dataclass\nfrom datetime import datetime\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import KFold\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom torchvision import models\nimport timm\nfrom tqdm import tqdm\n\n# ==================== Configuration Class (Updated) ====================\n@dataclass\nclass Config:\n    \"\"\"Centralized management of training configuration\"\"\"\n    \n    # Path settings\n    input_dir: str = '/kaggle/input/csiro-biomass/'\n    output_dir: str = '/kaggle/working/'\n    \n    # Image settings\n    image_size: tuple[int, int] = (512, 512)\n    \n    # Training settings\n    n_folds: int = 5\n    batch_size: int = 16\n    epochs: int = 5\n    learning_rate: float = 1e-4\n    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n    \n    # Target settings\n    target_cols: list[str] = None\n    target_weights: dict[str, float] = None\n    \n    # Model saving settings\n    save_models: bool = True                # Whether to save trained models\n    save_best_only: bool = True             # Save only best model per fold\n    model_save_dir: str = 'models'          # Subdirectory for model files\n    \n    # ========== Phase 1 Preprocessing Flags ==========\n    use_white_balance: bool = True          # White balance correction\n    use_clahe: bool = True                  # CLAHE contrast enhancement\n    use_vegetation_indices: bool = True     # Vegetation indices (ExG, VARI, NDI)\n    use_hsv_masking: bool = False           # HSV masking (original feature)\n    \n    # ========== Phase 2 Preprocessing Flags ==========\n    use_lab_color_space: bool = False       # LAB color space transformation\n    use_lbp_texture: bool = False           # LBP texture features\n    use_shadow_removal: bool = False        # Shadow removal\n    \n    # Preprocessing channel counts\n    n_vegetation_channels: int = 3          # ExG, VARI, NDI\n    n_lab_channels: int = 3                 # L, A, B\n    n_lbp_channels: int = 1                 # LBP\n    \n    def __post_init__(self):\n        \"\"\"Post-initialization processing\"\"\"\n        if self.target_cols is None:\n            self.target_cols = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n        \n        if self.target_weights is None:\n            self.target_weights = {\n                'Dry_Green_g': 0.1,\n                'Dry_Dead_g': 0.1,\n                'Dry_Clover_g': 0.1,\n                'GDM_g': 0.2,\n                'Dry_Total_g': 0.5\n            }\n        \n        # Create model save directory\n        if self.save_models:\n            self.full_model_save_path = os.path.join(self.output_dir, self.model_save_dir)\n            os.makedirs(self.full_model_save_path, exist_ok=True)\n    \n    @property\n    def total_input_channels(self) -> int:\n        \"\"\"\n        Calculate total number of input channels based on enabled preprocessing\n        \n        Returns:\n            Total number of channels\n        \"\"\"\n        n_channels = 3  # Base RGB channels\n        \n        if self.use_vegetation_indices:\n            n_channels += self.n_vegetation_channels\n        \n        if self.use_lab_color_space:\n            n_channels += self.n_lab_channels\n        \n        if self.use_lbp_texture:\n            n_channels += self.n_lbp_channels\n        \n        return n_channels\n    \n    def get_preprocessing_summary(self) -> str:\n        \"\"\"\n        Get summary of enabled preprocessing\n        \n        Returns:\n            Summary string\n        \"\"\"\n        enabled = []\n        \n        # Phase 1\n        if self.use_white_balance:\n            enabled.append(\"WhiteBalance\")\n        if self.use_clahe:\n            enabled.append(\"CLAHE\")\n        if self.use_vegetation_indices:\n            enabled.append(\"VegetationIndices\")\n        if self.use_hsv_masking:\n            enabled.append(\"HSVMask\")\n        \n        # Phase 2\n        if self.use_shadow_removal:\n            enabled.append(\"ShadowRemoval\")\n        if self.use_lab_color_space:\n            enabled.append(\"LAB\")\n        if self.use_lbp_texture:\n            enabled.append(\"LBP\")\n        \n        return \", \".join(enabled) if enabled else \"None\"\n\n\n# ==================== Phase 1 Preprocessing Classes ====================\nclass VegetationIndices:\n    \"\"\"Calculate RGB-based vegetation indices\"\"\"\n    \n    @staticmethod\n    def calculate_exg(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        ExG (Excess Green Index)\n        Emphasizes green vegetation, effective for separating from soil\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            ExG index map (H, W)\n        \"\"\"\n        image = image.astype(np.float32)\n        r, g, b = image[:, :, 0], image[:, :, 1], image[:, :, 2]\n        \n        exg = 2 * g - r - b\n        exg = np.clip(exg, 0, 255).astype(np.uint8)\n        \n        return exg\n    \n    @staticmethod\n    def calculate_vari(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        VARI (Visible Atmospherically Resistant Index)\n        Resistant to lighting variations\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            VARI index map (H, W)\n        \"\"\"\n        r, g, b = image[:, :, 0], image[:, :, 1], image[:, :, 2]\n        \n        denominator = g.astype(np.float32) + r.astype(np.float32) - b.astype(np.float32)\n        denominator = np.where(denominator == 0, 1, denominator)\n        vari = (g.astype(np.float32) - r.astype(np.float32)) / denominator\n        \n        vari = ((vari + 1) * 127.5).astype(np.uint8)\n        \n        return vari\n    \n    @staticmethod\n    def calculate_ndi(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        NDI (Normalized Difference Index)\n        Detects dead vegetation\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            NDI index map (H, W)\n        \"\"\"\n        r, g, b = image[:, :, 0], image[:, :, 1], image[:, :, 2]\n        \n        numerator = r.astype(np.float32) - g.astype(np.float32)\n        denominator = r.astype(np.float32) + g.astype(np.float32)\n        denominator = np.where(denominator == 0, 1, denominator)\n        ndi = numerator / denominator\n        \n        ndi = ((ndi + 1) * 127.5).astype(np.uint8)\n        \n        return ndi\n    \n    @staticmethod\n    def calculate_all_indices(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Calculate all vegetation indices at once\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            Stacked indices (H, W, 3) - [ExG, VARI, NDI]\n        \"\"\"\n        exg = VegetationIndices.calculate_exg(image)\n        vari = VegetationIndices.calculate_vari(image)\n        ndi = VegetationIndices.calculate_ndi(image)\n        \n        indices = np.stack([exg, vari, ndi], axis=-1)\n        \n        return indices\n\n\n# ==================== Phase 2 Preprocessing Classes (NEW) ====================\nclass ColorSpaceTransforms:\n    \"\"\"\n    Color space transformations for better plant species discrimination\n    \"\"\"\n    \n    @staticmethod\n    def apply_lab_space(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        LAB color space: Perceptually uniform color space\n        \n        Purpose: Detect subtle color differences between clover vs other green plants\n        \n        L channel: Lightness (0-255)\n        A channel: Green-Red axis (-128 to 127, mapped to 0-255)\n        B channel: Blue-Yellow axis (-128 to 127, mapped to 0-255)\n        \n        Why effective:\n        - Clover may have different A-B distribution\n        - Perceptually uniform color space\n        - Better separation between similar greens\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            LAB image (H, W, 3)\n        \"\"\"\n        lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)\n        \n        return lab\n\n\nclass TextureFeatures:\n    \"\"\"\n    Texture features to distinguish living vs dead vegetation\n    \"\"\"\n    \n    @staticmethod\n    def calculate_lbp_features(image: np.ndarray, radius: int = 3, n_points: int = 24) -> np.ndarray:\n        \"\"\"\n        LBP (Local Binary Patterns) for local texture\n        \n        Purpose: Lightweight and effective texture descriptor\n        \n        Why effective for this task:\n        - Living vegetation has regular texture patterns\n        - Dead vegetation has irregular, chaotic patterns\n        - Effective for Dry_Dead_g discrimination\n        \n        Args:\n            image: RGB image (H, W, 3) or grayscale (H, W)\n            radius: Radius of circular LBP\n            n_points: Number of sampling points\n        \n        Returns:\n            LBP texture map (H, W)\n        \"\"\"\n        # Convert to grayscale if needed\n        if len(image.shape) == 3:\n            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        else:\n            gray = image\n        \n        # Calculate LBP using uniform pattern\n        # Note: scikit-image's local_binary_pattern requires scipy\n        # For Kaggle compatibility, use simpler implementation\n        \n        h, w = gray.shape\n        lbp = np.zeros((h, w), dtype=np.uint8)\n        \n        # Calculate offsets for circular pattern\n        for i in range(n_points):\n            angle = 2 * np.pi * i / n_points\n            dy = -radius * np.sin(angle)\n            dx = radius * np.cos(angle)\n            \n            # Bilinear interpolation for sub-pixel coordinates\n            fy = np.clip(np.arange(h)[:, None] + dy, 0, h - 1.001)\n            fx = np.clip(np.arange(w)[None, :] + dx, 0, w - 1.001)\n            \n            y0, x0 = fy.astype(int), fx.astype(int)\n            y1, x1 = np.clip(y0 + 1, 0, h - 1), np.clip(x0 + 1, 0, w - 1)\n            \n            # Interpolation weights\n            wy = fy - y0\n            wx = fx - x0\n            \n            # Bilinear interpolation\n            neighbor = (gray[y0, x0] * (1 - wy) * (1 - wx) +\n                       gray[y1, x0] * wy * (1 - wx) +\n                       gray[y0, x1] * (1 - wy) * wx +\n                       gray[y1, x1] * wy * wx)\n            \n            # Compare with center pixel\n            lbp += ((neighbor >= gray) * (2 ** i)).astype(np.uint8)\n        \n        return lbp\n    \n    @staticmethod\n    def calculate_lbp_simple(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Simplified LBP implementation for faster computation\n        Uses 8-point circular pattern with radius=1\n        \n        Args:\n            image: RGB image (H, W, 3) or grayscale (H, W)\n        \n        Returns:\n            LBP texture map (H, W)\n        \"\"\"\n        if len(image.shape) == 3:\n            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        else:\n            gray = image\n        \n        h, w = gray.shape\n        lbp = np.zeros((h - 2, w - 2), dtype=np.uint8)\n        \n        # 8 neighbors in clockwise order starting from top-left\n        offsets = [(-1, -1), (-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1)]\n        \n        for i in range(1, h - 1):\n            for j in range(1, w - 1):\n                center = gray[i, j]\n                pattern = 0\n                \n                for idx, (dy, dx) in enumerate(offsets):\n                    if gray[i + dy, j + dx] >= center:\n                        pattern |= (1 << idx)\n                \n                lbp[i - 1, j - 1] = pattern\n        \n        # Pad to original size\n        lbp_padded = np.pad(lbp, ((1, 1), (1, 1)), mode='edge')\n        \n        return lbp_padded\n\n\nclass ShadowRemoval:\n    \"\"\"\n    Shadow removal techniques\n    \"\"\"\n    \n    @staticmethod\n    def remove_shadow(image: np.ndarray, method: str = 'hsv') -> np.ndarray:\n        \"\"\"\n        Remove shadows from outdoor images\n        \n        Purpose: Outdoor images have significant shadow issues that distort biomass estimation\n        \n        Why effective:\n        - Shadows reduce apparent green color\n        - Affects vegetation index calculations\n        - Creates false low biomass regions\n        \n        Args:\n            image: RGB image (H, W, 3)\n            method: 'hsv' or 'lab' or 'morphological'\n        \n        Returns:\n            Shadow-corrected image (H, W, 3)\n        \"\"\"\n        if method == 'hsv':\n            return ShadowRemoval._remove_shadow_hsv(image)\n        elif method == 'lab':\n            return ShadowRemoval._remove_shadow_lab(image)\n        elif method == 'morphological':\n            return ShadowRemoval._remove_shadow_morphological(image)\n        else:\n            return image\n    \n    @staticmethod\n    def _remove_shadow_hsv(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Shadow removal using HSV color space\n        Corrects V (value) channel in shadow regions\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            Shadow-corrected image (H, W, 3)\n        \"\"\"\n        hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)\n        h, s, v = cv2.split(hsv)\n        \n        # Detect shadow regions (low V value)\n        v_mean = np.mean(v)\n        v_std = np.std(v)\n        shadow_threshold = max(0, v_mean - 0.5 * v_std)\n        \n        shadow_mask = v < shadow_threshold\n        \n        # Correct V channel in shadow regions\n        v_corrected = v.copy().astype(np.float32)\n        \n        if shadow_mask.any():\n            # Calculate correction factor\n            non_shadow_mean = np.mean(v[~shadow_mask]) if (~shadow_mask).any() else v_mean\n            shadow_mean = np.mean(v[shadow_mask])\n            \n            if shadow_mean > 0:\n                correction_factor = non_shadow_mean / shadow_mean\n                correction_factor = min(correction_factor, 2.0)  # Limit correction\n                \n                v_corrected[shadow_mask] = np.clip(\n                    v[shadow_mask] * correction_factor,\n                    0,\n                    255\n                )\n        \n        # Merge and convert back\n        hsv_corrected = cv2.merge([h, s, v_corrected.astype(np.uint8)])\n        rgb_corrected = cv2.cvtColor(hsv_corrected, cv2.COLOR_HSV2RGB)\n        \n        return rgb_corrected\n    \n    @staticmethod\n    def _remove_shadow_lab(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Shadow removal using LAB color space\n        Corrects L (lightness) channel\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            Shadow-corrected image (H, W, 3)\n        \"\"\"\n        lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)\n        l, a, b = cv2.split(lab)\n        \n        # Apply CLAHE to L channel to reduce shadow effects\n        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))\n        l_corrected = clahe.apply(l)\n        \n        # Merge and convert back\n        lab_corrected = cv2.merge([l_corrected, a, b])\n        rgb_corrected = cv2.cvtColor(lab_corrected, cv2.COLOR_LAB2RGB)\n        \n        return rgb_corrected\n    \n    @staticmethod\n    def _remove_shadow_morphological(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Shadow removal using morphological operations\n        Removes shadow through background estimation\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            Shadow-corrected image (H, W, 3)\n        \"\"\"\n        # Convert to grayscale for processing\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        \n        # Morphological dilation to estimate background\n        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))\n        background = cv2.dilate(gray, kernel)\n        \n        # Calculate difference\n        diff = cv2.subtract(background, gray)\n        \n        # Normalize\n        diff_normalized = cv2.normalize(diff, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n        \n        # Add to original image\n        result = cv2.add(image, cv2.cvtColor(diff_normalized, cv2.COLOR_GRAY2RGB))\n        \n        return result\n\n\nclass AdaptivePreprocessor:\n    \"\"\"Adaptive preprocessing adjusted to image characteristics\"\"\"\n    \n    @staticmethod\n    def white_balance_correction(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Automatic white balance correction using Gray World algorithm\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            White balance corrected image (H, W, 3)\n        \"\"\"\n        result = image.copy().astype(np.float32)\n        \n        avg_r = np.mean(result[:, :, 0])\n        avg_g = np.mean(result[:, :, 1])\n        avg_b = np.mean(result[:, :, 2])\n        \n        avg_gray = (avg_r + avg_g + avg_b) / 3\n        \n        if avg_r > 0:\n            result[:, :, 0] = np.clip(result[:, :, 0] * (avg_gray / avg_r), 0, 255)\n        if avg_g > 0:\n            result[:, :, 1] = np.clip(result[:, :, 1] * (avg_gray / avg_g), 0, 255)\n        if avg_b > 0:\n            result[:, :, 2] = np.clip(result[:, :, 2] * (avg_gray / avg_b), 0, 255)\n        \n        return result.astype(np.uint8)\n    \n    @staticmethod\n    def adaptive_histogram_equalization(image: np.ndarray) -> np.ndarray:\n        \"\"\"\n        CLAHE (Contrast Limited Adaptive Histogram Equalization)\n        \n        Args:\n            image: RGB image (H, W, 3)\n        \n        Returns:\n            CLAHE applied image (H, W, 3)\n        \"\"\"\n        lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)\n        l, a, b = cv2.split(lab)\n        \n        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n        l_clahe = clahe.apply(l)\n        \n        lab_clahe = cv2.merge([l_clahe, a, b])\n        rgb_clahe = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB)\n        \n        return rgb_clahe\n\n\nclass ImagePreprocessor:\n    \"\"\"Comprehensive image preprocessing class\"\"\"\n    \n    @staticmethod\n    def apply_hsv_mask(\n        image: np.ndarray,\n        lower_h: int = 35,\n        lower_s: int = 45,\n        lower_v: int = 45,\n        upper_h: int = 90,\n        upper_s: int = 200,\n        upper_v: int = 200\n    ) -> np.ndarray:\n        \"\"\"\n        Extract green vegetation using HSV masking\n        \n        Args:\n            image: BGR image\n            lower_h, lower_s, lower_v: HSV lower bounds\n            upper_h, upper_s, upper_v: HSV upper bounds\n        \n        Returns:\n            Masked image\n        \"\"\"\n        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n        lower_bound = np.array([lower_h, lower_s, lower_v])\n        upper_bound = np.array([upper_h, upper_s, upper_v])\n        mask = cv2.inRange(hsv, lower_bound, upper_bound)\n        result = cv2.bitwise_and(image, image, mask=mask)\n        return result\n    \n    @staticmethod\n    def preprocess_image(\n        image: np.ndarray,\n        config: Config,\n        is_train: bool = True\n    ) -> np.ndarray:\n        \"\"\"\n        Apply preprocessing based on Config settings\n        \n        Processing order:\n        1. Shadow removal (if enabled) - should be first\n        2. White balance correction\n        3. CLAHE\n        4. HSV masking (probabilistic during training)\n        5. Calculate additional features (vegetation indices, LAB, LBP)\n        \n        Args:\n            image: RGB image (H, W, 3)\n            config: Configuration object\n            is_train: Whether in training mode\n        \n        Returns:\n            Preprocessed image (H, W, C) where C depends on config\n        \"\"\"\n        # Phase 2: Shadow removal (should be applied first)\n        if config.use_shadow_removal:\n            image = ShadowRemoval.remove_shadow(image, method='hsv')\n        \n        # Phase 1: White balance correction\n        if config.use_white_balance:\n            image = AdaptivePreprocessor.white_balance_correction(image)\n        \n        # Phase 1: CLAHE\n        if config.use_clahe:\n            image = AdaptivePreprocessor.adaptive_histogram_equalization(image)\n        \n        # Phase 1: HSV masking (optional, probabilistic during training)\n        if config.use_hsv_masking and is_train and np.random.random() < 0.5:\n            image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n            image_bgr = ImagePreprocessor.apply_hsv_mask(image_bgr)\n            image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n        \n        # Collect all channels\n        channels = [image]  # Start with RGB (3 channels)\n        \n        # Phase 1: Vegetation indices\n        if config.use_vegetation_indices:\n            indices = VegetationIndices.calculate_all_indices(image)\n            channels.append(indices)\n        \n        # Phase 2: LAB color space\n        if config.use_lab_color_space:\n            lab = ColorSpaceTransforms.apply_lab_space(image)\n            channels.append(lab)\n        \n        # Phase 2: LBP texture\n        if config.use_lbp_texture:\n            lbp = TextureFeatures.calculate_lbp_simple(image)\n            channels.append(np.expand_dims(lbp, axis=-1))\n        \n        # Concatenate all channels\n        if len(channels) > 1:\n            multi_channel = np.concatenate(channels, axis=-1)\n            return multi_channel\n        else:\n            return image\n\n\n# ==================== Dataset Class (Updated) ====================\nclass BiomassDataset(Dataset):\n    \"\"\"Pasture Biomass Dataset with Phase 1 & 2 preprocessing\"\"\"\n    \n    def __init__(\n        self,\n        df: pd.DataFrame,\n        config: Config,\n        transform: Optional[transforms.Compose] = None,\n        is_train: bool = True\n    ):\n        self.df = df.reset_index(drop=True)\n        self.config = config\n        self.transform = transform\n        self.is_train = is_train\n        \n    def __len__(self) -> int:\n        return len(self.df)\n    \n    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:\n        row = self.df.iloc[idx]\n        \n        # Load image\n        image_path = os.path.join(self.config.input_dir, row['image_path'])\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Apply preprocessing\n        image = ImagePreprocessor.preprocess_image(\n            image,\n            self.config,\n            is_train=self.is_train\n        )\n        \n        # Handle multi-channel images\n        if image.shape[-1] > 3:\n            rgb_image = image[:, :, :3]\n            extra_channels = image[:, :, 3:]\n            \n            rgb_pil = Image.fromarray(rgb_image.astype(np.uint8))\n            if self.transform:\n                rgb_tensor = self.transform(rgb_pil)\n            else:\n                rgb_tensor = transforms.ToTensor()(rgb_pil)\n            \n            extra_tensor = torch.from_numpy(extra_channels.transpose(2, 0, 1)).float()\n            extra_tensor = extra_tensor / 255.0\n            \n            if self.transform:\n                extra_tensor = torch.nn.functional.interpolate(\n                    extra_tensor.unsqueeze(0),\n                    size=self.config.image_size,\n                    mode='bilinear',\n                    align_corners=False\n                ).squeeze(0)\n            \n            image_tensor = torch.cat([rgb_tensor, extra_tensor], dim=0)\n        else:\n            image_pil = Image.fromarray(image.astype(np.uint8))\n            if self.transform:\n                image_tensor = self.transform(image_pil)\n            else:\n                image_tensor = transforms.ToTensor()(image_pil)\n        \n        # Log-transform targets\n        targets = []\n        for col in self.config.target_cols:\n            value = row[col]\n            log_value = np.log1p(value)\n            targets.append(log_value)\n        \n        targets = torch.tensor(targets, dtype=torch.float32)\n        \n        return image_tensor, targets\n\n\n# ==================== Data Augmentation ====================\ndef get_transforms(config: Config, is_train: bool = True) -> transforms.Compose:\n    \"\"\"Get image transformation pipeline\"\"\"\n    if is_train:\n        return transforms.Compose([\n            transforms.Resize(config.image_size),\n            transforms.RandomHorizontalFlip(p=0.5),\n            transforms.RandomVerticalFlip(p=0.5),\n            transforms.RandomRotation(15),\n            transforms.ColorJitter(\n                brightness=0.2,\n                contrast=0.2,\n                saturation=0.2,\n                hue=0.1\n            ),\n            transforms.ToTensor(),\n            transforms.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225]\n            )\n        ])\n    else:\n        return transforms.Compose([\n            transforms.Resize(config.image_size),\n            transforms.ToTensor(),\n            transforms.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225]\n            )\n        ])\n\n\n# ==================== Model Definition (Updated) ====================\nclass BiomassModel(nn.Module):\n    \"\"\"Multi-target regression model with flexible input channels\"\"\"\n    \n    def __init__(self, config: Config, pretrained: bool = True):\n        super().__init__()\n        self.config = config\n        \n        n_input_channels = config.total_input_channels\n        \n        if n_input_channels == 3:\n            self.backbone = timm.create_model(\n                'tf_efficientnetv2_s',\n                pretrained=pretrained,\n                num_classes=0,\n                global_pool='avg'\n            )\n        else:\n            self.backbone = timm.create_model(\n                'tf_efficientnetv2_s',\n                pretrained=False,\n                num_classes=0,\n                global_pool='avg',\n                in_chans=n_input_channels\n            )\n            \n            if pretrained:\n                pretrained_model = timm.create_model(\n                    'tf_efficientnetv2_s',\n                    pretrained=True,\n                    num_classes=0\n                )\n                \n                pretrained_dict = pretrained_model.state_dict()\n                model_dict = self.backbone.state_dict()\n                \n                pretrained_dict = {\n                    k: v for k, v in pretrained_dict.items()\n                    if k in model_dict and 'conv_stem' not in k\n                }\n                \n                model_dict.update(pretrained_dict)\n                self.backbone.load_state_dict(model_dict, strict=False)\n                \n                print(f\"Loaded pretrained weights (except first conv layer)\")\n                print(f\"Input channels: {n_input_channels}\")\n        \n        n_features = self.backbone.num_features\n        \n        self.head = nn.Sequential(\n            nn.Linear(n_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 256),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(256, len(config.target_cols))\n        )\n    \n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        features = self.backbone(x)\n        predictions = self.head(features)\n        return predictions\n\n\n# ==================== Loss Function ====================\nclass WeightedMSELoss(nn.Module):\n    \"\"\"Weighted MSE loss\"\"\"\n    \n    def __init__(self, config: Config):\n        super().__init__()\n        weights = [config.target_weights[col] for col in config.target_cols]\n        self.weights = torch.tensor(weights, dtype=torch.float32)\n    \n    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:\n        self.weights = self.weights.to(predictions.device)\n        mse = (predictions - targets) ** 2\n        weighted_mse = mse * self.weights\n        return weighted_mse.mean()\n\n\n# ==================== Evaluation Metrics ====================\ndef calculate_weighted_r2(\n    predictions: np.ndarray,\n    targets: np.ndarray,\n    config: Config\n) -> tuple[float, dict[str, float]]:\n    \"\"\"Calculate weighted R² score\"\"\"\n    r2_dict = {}\n    weighted_r2 = 0.0\n    \n    for i, col in enumerate(config.target_cols):\n        pred = predictions[:, i]\n        true = targets[:, i]\n        \n        ss_res = np.sum((true - pred) ** 2)\n        ss_tot = np.sum((true - true.mean()) ** 2)\n        r2 = 1 - (ss_res / ss_tot) if ss_tot > 0 else 0\n        \n        r2_dict[col] = r2\n        weighted_r2 += r2 * config.target_weights[col]\n    \n    return weighted_r2, r2_dict\n\n\n# ==================== Model Saving Utilities (NEW) ====================\nclass ModelSaver:\n    \"\"\"Utility class for saving trained models\"\"\"\n    \n    @staticmethod\n    def save_model(\n        model: nn.Module,\n        config: Config,\n        fold: int,\n        score: float,\n        epoch: Optional[int] = None\n    ) -> str:\n        \"\"\"\n        Save model checkpoint\n        \n        Args:\n            model: Model to save\n            config: Configuration object\n            fold: Fold number\n            score: Validation score\n            epoch: Epoch number (optional)\n        \n        Returns:\n            Path to saved model file\n        \"\"\"\n        if not config.save_models:\n            return \"\"\n        \n        # Create filename\n        # timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n        # preprocessing_tag = config.get_preprocessing_summary().replace(\", \", \"_\").replace(\" \", \"\")\n        \n        filename = f\"model_fold{fold}.pt\"\n        \n        filepath = os.path.join(config.full_model_save_path, filename)\n        \n        # Prepare checkpoint\n        checkpoint = {\n            'model_state_dict': model.state_dict(),\n            'config': config,\n            'fold': fold,\n            'score': score,\n            'epoch': epoch,\n            'preprocessing_summary': config.get_preprocessing_summary(),\n            'timestamp': timestamp\n        }\n        \n        # Save\n        torch.save(checkpoint, filepath)\n        print(f\"  💾 Model saved: {filename}\")\n        \n        return filepath\n    \n    @staticmethod\n    def load_model(filepath: str, device: str = 'cuda') -> tuple[nn.Module, Config, dict]:\n        \"\"\"\n        Load model checkpoint\n        \n        Args:\n            filepath: Path to checkpoint file\n            device: Device to load model on\n        \n        Returns:\n            model: Loaded model\n            config: Configuration object\n            metadata: Checkpoint metadata\n        \"\"\"\n        checkpoint = torch.load(filepath, map_location=device)\n        \n        config = checkpoint['config']\n        model = BiomassModel(config, pretrained=False)\n        model.load_state_dict(checkpoint['model_state_dict'])\n        model.to(device)\n        \n        metadata = {\n            'fold': checkpoint.get('fold'),\n            'score': checkpoint.get('score'),\n            'epoch': checkpoint.get('epoch'),\n            'preprocessing_summary': checkpoint.get('preprocessing_summary'),\n            'timestamp': checkpoint.get('timestamp')\n        }\n        \n        return model, config, metadata\n\n\n# ==================== Training Loop (Updated with Model Saving) ====================\nclass Trainer:\n    \"\"\"Manages model training and evaluation\"\"\"\n    \n    def __init__(self, config: Config):\n        self.config = config\n        self.device = torch.device(config.device)\n    \n    def train_epoch(\n        self,\n        model: nn.Module,\n        dataloader: DataLoader,\n        criterion: nn.Module,\n        optimizer: torch.optim.Optimizer,\n        epoch: int\n    ) -> float:\n        \"\"\"Train for one epoch\"\"\"\n        model.train()\n        total_loss = 0.0\n        \n        # tqdmの設定を修正\n        pbar = tqdm(\n            dataloader,\n            desc=f'Epoch {epoch+1} Train',\n            leave=False,           # 完了後に消す\n            dynamic_ncols=True,    # 動的に幅を調整\n            position=1             # 位置を指定（Fold barの下）\n        )\n        \n        for images, targets in pbar:\n            images = images.to(self.device)\n            targets = targets.to(self.device)\n            \n            outputs = model(images)\n            loss = criterion(outputs, targets)\n            \n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            \n            total_loss += loss.item()\n            \n            # リアルタイムでlossを更新\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'}, refresh=True)\n        \n        pbar.close()  # 明示的にクローズ\n        return total_loss / len(dataloader)\n    \n    def validate(\n        self,\n        model: nn.Module,\n        dataloader: DataLoader,\n        criterion: nn.Module,\n        epoch: int\n    ) -> tuple[float, float, dict[str, float]]:\n        \"\"\"Validation\"\"\"\n        model.eval()\n        total_loss = 0.0\n        all_predictions = []\n        all_targets = []\n        \n        # tqdmの設定を修正\n        pbar = tqdm(\n            dataloader,\n            desc=f'Epoch {epoch+1} Valid',\n            leave=False,\n            dynamic_ncols=True,\n            position=1\n        )\n        \n        with torch.no_grad():\n            for images, targets in pbar:\n                images = images.to(self.device)\n                targets = targets.to(self.device)\n                \n                outputs = model(images)\n                loss = criterion(outputs, targets)\n                \n                total_loss += loss.item()\n                all_predictions.append(outputs.cpu().numpy())\n                all_targets.append(targets.cpu().numpy())\n                \n                pbar.set_postfix({'loss': f'{loss.item():.4f}'}, refresh=True)\n        \n        pbar.close()  # 明示的にクローズ\n        \n        all_predictions = np.concatenate(all_predictions, axis=0)\n        all_targets = np.concatenate(all_targets, axis=0)\n        \n        weighted_r2, r2_dict = calculate_weighted_r2(\n            all_predictions,\n            all_targets,\n            self.config\n        )\n        \n        avg_loss = total_loss / len(dataloader)\n        return avg_loss, weighted_r2, r2_dict\n    \n    def train_fold(\n        self,\n        train_df: pd.DataFrame,\n        val_df: pd.DataFrame,\n        fold: int\n    ) -> tuple[nn.Module, float]:\n        \"\"\"Train one fold\"\"\"\n        print(f'\\n{\"=\"*60}')\n        print(f'Fold {fold + 1} / {self.config.n_folds}')\n        print(f'Preprocessing: {self.config.get_preprocessing_summary()}')\n        print(f'Input channels: {self.config.total_input_channels}')\n        print(f'{\"=\"*60}\\n')\n        \n        train_dataset = BiomassDataset(\n            train_df,\n            self.config,\n            transform=get_transforms(self.config, is_train=True),\n            is_train=True\n        )\n        val_dataset = BiomassDataset(\n            val_df,\n            self.config,\n            transform=get_transforms(self.config, is_train=False),\n            is_train=False\n        )\n        \n        train_loader = DataLoader(\n            train_dataset,\n            batch_size=self.config.batch_size,\n            shuffle=True,\n            num_workers=2,\n            pin_memory=True\n        )\n        val_loader = DataLoader(\n            val_dataset,\n            batch_size=self.config.batch_size,\n            shuffle=False,\n            num_workers=2,\n            pin_memory=True\n        )\n        \n        model = BiomassModel(self.config).to(self.device)\n        criterion = WeightedMSELoss(self.config).to(self.device)\n        optimizer = torch.optim.AdamW(model.parameters(), lr=self.config.learning_rate)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            optimizer,\n            T_max=self.config.epochs\n        )\n        \n        best_score = -np.inf\n        best_model_state = None\n        \n        # エポックループのプログレスバー\n        epoch_pbar = tqdm(\n            range(self.config.epochs),\n            desc=f'Fold {fold+1}/{self.config.n_folds}',\n            position=0,           # 最上位\n            leave=True,           # 残す\n            dynamic_ncols=True\n        )\n        \n        for epoch in epoch_pbar:\n            # train_epochとvalidateにepochを渡す\n            train_loss = self.train_epoch(model, train_loader, criterion, optimizer, epoch)\n            val_loss, val_r2, r2_dict = self.validate(model, val_loader, criterion, epoch)\n            scheduler.step()\n            \n            # エポックバーの情報を更新\n            epoch_pbar.set_postfix({\n                'train_loss': f'{train_loss:.4f}',\n                'val_loss': f'{val_loss:.4f}',\n                'val_r2': f'{val_r2:.4f}',\n                'best_r2': f'{best_score:.4f}' if best_score > -np.inf else 'N/A'\n            }, refresh=True)\n            \n            # ベストモデル更新\n            if val_r2 > best_score:\n                best_score = val_r2\n                best_model_state = model.state_dict().copy()\n                \n                # tqdm.write()を使って出力（プログレスバーを崩さない）\n                tqdm.write(f'  ✓ Epoch {epoch+1:02d}: Best updated! R²={best_score:.4f}')\n                \n                # f-stringのネストを回避\n                target_scores = \", \".join([f\"{k.split('_')[1]}: {v:.3f}\" for k, v in r2_dict.items()])\n                tqdm.write(f'    └─ {target_scores}')\n        # エポックバーを閉じる\n        epoch_pbar.close()\n        \n        # Load best model\n        model.load_state_dict(best_model_state)\n        \n        # Save best model\n        if self.config.save_models and self.config.save_best_only:\n            ModelSaver.save_model(model, self.config, fold, best_score)\n        \n        # 最終結果を表示\n        print(f'\\nFold {fold+1} completed: Best R² = {best_score:.4f}\\n')\n        \n        return model, best_score\n\n\n# ==================== Main Execution ====================\ndef main():\n    \"\"\"Main processing\"\"\"\n    \n    # Initialize configuration\n    config = Config()\n    \n    # ========== Phase 1 & 2 Preprocessing Toggles ==========\n    # Phase 1\n    config.use_white_balance = True\n    config.use_clahe = True\n    config.use_vegetation_indices = True\n    config.use_hsv_masking = True\n    \n    # Phase 2\n    config.use_shadow_removal = False      # NEW: Shadow removal\n    config.use_lab_color_space = False     # NEW: LAB color space\n    config.use_lbp_texture = False         # NEW: LBP texture\n    \n    # Model saving\n    config.save_models = True             # NEW: Enable model saving\n    config.save_best_only = True          # NEW: Save only best models\n    # =======================================================\n    \n    print(\"=\"*60)\n    print(\"Configuration Summary:\")\n    print(f\"  Phase 1 Preprocessing:\")\n    print(f\"    - White Balance: {config.use_white_balance}\")\n    print(f\"    - CLAHE: {config.use_clahe}\")\n    print(f\"    - Vegetation Indices: {config.use_vegetation_indices}\")\n    print(f\"    - HSV Masking: {config.use_hsv_masking}\")\n    print(f\"  Phase 2 Preprocessing:\")\n    print(f\"    - Shadow Removal: {config.use_shadow_removal}\")\n    print(f\"    - LAB Color Space: {config.use_lab_color_space}\")\n    print(f\"    - LBP Texture: {config.use_lbp_texture}\")\n    print(f\"  Total Input Channels: {config.total_input_channels}\")\n    print(f\"  Model Saving: {config.save_models}\")\n    if config.save_models:\n        print(f\"  Save Directory: {config.full_model_save_path}\")\n    print(\"=\"*60)\n    \n    # Load data\n    print('\\nLoading data...')\n    train = pd.read_csv(os.path.join(config.input_dir, 'train.csv'))\n    \n    train_grouped = train.pivot_table(\n        index=['image_path', 'Sampling_Date', 'State', 'Species', \n               'Pre_GSHH_NDVI', 'Height_Ave_cm'],\n        columns='target_name',\n        values='target',\n        aggfunc='first'\n    ).reset_index()\n    \n    print(f'Total samples: {len(train_grouped)}')\n    print(f'Targets: {config.target_cols}')\n    \n    # Cross Validation\n    kfold = KFold(n_splits=config.n_folds, shuffle=True, random_state=42)\n    trainer = Trainer(config)\n    \n    fold_scores = []\n    saved_models = []\n    \n    for fold, (train_idx, val_idx) in enumerate(kfold.split(train_grouped)):\n        train_df = train_grouped.iloc[train_idx]\n        val_df = train_grouped.iloc[val_idx]\n        \n        model, score = trainer.train_fold(train_df, val_df, fold)\n        fold_scores.append(score)\n        \n        # Free memory\n        del model\n        torch.cuda.empty_cache()\n        gc.collect()\n    \n    # CV results summary\n    print(f'\\n{\"=\"*60}')\n    print('Cross Validation Results')\n    print(f'{\"=\"*60}')\n    for i, score in enumerate(fold_scores):\n        print(f'Fold {i+1}: {score:.4f}')\n    print(f'Mean: {np.mean(fold_scores):.4f} ± {np.std(fold_scores):.4f}')\n    print(f'{\"=\"*60}')\n    \n    if config.save_models:\n        print(f'\\n✅ All models saved to: {config.full_model_save_path}')\n\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-20T00:55:17.006792Z","iopub.execute_input":"2025-11-20T00:55:17.007459Z","execution_failed":"2025-11-20T00:56:24.904Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}