{"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":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"I modified the learning model, visualization PNG, and config output settings.","metadata":{}},{"cell_type":"markdown","source":"### Inference Notebook: https://www.kaggle.com/code/takahitomizunobyts/lb-0-61-csiro-infer\n### EDA Notebook: https://www.kaggle.com/code/takahitomizunobyts/csiro-customizable-eda","metadata":{}},{"cell_type":"code","source":"\"\"\"\nCSIRO Biomass Competition - Two-Stream Three-Head Regression Model with Visualization\n================================================================================\nThis script implements a pipeline for predicting biomass from images with comprehensive visualization.\n\nPipeline Overview:\n1. Data Preparation: Load CSV → Pivot → Stratified K-Fold split\n2. Preprocessing: Image split (left/right) → Augmentation → Normalization\n3. Model: Shared Backbone → Feature concatenation → 3 dedicated heads\n4. Training: Two-stage learning (Freeze→Unfreeze) → Weighted loss → R² evaluation\n5. Visualization: Learning curves, Fold comparison, Prediction scatter plots\n\"\"\"\n\nfrom __future__ import annotations\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import Callable, Optional\nimport os\nimport gc\nimport time\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nimport cv2\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import r2_score\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n\n# ============================================================================\n# Configuration Management\n# ============================================================================\n\n@dataclass\nclass Config:\n    \"\"\"\n    Data class for managing pipeline-wide configuration.\n    \n    Centralizing all hyperparameters and constants ensures\n    experiment reproducibility and facilitates configuration changes.\n    \"\"\"\n    \n    # Path settings\n    base_path: Path = Path('/kaggle/input/csiro-biomass')\n    train_csv: Path = field(init=False)\n    image_dir: Path = field(init=False)\n    output_dir: Path = Path('/kaggle/working/')  # Directory for saving visualizations\n    \n    # Model settings\n    model_name: str = 'convnext_tiny'  # timm-compatible model name\n    pretrained: bool = True\n    img_size: int = 500\n    \n    # Device settings\n    device: torch.device = field(default_factory=lambda: torch.device(\n        'cuda' if torch.cuda.is_available() else 'cpu'\n    ))\n    \n    # Training settings\n    batch_size: int = 4  # Adjust according to GPU memory\n    epochs: int = 30\n    learning_rate: float = 1e-4  # Learning rate for Stage 1\n    finetune_lr: float = 1e-5     # Learning rate for Stage 2\n    freeze_epochs: int = 10        # Number of epochs to freeze backbone\n    num_workers: int = 2\n    \n    # Cross-validation settings\n    n_folds: int = 5\n    random_state: int = 42\n    \n    # Target settings\n    # Three targets used for training (loss calculation)\n    train_target_cols: list[str] = field(default_factory=lambda: [\n        'Dry_Total_g', 'GDM_g', 'Dry_Green_g'\n    ])\n    \n    # Five targets used for evaluation (R² score calculation)\n    all_target_cols: list[str] = field(default_factory=lambda: [\n        'Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g'\n    ])\n    \n    # Loss function weights (corresponding to 3 training targets)\n    loss_weights: dict[str, float] = field(default_factory=lambda: {\n        'total_loss': 0.5,\n        'gdm_loss': 0.2,\n        'green_loss': 0.1\n    })\n    \n    # R² score weights (corresponding to 5 evaluation targets)\n    r2_weights: list[float] = field(default_factory=lambda: [\n        0.1, 0.1, 0.1, 0.2, 0.5\n    ])\n    \n    def __post_init__(self) -> None:\n        \"\"\"Construct paths after initialization and create output directory\"\"\"\n        self.train_csv = self.base_path / 'train.csv'\n        self.image_dir = self.base_path / 'train'\n        \n        # Create output directory for visualizations\n        self.output_dir.mkdir(parents=True, exist_ok=True)\n        \n    def display_info(self) -> None:\n        \"\"\"Display configuration information\"\"\"\n        print(f\"{'='*70}\")\n        print(f\"Configuration\")\n        print(f\"{'='*70}\")\n        print(f\"Device: {self.device}\")\n        print(f\"Backbone: {self.model_name}\")\n        print(f\"Image Size: {self.img_size}x{self.img_size}\")\n        print(f\"Batch Size: {self.batch_size}\")\n        print(f\"Total Epochs: {self.epochs}\")\n        print(f\"  - Stage 1 (Freeze): Epoch 1-{self.freeze_epochs} (LR={self.learning_rate})\")\n        print(f\"  - Stage 2 (Finetune): Epoch {self.freeze_epochs+1}-{self.epochs} (LR={self.finetune_lr})\")\n        print(f\"Cross-Validation: {self.n_folds}-Fold\")\n        print(f\"Output Directory: {self.output_dir}\")\n        print(f\"{'='*70}\\n\")\n\n\n# ============================================================================\n# Data Preparation\n# ============================================================================\n\nclass DataPreparator:\n    \"\"\"\n    Class responsible for data loading and preprocessing.\n    \n    Main functions:\n    - Load and pivot CSV\n    - Stratified K-Fold splitting\n    \"\"\"\n    \n    def __init__(self, config: Config):\n        \"\"\"\n        Args:\n            config: Configuration object\n        \"\"\"\n        self.config = config\n        self.df_wide: Optional[pd.DataFrame] = None\n        \n    def load_and_pivot(self) -> pd.DataFrame:\n        \"\"\"\n        Load CSV and convert from long format to wide format.\n        \n        Returns:\n            Wide-format DataFrame (each row is one image, each column is one target)\n            \n        Why not: Using pivot() instead of pivot_table()\n            → Image paths are guaranteed to be unique\n        \"\"\"\n        print(f\"Loading CSV: {self.config.train_csv}\")\n        \n        try:\n            df_long = pd.read_csv(self.config.train_csv)\n            print(f\"Long format: {len(df_long)} rows\")\n            \n            # Pivot transformation: image_path × target_name → values\n            df_wide = df_long.pivot(\n                index='image_path',\n                columns='target_name',\n                values='target'\n            ).reset_index()\n            \n            df_wide.columns.name = None  # Clean up column names\n            print(f\"Wide format: {len(df_wide)} rows × {len(df_wide.columns)} columns\")\n            print(f\"\\nFirst 5 rows:\\n{df_wide.head()}\\n\")\n            \n            self.df_wide = df_wide\n            return df_wide\n            \n        except FileNotFoundError:\n            print(f\"Error: {self.config.train_csv} not found\")\n            # Return dummy DataFrame on error (prevent downstream crashes)\n            return pd.DataFrame(columns=['image_path'] + self.config.all_target_cols)\n    \n    def create_stratified_folds(self, df: pd.DataFrame) -> pd.DataFrame:\n        \"\"\"\n        Assign fold numbers for cross-validation using Stratified K-Fold.\n        \n        Args:\n            df: Wide-format DataFrame\n            \n        Returns:\n            DataFrame with 'fold' column added\n            \n        Why not: Using StratifiedKFold instead of GroupKFold\n            → For regression, stratify after binning to maintain target distribution\n        \"\"\"\n        print(f\"\\nPreparing {self.config.n_folds}-Fold Cross-Validation...\")\n        \n        df = df.copy()\n        df['fold'] = -1\n        \n        # Bin targets (continuous → discrete)\n        # Determine number of bins using Sturges' formula\n        num_bins = min(10, int(np.floor(1 + np.log2(len(df)))))\n        print(f\"Stratifying Dry_Total_g into {num_bins} bins\")\n        \n        df['total_bin'] = pd.cut(\n            df['Dry_Total_g'], \n            bins=num_bins, \n            labels=False,\n            duplicates='drop'  # Remove duplicate edges\n        )\n        \n        # Stratified K-Fold split\n        skf = StratifiedKFold(\n            n_splits=self.config.n_folds,\n            shuffle=True,\n            random_state=self.config.random_state\n        )\n        \n        for fold_num, (_, valid_idx) in enumerate(skf.split(df, df['total_bin'])):\n            df.loc[valid_idx, 'fold'] = fold_num\n        \n        # Remove binning column (no longer needed)\n        df = df.drop(columns=['total_bin'])\n        \n        print(\"\\nFold distribution:\")\n        print(df['fold'].value_counts().sort_index())\n        \n        self.df_wide = df\n        return df\n\n\n# ============================================================================\n# Data Augmentation\n# ============================================================================\n\nclass AugmentationFactory:\n    \"\"\"\n    Factory class for generating Albumentations pipelines.\n    \n    Provides different pipelines for training and validation.\n    \"\"\"\n    \n    def __init__(self, img_size: int):\n        \"\"\"\n        Args:\n            img_size: Image size after resizing\n        \"\"\"\n        self.img_size = img_size\n    \n    def get_train_transforms(self) -> A.Compose:\n        \"\"\"\n        Augmentation pipeline for training.\n        \n        Returns:\n            Albumentations Compose object\n            \n        Why not: Not adding stronger augmentations\n            → Balancing overfitting risk and training time\n        \"\"\"\n        return A.Compose([\n            # Geometric transforms\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            \n            # Color transforms\n            A.ColorJitter(\n                brightness=0.2,\n                contrast=0.2,\n                saturation=0.2,\n                hue=0.1,\n                p=0.75\n            ),\n            \n            # Resize and normalize\n            A.Resize(self.img_size, self.img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2()\n        ])\n    \n    def get_valid_transforms(self) -> A.Compose:\n        \"\"\"\n        Pipeline for validation (no augmentation).\n        \n        Returns:\n            Albumentations Compose object\n        \"\"\"\n        return A.Compose([\n            A.Resize(self.img_size, self.img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2()\n        ])\n\n\n# ============================================================================\n# Dataset\n# ============================================================================\n\nclass BiomassDataset(Dataset):\n    \"\"\"\n    Custom Dataset for two-stream architecture.\n    \n    Splits 2000x1000 images into left and right halves,\n    applying independent augmentation to each.\n    \n    Returns:\n        tuple: (img_left, img_right, train_targets, all_targets)\n            - img_left: Left image tensor [C, H, W]\n            - img_right: Right image tensor [C, H, W]\n            - train_targets: Training targets [3]\n            - all_targets: Evaluation targets [5]\n    \"\"\"\n    \n    def __init__(\n        self,\n        df: pd.DataFrame,\n        transforms_fn: Callable[[], A.Compose],\n        image_dir: Path,\n        train_target_cols: list[str],\n        all_target_cols: list[str]\n    ):\n        \"\"\"\n        Args:\n            df: DataFrame containing image paths and targets\n            transforms_fn: Function returning augmentation pipeline\n            image_dir: Image directory path\n            train_target_cols: Training target column names\n            all_target_cols: Evaluation target column names\n        \"\"\"\n        self.df = df\n        self.transforms_fn = transforms_fn\n        self.image_dir = image_dir\n        \n        # Convert to numpy arrays for fast access\n        self.image_paths = df['image_path'].values\n        self.train_targets = df[train_target_cols].values\n        self.all_targets = df[all_target_cols].values\n    \n    def __len__(self) -> int:\n        return len(self.df)\n    \n    def __getitem__(\n        self, \n        idx: int\n    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Get one sample.\n        \n        Args:\n            idx: Sample index\n            \n        Returns:\n            (left_image, right_image, train_targets, eval_targets)\n            \n        Why not: Not pre-splitting and saving images\n            → Considering trade-off between storage and I/O time\n        \"\"\"\n        # Get image path and targets\n        img_path = self.image_paths[idx]\n        train_target = self.train_targets[idx]\n        all_target = self.all_targets[idx]\n        \n        # Load image\n        full_path = self.image_dir / Path(img_path).name\n        image = cv2.imread(str(full_path))\n        \n        if image is None:\n            raise FileNotFoundError(f\"Failed to load image: {full_path}\")\n        \n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Split into left and right (2000x1000 → 2 x 1000x1000)\n        height, width = image.shape[:2]\n        mid_point = width // 2\n        img_left = image[:, :mid_point]\n        img_right = image[:, mid_point:]\n        \n        # Apply augmentation independently\n        # Why not: Not reusing the same transform\n        #   → Applying different augmentations to left/right improves data diversity\n        transform_left = self.transforms_fn()\n        transform_right = self.transforms_fn()\n        \n        img_left_tensor = transform_left(image=img_left)['image']\n        img_right_tensor = transform_right(image=img_right)['image']\n        \n        # Convert targets to tensors\n        train_target_tensor = torch.tensor(train_target, dtype=torch.float32)\n        all_target_tensor = torch.tensor(all_target, dtype=torch.float32)\n        \n        return img_left_tensor, img_right_tensor, train_target_tensor, all_target_tensor\n\n\n# ============================================================================\n# Model\n# ============================================================================\n\nclass BiomassModel(nn.Module):\n    \"\"\"\n    Two-stream, three-head regression model.\n    \n    Architecture:\n    1. Shared Backbone (ConvNeXt etc.) extracts features from left/right images\n    2. Concatenate two feature vectors\n    3. Three dedicated MLP heads predict each target\n    \"\"\"\n    \n    def __init__(self, model_name: str, pretrained: bool):\n        \"\"\"\n        Args:\n            model_name: timm model name\n            pretrained: Whether to use ImageNet pretrained weights\n        \"\"\"\n        super().__init__()\n        \n        # Shared Backbone (no classification layer, with Global Average Pooling)\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=pretrained,\n            num_classes=0,      # Remove classification layer\n            global_pool='avg'   # Add GAP\n        )\n        \n        # Get feature dimension\n        self.n_features = self.backbone.num_features\n        self.n_combined = self.n_features * 2  # Dimension after concatenation\n        \n        # Three dedicated heads\n        # Why not: Not using a single shared head\n        #   → Dedicated heads improve accuracy as each target has different characteristics\n        self.head_total = self._create_head()\n        self.head_gdm = self._create_head()\n        self.head_green = self._create_head()\n    \n    def _create_head(self) -> nn.Sequential:\n        \"\"\"\n        Generate MLP structure for a single head.\n        \n        Returns:\n            Two-layer MLP (with ReLU + Dropout in middle layer)\n        \"\"\"\n        return nn.Sequential(\n            nn.Linear(self.n_combined, self.n_combined // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined // 2, 1)\n        )\n    \n    def forward(\n        self, \n        img_left: torch.Tensor, \n        img_right: torch.Tensor\n    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Forward pass.\n        \n        Args:\n            img_left: Left image [B, C, H, W]\n            img_right: Right image [B, C, H, W]\n            \n        Returns:\n            (total_pred, gdm_pred, green_pred) each [B, 1]\n        \"\"\"\n        # Feature extraction (shared backbone)\n        feat_left = self.backbone(img_left)    # [B, n_features]\n        feat_right = self.backbone(img_right)  # [B, n_features]\n        \n        # Concatenate features\n        combined = torch.cat([feat_left, feat_right], dim=1)  # [B, n_combined]\n        \n        # Predict with each head\n        out_total = self.head_total(combined)\n        out_gdm = self.head_gdm(combined)\n        out_green = self.head_green(combined)\n        \n        return out_total, out_gdm, out_green\n\n\n# ============================================================================\n# Loss Function\n# ============================================================================\n\nclass WeightedBiomassLoss(nn.Module):\n    \"\"\"\n    Weighted loss for three targets.\n    \n    Weight losses according to importance of each target,\n    using SmoothL1Loss for robust learning against outliers.\n    \"\"\"\n    \n    def __init__(self, loss_weights: dict[str, float]):\n        \"\"\"\n        Args:\n            loss_weights: Weights for each target\n        \"\"\"\n        super().__init__()\n        self.criterion = nn.SmoothL1Loss()  # A variant of Huber loss\n        self.weights = loss_weights\n    \n    def forward(\n        self,\n        predictions: tuple[torch.Tensor, torch.Tensor, torch.Tensor],\n        targets: torch.Tensor\n    ) -> torch.Tensor:\n        \"\"\"\n        Calculate weighted loss.\n        \n        Args:\n            predictions: (total, gdm, green) each [B, 1]\n            targets: [B, 3] (in order: total, gdm, green)\n            \n        Returns:\n            Weighted total loss\n        \"\"\"\n        pred_total, pred_gdm, pred_green = predictions\n        \n        # Decompose targets\n        true_total = targets[:, 0:1]  # Maintain [B, 1] shape\n        true_gdm = targets[:, 1:2]\n        true_green = targets[:, 2:3]\n        \n        # Calculate each loss\n        loss_total = self.criterion(pred_total, true_total)\n        loss_gdm = self.criterion(pred_gdm, true_gdm)\n        loss_green = self.criterion(pred_green, true_green)\n        \n        # Weighted sum\n        total_loss = (\n            self.weights['total_loss'] * loss_total +\n            self.weights['gdm_loss'] * loss_gdm +\n            self.weights['green_loss'] * loss_green\n        )\n        \n        return total_loss\n\n\n# ============================================================================\n# Evaluation Metrics\n# ============================================================================\n\nclass CompetitionScorer:\n    \"\"\"\n    Class for calculating competition evaluation metric (weighted R²).\n    \n    Reconstructs five targets from three predictions,\n    then calculates weighted average of R² scores for each target.\n    \"\"\"\n    \n    def __init__(self, r2_weights: list[float]):\n        \"\"\"\n        Args:\n            r2_weights: R² weights for five targets\n        \"\"\"\n        self.r2_weights = np.array(r2_weights)\n    \n    def calculate_score(\n        self,\n        preds_dict: dict[str, np.ndarray],\n        targets_5: np.ndarray\n    ) -> float:\n        \"\"\"\n        Calculate weighted R² score.\n        \n        Args:\n            preds_dict: {'total': [N], 'gdm': [N], 'green': [N]}\n            targets_5: [N, 5] (in order: green, dead, clover, gdm, total)\n            \n        Returns:\n            Weighted R² score\n            \n        Why not: Not using simple MSE\n            → Following competition evaluation rules\n        \"\"\"\n        # Get predictions\n        pred_total = preds_dict['total']\n        pred_gdm = preds_dict['gdm']\n        pred_green = preds_dict['green']\n        \n        # Estimate remaining two (clip negative values)\n        pred_clover = np.maximum(0, pred_gdm - pred_green)\n        pred_dead = np.maximum(0, pred_total - pred_gdm)\n        \n        # Combine five predictions (green, dead, clover, gdm, total)\n        y_preds = np.stack([\n            pred_green, pred_dead, pred_clover, pred_gdm, pred_total\n        ], axis=1)\n        \n        # Calculate R² for each target\n        r2_scores = r2_score(targets_5, y_preds, multioutput='raw_values')\n        \n        # Weighted sum\n        weighted_score = np.sum(r2_scores * self.r2_weights)\n        \n        return float(weighted_score)\n    \n    def calculate_individual_scores(\n        self,\n        preds_dict: dict[str, np.ndarray],\n        targets_5: np.ndarray\n    ) -> dict[str, float]:\n        \"\"\"\n        Calculate individual R² scores for each target.\n        \n        Args:\n            preds_dict: {'total': [N], 'gdm': [N], 'green': [N]}\n            targets_5: [N, 5] (in order: green, dead, clover, gdm, total)\n            \n        Returns:\n            Dictionary of R² scores for each target\n        \"\"\"\n        # Get predictions\n        pred_total = preds_dict['total']\n        pred_gdm = preds_dict['gdm']\n        pred_green = preds_dict['green']\n        \n        # Estimate remaining two (clip negative values)\n        pred_clover = np.maximum(0, pred_gdm - pred_green)\n        pred_dead = np.maximum(0, pred_total - pred_gdm)\n        \n        # Combine five predictions (green, dead, clover, gdm, total)\n        y_preds = np.stack([\n            pred_green, pred_dead, pred_clover, pred_gdm, pred_total\n        ], axis=1)\n        \n        # Calculate R² for each target\n        r2_scores = r2_score(targets_5, y_preds, multioutput='raw_values')\n        \n        target_names = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n        \n        return {name: float(score) for name, score in zip(target_names, r2_scores)}\n\n\n# ============================================================================\n# Training History Tracker\n# ============================================================================\n\n@dataclass\nclass TrainingHistory:\n    \"\"\"\n    Data class for tracking training history.\n    \n    Stores losses, scores, and predictions for visualization.\n    \"\"\"\n    train_losses: list[float] = field(default_factory=list)\n    valid_losses: list[float] = field(default_factory=list)\n    valid_scores: list[float] = field(default_factory=list)\n    \n    # For final validation predictions\n    final_preds: Optional[dict[str, np.ndarray]] = None\n    final_targets: Optional[np.ndarray] = None\n    \n    def add_epoch(\n        self,\n        train_loss: float,\n        valid_loss: float,\n        valid_score: float\n    ) -> None:\n        \"\"\"Add one epoch's metrics.\"\"\"\n        self.train_losses.append(train_loss)\n        self.valid_losses.append(valid_loss)\n        self.valid_scores.append(valid_score)\n    \n    def set_final_predictions(\n        self,\n        preds: dict[str, np.ndarray],\n        targets: np.ndarray\n    ) -> None:\n        \"\"\"Store final predictions for scatter plots.\"\"\"\n        self.final_preds = preds\n        self.final_targets = targets\n\n\n# ============================================================================\n# Training and Validation Loops\n# ============================================================================\n\nclass Trainer:\n    \"\"\"\n    Class for managing training and validation loops with history tracking.\n    \"\"\"\n    \n    def __init__(\n        self,\n        model: nn.Module,\n        criterion: nn.Module,\n        device: torch.device,\n        scorer: CompetitionScorer\n    ):\n        \"\"\"\n        Args:\n            model: Model to train\n            criterion: Loss function\n            device: Device (CPU/GPU)\n            scorer: Evaluation metric calculator\n        \"\"\"\n        self.model = model\n        self.criterion = criterion\n        self.device = device\n        self.scorer = scorer\n        self.history = TrainingHistory()\n    \n    def train_one_epoch(\n        self,\n        loader: DataLoader,\n        optimizer: optim.Optimizer\n    ) -> float:\n        \"\"\"\n        Train for one epoch.\n        \n        Args:\n            loader: Training DataLoader\n            optimizer: Optimizer\n            \n        Returns:\n            Average loss for the epoch\n        \"\"\"\n        self.model.train()\n        epoch_loss = 0.0\n        \n        pbar = tqdm(loader, desc=\"Training\", leave=False)\n        \n        for img_left, img_right, train_targets, _ in pbar:\n            # Transfer to device\n            img_left = img_left.to(self.device)\n            img_right = img_right.to(self.device)\n            targets = train_targets.to(self.device)\n            \n            # Zero gradients\n            optimizer.zero_grad()\n            \n            # Forward pass\n            predictions = self.model(img_left, img_right)\n            \n            # Calculate loss\n            loss = self.criterion(predictions, targets)\n            \n            # Backward pass\n            loss.backward()\n            optimizer.step()\n            \n            # Record\n            epoch_loss += loss.item()\n            pbar.set_postfix(loss=f'{loss.item():.4f}')\n        \n        return epoch_loss / len(loader)\n    \n    def validate_one_epoch(\n        self,\n        loader: DataLoader,\n        store_predictions: bool = False\n    ) -> tuple[float, float]:\n        \"\"\"\n        Validate for one epoch.\n        \n        Args:\n            loader: Validation DataLoader\n            store_predictions: Whether to store predictions for visualization\n            \n        Returns:\n            (average loss for epoch, R² score)\n        \"\"\"\n        self.model.eval()\n        epoch_loss = 0.0\n        \n        # Collect predictions and targets\n        all_preds = {'total': [], 'gdm': [], 'green': []}\n        all_targets = []\n        \n        with torch.no_grad():\n            pbar = tqdm(loader, desc=\"Validating\", leave=False)\n            \n            for img_left, img_right, train_targets, all_targets_5 in pbar:\n                img_left = img_left.to(self.device)\n                img_right = img_right.to(self.device)\n                train_targets = train_targets.to(self.device)\n                \n                # Forward pass\n                pred_total, pred_gdm, pred_green = self.model(img_left, img_right)\n                \n                # Calculate loss\n                predictions = (pred_total, pred_gdm, pred_green)\n                loss = self.criterion(predictions, train_targets)\n                epoch_loss += loss.item()\n                \n                # Collect predictions (convert to numpy on CPU)\n                all_preds['total'].append(pred_total.cpu().numpy())\n                all_preds['gdm'].append(pred_gdm.cpu().numpy())\n                all_preds['green'].append(pred_green.cpu().numpy())\n                all_targets.append(all_targets_5.cpu().numpy())\n        \n        # Concatenate batches\n        preds_np = {\n            k: np.concatenate(v).flatten() \n            for k, v in all_preds.items()\n        }\n        targets_np = np.concatenate(all_targets)\n        \n        # Store predictions if requested\n        if store_predictions:\n            self.history.set_final_predictions(preds_np, targets_np)\n        \n        # Calculate R² score\n        score = self.scorer.calculate_score(preds_np, targets_np)\n        \n        avg_loss = epoch_loss / len(loader)\n        return avg_loss, score\n\n\n# ============================================================================\n# Visualization\n# ============================================================================\n\nclass TrainingVisualizer:\n    \"\"\"\n    Class for creating comprehensive training visualizations.\n    \n    Generates and saves:\n    1. Learning curves (loss and R² over epochs)\n    2. Fold comparison (bar chart of scores across folds)\n    3. Prediction scatter plots (predicted vs. actual for all targets)\n    \"\"\"\n    \n    def __init__(self, output_dir: Path, target_names: list[str]):\n        \"\"\"\n        Args:\n            output_dir: Directory to save plots\n            target_names: List of target names for labeling\n        \"\"\"\n        self.output_dir = output_dir\n        self.target_names = target_names\n        \n        # Set style\n        sns.set_style(\"whitegrid\")\n        plt.rcParams['figure.dpi'] = 100\n        plt.rcParams['savefig.dpi'] = 300\n        plt.rcParams['font.size'] = 10\n    \n    def plot_learning_curves(\n        self,\n        history: TrainingHistory,\n        fold: int,\n        freeze_epoch: int\n    ) -> None:\n        \"\"\"\n        Plot learning curves showing loss and R² score over epochs.\n        \n        Args:\n            history: Training history object\n            fold: Fold number\n            freeze_epoch: Epoch where backbone was unfrozen\n        \"\"\"\n        fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n        \n        epochs = range(1, len(history.train_losses) + 1)\n        \n        # Plot 1: Loss curves\n        axes[0].plot(epochs, history.train_losses, label='Train Loss', marker='o', linewidth=2)\n        axes[0].plot(epochs, history.valid_losses, label='Valid Loss', marker='s', linewidth=2)\n        axes[0].axvline(x=freeze_epoch, color='red', linestyle='--', \n                        label=f'Unfreeze (Epoch {freeze_epoch})', alpha=0.7)\n        axes[0].set_xlabel('Epoch')\n        axes[0].set_ylabel('Loss')\n        axes[0].set_title(f'Fold {fold}: Training and Validation Loss')\n        axes[0].legend()\n        axes[0].grid(True, alpha=0.3)\n        \n        # Plot 2: R² score curve\n        axes[1].plot(epochs, history.valid_scores, label='Valid R² Score', \n                     marker='D', color='green', linewidth=2)\n        axes[1].axvline(x=freeze_epoch, color='red', linestyle='--', \n                        label=f'Unfreeze (Epoch {freeze_epoch})', alpha=0.7)\n        axes[1].axhline(y=max(history.valid_scores), color='orange', linestyle=':', \n                        label=f'Best R² = {max(history.valid_scores):.4f}', alpha=0.7)\n        axes[1].set_xlabel('Epoch')\n        axes[1].set_ylabel('R² Score')\n        axes[1].set_title(f'Fold {fold}: Validation R² Score')\n        axes[1].legend()\n        axes[1].grid(True, alpha=0.3)\n        \n        plt.tight_layout()\n        save_path = self.output_dir / f'fold{fold}_learning_curves.png'\n        plt.savefig(save_path, bbox_inches='tight')\n        plt.close()\n        \n        print(f\"  📊 Saved learning curves: {save_path}\")\n    \n    def plot_prediction_scatter(\n        self,\n        history: TrainingHistory,\n        fold: int,\n        scorer: CompetitionScorer\n    ) -> None:\n        \"\"\"\n        Plot scatter plots of predicted vs. actual values for all targets.\n        \n        Args:\n            history: Training history with final predictions\n            fold: Fold number\n            scorer: Scorer to calculate individual R² scores\n        \"\"\"\n        if history.final_preds is None or history.final_targets is None:\n            print(f\"  ⚠️  No predictions stored for fold {fold}\")\n            return\n        \n        # Reconstruct all 5 predictions\n        pred_total = history.final_preds['total']\n        pred_gdm = history.final_preds['gdm']\n        pred_green = history.final_preds['green']\n        pred_clover = np.maximum(0, pred_gdm - pred_green)\n        pred_dead = np.maximum(0, pred_total - pred_gdm)\n        \n        predictions = [pred_green, pred_dead, pred_clover, pred_gdm, pred_total]\n        targets = history.final_targets\n        \n        # Calculate individual R² scores\n        individual_r2 = scorer.calculate_individual_scores(\n            history.final_preds,\n            history.final_targets\n        )\n        \n        # Create subplots\n        fig, axes = plt.subplots(2, 3, figsize=(18, 12))\n        axes = axes.flatten()\n        \n        for idx, (pred, target_name) in enumerate(zip(predictions, self.target_names)):\n            ax = axes[idx]\n            true_vals = targets[:, idx]\n            \n            # Scatter plot\n            ax.scatter(true_vals, pred, alpha=0.5, s=20)\n            \n            # Perfect prediction line\n            min_val = min(true_vals.min(), pred.min())\n            max_val = max(true_vals.max(), pred.max())\n            ax.plot([min_val, max_val], [min_val, max_val], \n                   'r--', linewidth=2, label='Perfect Prediction')\n            \n            # Labels and title\n            ax.set_xlabel('Actual Value (g)')\n            ax.set_ylabel('Predicted Value (g)')\n            r2 = individual_r2[target_name]\n            ax.set_title(f'{target_name}\\nR² = {r2:.4f}')\n            ax.legend()\n            ax.grid(True, alpha=0.3)\n        \n        # Remove extra subplot\n        axes[5].axis('off')\n        \n        plt.suptitle(f'Fold {fold}: Prediction vs. Actual Scatter Plots', \n                     fontsize=14, fontweight='bold', y=0.995)\n        plt.tight_layout()\n        \n        save_path = self.output_dir / f'fold{fold}_scatter_plots.png'\n        plt.savefig(save_path, bbox_inches='tight')\n        plt.close()\n        \n        print(f\"  📊 Saved scatter plots: {save_path}\")\n    \n    def plot_fold_comparison(\n        self,\n        fold_scores: dict[int, float]\n    ) -> None:\n        \"\"\"\n        Plot bar chart comparing R² scores across all folds.\n        \n        Args:\n            fold_scores: Dictionary mapping fold number to best R² score\n        \"\"\"\n        folds = sorted(fold_scores.keys())\n        scores = [fold_scores[f] for f in folds]\n        mean_score = np.mean(scores)\n        std_score = np.std(scores)\n        \n        fig, ax = plt.subplots(figsize=(12, 6))\n        \n        # Bar plot\n        bars = ax.bar(folds, scores, color='skyblue', edgecolor='navy', linewidth=1.5)\n        \n        # Add value labels on bars\n        for fold, score, bar in zip(folds, scores, bars):\n            height = bar.get_height()\n            ax.text(bar.get_x() + bar.get_width() / 2., height,\n                   f'{score:.4f}',\n                   ha='center', va='bottom', fontweight='bold')\n        \n        # Mean line\n        ax.axhline(y=mean_score, color='red', linestyle='--', linewidth=2,\n                  label=f'Mean R² = {mean_score:.4f} ± {std_score:.4f}')\n        \n        # Labels and title\n        ax.set_xlabel('Fold', fontsize=12, fontweight='bold')\n        ax.set_ylabel('Best R² Score', fontsize=12, fontweight='bold')\n        ax.set_title('Cross-Validation Performance: R² Score Comparison Across Folds', \n                    fontsize=14, fontweight='bold')\n        ax.set_xticks(folds)\n        ax.set_xticklabels([f'Fold {f}' for f in folds])\n        ax.legend(fontsize=11)\n        ax.grid(True, alpha=0.3, axis='y')\n        \n        plt.tight_layout()\n        save_path = self.output_dir / 'fold_comparison.png'\n        plt.savefig(save_path, bbox_inches='tight')\n        plt.close()\n        \n        print(f\"\\n📊 Saved fold comparison: {save_path}\")\n        print(f\"📈 Overall CV Score: {mean_score:.4f} ± {std_score:.4f}\")\n\n\n# ============================================================================\n# Two-Stage Training Pipeline with Visualization\n# ============================================================================\n\nclass TwoStageTrainingPipeline:\n    \"\"\"\n    Pipeline executing two-stage training: Freeze→Unfreeze with visualization.\n    \n    Stage 1: Freeze backbone and train heads only\n    Stage 2: Fine-tune entire model\n    \"\"\"\n    \n    def __init__(self, config: Config):\n        \"\"\"\n        Args:\n            config: Configuration object\n        \"\"\"\n        self.config = config\n        self.data_prep = DataPreparator(config)\n        self.aug_factory = AugmentationFactory(config.img_size)\n        self.scorer = CompetitionScorer(config.r2_weights)\n        self.visualizer = TrainingVisualizer(config.output_dir, config.all_target_cols)\n        \n        # Store fold scores for comparison\n        self.fold_scores: dict[int, float] = {}\n    \n    def run_fold(self, fold: int) -> None:\n        \"\"\"\n        Execute training for one fold with visualization.\n        \n        Args:\n            fold: Fold number\n        \"\"\"\n        print(f\"\\n{'='*70}\")\n        print(f\"🚀 Starting Fold {fold} Training (Two-Stage)\")\n        print(f\"{'='*70}\")\n        \n        start_time = time.time()\n        \n        # Prepare data\n        df = self.data_prep.df_wide\n        train_df = df[df['fold'] != fold].reset_index(drop=True)\n        valid_df = df[df['fold'] == fold].reset_index(drop=True)\n        \n        print(f\"Training data: {len(train_df)} images | Validation data: {len(valid_df)} images\")\n        \n        # Create datasets\n        train_dataset = BiomassDataset(\n            train_df,\n            self.aug_factory.get_train_transforms,\n            self.config.image_dir,\n            self.config.train_target_cols,\n            self.config.all_target_cols\n        )\n        \n        valid_dataset = BiomassDataset(\n            valid_df,\n            self.aug_factory.get_valid_transforms,\n            self.config.image_dir,\n            self.config.train_target_cols,\n            self.config.all_target_cols\n        )\n        \n        # Create DataLoaders\n        train_loader = DataLoader(\n            train_dataset,\n            batch_size=self.config.batch_size,\n            shuffle=True,\n            num_workers=self.config.num_workers,\n            pin_memory=True\n        )\n        \n        valid_loader = DataLoader(\n            valid_dataset,\n            batch_size=self.config.batch_size * 2,\n            shuffle=False,\n            num_workers=self.config.num_workers,\n            pin_memory=True\n        )\n        \n        # Initialize model\n        model_base = BiomassModel(\n            self.config.model_name,\n            self.config.pretrained\n        )\n        \n        # Multi-GPU support\n        if torch.cuda.device_count() > 1:\n            print(f\"Using {torch.cuda.device_count()} GPUs (DataParallel)\")\n            model = nn.DataParallel(model_base)\n        else:\n            model = model_base\n        \n        model.to(self.config.device)\n        \n        # Loss function\n        criterion = WeightedBiomassLoss(self.config.loss_weights).to(self.config.device)\n        \n        # Trainer instance\n        trainer = Trainer(model, criterion, self.config.device, self.scorer)\n        \n        # ===== Stage 1: Freeze Backbone =====\n        print(f\"\\n--- Stage 1: Backbone Frozen (Epoch 1-{self.config.freeze_epochs}) ---\")\n        self._freeze_backbone(model)\n        \n        optimizer = optim.Adam(\n            filter(lambda p: p.requires_grad, model.parameters()),\n            lr=self.config.learning_rate\n        )\n        scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, mode='min', factor=0.1, patience=2\n        )\n        \n        best_score = -float('inf')\n        \n        for epoch in range(1, self.config.freeze_epochs + 1):\n            print(f\"\\nEpoch {epoch}/{self.config.epochs} (Stage 1)\")\n            \n            train_loss = trainer.train_one_epoch(train_loader, optimizer)\n            valid_loss, score = trainer.validate_one_epoch(valid_loader)\n            \n            # Record history\n            trainer.history.add_epoch(train_loss, valid_loss, score)\n            \n            scheduler.step(valid_loss)\n            \n            print(f\"Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R²: {score:.4f}\")\n            \n            if score > best_score:\n                best_score = score\n                self._save_model(model, fold)\n                print(\"✨ R² score improved! Saving model\")\n        \n        # ===== Stage 2: Fine-tune Entire Model =====\n        print(f\"\\n--- Stage 2: Full Fine-tuning (Epoch {self.config.freeze_epochs+1}-{self.config.epochs}) ---\")\n        self._unfreeze_backbone(model)\n        \n        optimizer = optim.Adam(\n            model.parameters(),\n            lr=self.config.finetune_lr\n        )\n        scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, mode='min', factor=0.2, patience=3\n        )\n        \n        for epoch in range(self.config.freeze_epochs + 1, self.config.epochs + 1):\n            print(f\"\\nEpoch {epoch}/{self.config.epochs} (Stage 2)\")\n            \n            train_loss = trainer.train_one_epoch(train_loader, optimizer)\n            \n            # Store predictions on last epoch for visualization\n            store_preds = (epoch == self.config.epochs)\n            valid_loss, score = trainer.validate_one_epoch(valid_loader, store_predictions=store_preds)\n            \n            # Record history\n            trainer.history.add_epoch(train_loss, valid_loss, score)\n            \n            scheduler.step(valid_loss)\n            \n            print(f\"Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R²: {score:.4f}\")\n            \n            if score > best_score:\n                best_score = score\n                self._save_model(model, fold)\n                print(\"✨ R² score improved! Saving model\")\n        \n        # Store fold score\n        self.fold_scores[fold] = best_score\n        \n        # Generate visualizations for this fold\n        print(f\"\\n📊 Generating visualizations for Fold {fold}...\")\n        self.visualizer.plot_learning_curves(trainer.history, fold, self.config.freeze_epochs)\n        self.visualizer.plot_prediction_scatter(trainer.history, fold, self.scorer)\n        \n        # Finish\n        elapsed = (time.time() - start_time) / 60\n        print(f\"\\n🎉 Fold {fold} completed ({elapsed:.2f} minutes)\")\n        print(f\"Best R² score: {best_score:.4f}\")\n        \n        # Free memory\n        del model, train_loader, valid_loader\n        gc.collect()\n        torch.cuda.empty_cache()\n    \n    def _freeze_backbone(self, model: nn.Module) -> None:\n        \"\"\"Freeze backbone parameters\"\"\"\n        backbone = model.module.backbone if isinstance(model, nn.DataParallel) else model.backbone\n        for param in backbone.parameters():\n            param.requires_grad = False\n    \n    def _unfreeze_backbone(self, model: nn.Module) -> None:\n        \"\"\"Unfreeze backbone parameters\"\"\"\n        backbone = model.module.backbone if isinstance(model, nn.DataParallel) else model.backbone\n        for param in backbone.parameters():\n            param.requires_grad = True\n    \n    def _save_model(self, model: nn.Module, fold: int) -> None:\n        \"\"\"Save model to output directory\"\"\"\n        state_dict = model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict()\n        save_path = self.config.output_dir / f'best_model_fold{fold}.pth'\n        torch.save(state_dict, save_path)\n    \n    def run_all_folds(self) -> None:\n        \"\"\"Execute training for all folds and generate summary visualization\"\"\"\n        # Prepare data\n        self.data_prep.load_and_pivot()\n        self.data_prep.create_stratified_folds(self.data_prep.df_wide)\n        \n        # Train each fold\n        for fold in range(self.config.n_folds):\n            try:\n                self.run_fold(fold)\n            except Exception as e:\n                print(f\"\\n❌ Error occurred in Fold {fold}: {e}\")\n                gc.collect()\n                torch.cuda.empty_cache()\n                raise\n        \n        # Generate fold comparison plot\n        print(f\"\\n{'='*70}\")\n        print(\"📊 Generating cross-validation summary...\")\n        print(f\"{'='*70}\")\n        self.visualizer.plot_fold_comparison(self.fold_scores)\n\n\n# ============================================================================\n# Main Execution\n# ============================================================================\n\nif __name__ == '__main__':\n    # Initialize configuration\n    config = Config()\n    config.display_info()\n    \n    # Execute pipeline\n    pipeline = TwoStageTrainingPipeline(config)\n    pipeline.run_all_folds()\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"🎊 All Fold Training Completed!\")\n    print(\"=\"*70)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T06:14:21.514962Z","iopub.execute_input":"2025-11-13T06:14:21.515233Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}