{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nCSIRO Pasture Biomass Prediction - Generalized Training Pipeline\n==================================================================\nA production-ready, fully configurable training pipeline for predicting\n5 biomass components from pasture images and metadata.\n\nAuthor: Custom ML Pipeline\nDate: 2025\n\"\"\"\n\nimport os\nimport json\nimport random\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple, Optional, Any\nfrom dataclasses import dataclass, asdict\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.metrics import r2_score\nimport cv2\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\n\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\n\n@dataclass\nclass Config:\n    \"\"\"Centralized configuration for the entire pipeline\"\"\"\n    \n    # ==================== PATHS ====================\n    train_csv: str = '/kaggle/input/csiro-biomass/train.csv'\n    train_image_dir: str = '/kaggle/input/csiro-biomass'\n    output_dir: str = './outputs'\n    \n    # ==================== MODEL ====================\n    model_name: str = 'convnext_tiny'  # Any timm model\n    img_size: int = 768\n    pretrained: bool = True\n    use_two_stream: bool = True  # Split image into left/right patches\n    \n    # ==================== TARGETS ====================\n    target_cols: List[str] = None  # Will be set automatically\n    target_weights: Dict[str, float] = None  # Will be set automatically\n    predict_primary_only: bool = True  # Predict Total, GDM, Green; calculate others\n    primary_targets: List[str] = None  # Will be set if predict_primary_only=True\n    \n    # ==================== FEATURES ====================\n    tabular_features: List[str] = None  # Will be set automatically\n    use_tabular_features: bool = True\n    scale_tabular: bool = True\n    scale_targets: bool = True\n    \n    # ==================== TRAINING ====================\n    n_folds: int = 5\n    fold_strategy: str = 'group'  # 'stratified', 'group', or 'simple'\n    group_col: str = 'Sampling_Date'  # For GroupKFold\n    \n    epochs: int = 25\n    batch_size: int = 8\n    accumulation_steps: int = 1  # Gradient accumulation\n    num_workers: int = 2\n    pin_memory: bool = True\n    \n    # ==================== OPTIMIZER ====================\n    optimizer: str = 'adamw'  # 'adam', 'adamw', 'sgd'\n    lr: float = 1e-4\n    weight_decay: float = 1e-5\n    \n    # ==================== SCHEDULER ====================\n    scheduler: str = 'cosine'  # 'cosine', 'step', 'reduce_on_plateau', None\n    warmup_epochs: int = 2\n    min_lr: float = 1e-6\n    \n    # ==================== LOSS ====================\n    loss_fn: str = 'smooth_l1'  # 'mse', 'mae', 'smooth_l1', 'huber'\n    use_weighted_loss: bool = True  # Weight by competition weights\n    add_physics_loss: bool = True  # Add constraint: Total = Green+Dead+Clover\n    physics_loss_weight: float = 0.1\n    \n    # ==================== AUGMENTATION ====================\n    use_augmentation: bool = True\n    augmentation_prob: float = 0.5\n    \n    # ==================== REGULARIZATION ====================\n    dropout: float = 0.3\n    use_mixup: bool = False\n    mixup_alpha: float = 0.2\n    \n    # ==================== EARLY STOPPING ====================\n    early_stopping: bool = True\n    patience: int = 7\n    \n    # ==================== MISC ====================\n    seed: int = 42\n    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n    fp16: bool = True  # Mixed precision training\n    save_best_only: bool = True\n    verbose: int = 1  # 0: quiet, 1: progress bars, 2: detailed\n    \n    def __post_init__(self):\n        \"\"\"Initialize derived attributes\"\"\"\n        # Set target columns and weights\n        if self.target_cols is None:\n            self.target_cols = [\n                'Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', \n                'GDM_g', 'Dry_Total_g'\n            ]\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        # Set primary targets if using hierarchical prediction\n        if self.predict_primary_only and self.primary_targets is None:\n            self.primary_targets = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n        \n        # Set tabular features\n        if self.tabular_features is None:\n            self.tabular_features = ['Pre_GSHH_NDVI', 'Height_Ave_cm']\n        \n        # Create output directory\n        Path(self.output_dir).mkdir(parents=True, exist_ok=True)\n    \n    def save(self, filepath: str):\n        \"\"\"Save configuration to JSON\"\"\"\n        with open(filepath, 'w') as f:\n            json.dump(asdict(self), f, indent=4)\n    \n    @classmethod\n    def load(cls, filepath: str):\n        \"\"\"Load configuration from JSON\"\"\"\n        with open(filepath, 'r') as f:\n            config_dict = json.load(f)\n        return cls(**config_dict)\n\n\n# ============================================================================\n# UTILITIES\n# ============================================================================\n\ndef set_seed(seed: int):\n    \"\"\"Set random seeds for reproducibility\"\"\"\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n\nclass AverageMeter:\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n    \n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n    \n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\ndef calculate_weighted_r2(y_true: np.ndarray, y_pred: np.ndarray, \n                         target_weights: Dict[str, float],\n                         target_cols: List[str]) -> float:\n    \"\"\"Calculate weighted R² score as per competition metric\"\"\"\n    weighted_r2 = 0.0\n    \n    for idx, target in enumerate(target_cols):\n        y_true_target = y_true[:, idx]\n        y_pred_target = y_pred[:, idx]\n        \n        # Calculate R² for this target\n        r2 = r2_score(y_true_target, y_pred_target)\n        \n        # Weight it\n        weight = target_weights.get(target, 0.0)\n        weighted_r2 += weight * r2\n    \n    return weighted_r2\n\n\n# ============================================================================\n# DATA PREPROCESSING\n# ============================================================================\n\nclass DataPreprocessor:\n    \"\"\"Handle data loading and preprocessing\"\"\"\n    \n    def __init__(self, config: Config):\n        self.config = config\n        self.tabular_scaler = None\n        self.target_scaler = None\n    \n    def load_and_prepare_data(self) -> pd.DataFrame:\n        \"\"\"Load CSV and convert from long to wide format\"\"\"\n        \n        # Load data\n        df_long = pd.read_csv(self.config.train_csv)\n        \n        if self.config.verbose > 0:\n            print(f\"📊 Loaded {len(df_long)} rows from CSV\")\n        \n        # Pivot from long to wide format\n        df_wide = df_long.pivot_table(\n            index='image_path',\n            columns='target_name',\n            values='target',\n            aggfunc='mean'\n        )\n        \n        # Reindex to ensure all targets are present\n        df_wide = df_wide.reindex(columns=self.config.target_cols, fill_value=0.0)\n        \n        # Add metadata\n        meta_cols = ['image_path'] + self.config.tabular_features + ['Sampling_Date']\n        meta_cols = [c for c in meta_cols if c in df_long.columns]\n        \n        meta = df_long.drop_duplicates(subset=['image_path'])[meta_cols]\n        meta = meta.set_index('image_path')\n        \n        df_wide = df_wide.join(meta, how='left')\n        df_wide = df_wide.reset_index()\n        \n        # Handle missing values in tabular features\n        for feat in self.config.tabular_features:\n            if feat in df_wide.columns:\n                df_wide[feat] = df_wide[feat].fillna(df_wide[feat].median())\n        \n        # Create stratification column for folding\n        df_wide = self._create_stratification_column(df_wide)\n        \n        if self.config.verbose > 0:\n            print(f\"✅ Prepared {len(df_wide)} unique images\")\n            print(f\"📋 Target statistics:\")\n            for target in self.config.target_cols:\n                if target in df_wide.columns:\n                    print(f\"   {target}: mean={df_wide[target].mean():.2f}, \"\n                          f\"std={df_wide[target].std():.2f}\")\n        \n        return df_wide\n    \n    def _create_stratification_column(self, df: pd.DataFrame) -> pd.DataFrame:\n        \"\"\"Create stratification bins for balanced folding\"\"\"\n        \n        # Use total biomass for stratification\n        if 'Dry_Total_g' in df.columns:\n            try:\n                df['strat_bin'] = pd.qcut(\n                    df['Dry_Total_g'], \n                    q=10, \n                    labels=False, \n                    duplicates='drop'\n                )\n            except ValueError:\n                df['strat_bin'] = pd.cut(\n                    df['Dry_Total_g'], \n                    bins=10, \n                    labels=False\n                )\n            \n            df['strat_bin'] = df['strat_bin'].fillna(0).astype(int)\n        else:\n            df['strat_bin'] = 0\n        \n        return df\n    \n    def fit_scalers(self, train_df: pd.DataFrame):\n        \"\"\"Fit scalers on training data\"\"\"\n        \n        # Tabular scaler\n        if self.config.scale_tabular and self.config.use_tabular_features:\n            self.tabular_scaler = StandardScaler()\n            tabular_data = train_df[self.config.tabular_features].values\n            self.tabular_scaler.fit(tabular_data)\n        \n        # Target scaler\n        if self.config.scale_targets:\n            self.target_scaler = StandardScaler()\n            \n            # Determine which targets to fit scaler on\n            if self.config.predict_primary_only:\n                target_cols = self.config.primary_targets\n            else:\n                target_cols = self.config.target_cols\n            \n            target_data = train_df[target_cols].values\n            self.target_scaler.fit(target_data)\n    \n    def transform_targets(self, targets: np.ndarray) -> np.ndarray:\n        \"\"\"Transform targets using fitted scaler\"\"\"\n        if self.config.scale_targets and self.target_scaler is not None:\n            return self.target_scaler.transform(targets)\n        return targets\n    \n    def inverse_transform_targets(self, targets: np.ndarray) -> np.ndarray:\n        \"\"\"Inverse transform targets\"\"\"\n        if self.config.scale_targets and self.target_scaler is not None:\n            return self.target_scaler.inverse_transform(targets)\n        return targets\n\n\n# ============================================================================\n# AUGMENTATIONS\n# ============================================================================\n\ndef get_train_transforms(config: Config) -> A.Compose:\n    \"\"\"Get training augmentations\"\"\"\n    \n    transforms = [A.Resize(config.img_size, config.img_size)]\n    \n    if config.use_augmentation:\n        transforms.extend([\n            A.HorizontalFlip(p=config.augmentation_prob),\n            A.VerticalFlip(p=config.augmentation_prob),\n            A.RandomRotate90(p=config.augmentation_prob),\n            A.ShiftScaleRotate(\n                shift_limit=0.1, \n                scale_limit=0.1, \n                rotate_limit=15, \n                p=config.augmentation_prob\n            ),\n            A.OneOf([\n                A.GaussNoise(p=1.0),\n                A.GaussianBlur(p=1.0),\n            ], p=0.2),\n            A.ColorJitter(\n                brightness=0.2, \n                contrast=0.2, \n                saturation=0.2, \n                hue=0.1, \n                p=0.3\n            ),\n        ])\n    \n    transforms.extend([\n        A.Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        ),\n        ToTensorV2()\n    ])\n    \n    return A.Compose(transforms)\n\n\ndef get_valid_transforms(config: Config) -> A.Compose:\n    \"\"\"Get validation augmentations\"\"\"\n    return A.Compose([\n        A.Resize(config.img_size, config.img_size),\n        A.Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        ),\n        ToTensorV2()\n    ])\n\n\n# ============================================================================\n# DATASET\n# ============================================================================\n\nclass BiomassDataset(Dataset):\n    \"\"\"Dataset for biomass prediction\"\"\"\n    \n    def __init__(self, \n                 df: pd.DataFrame,\n                 image_dir: str,\n                 config: Config,\n                 preprocessor: DataPreprocessor,\n                 transform: Optional[A.Compose] = None,\n                 is_training: bool = True):\n        \n        self.df = df.reset_index(drop=True)\n        self.image_dir = image_dir\n        self.config = config\n        self.preprocessor = preprocessor\n        self.transform = transform\n        self.is_training = is_training\n        \n        # Determine target columns to return\n        if config.predict_primary_only:\n            self.target_cols = config.primary_targets\n        else:\n            self.target_cols = config.target_cols\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        \n        # Load image\n        img_path = os.path.join(self.image_dir, row['image_path'])\n        image = cv2.imread(img_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Handle two-stream processing\n        if self.config.use_two_stream:\n            h, w = image.shape[:2]\n            mid = w // 2\n            \n            # Split into left and right\n            img_left = image[:, :mid, :]\n            img_right = image[:, mid:, :]\n            \n            # Apply transforms independently\n            if self.transform:\n                img_left = self.transform(image=img_left)['image']\n                img_right = self.transform(image=img_right)['image']\n            \n            # Stack along channel dimension\n            image = torch.stack([img_left, img_right], dim=0)  # Shape: (2, C, H, W)\n        else:\n            # Single stream\n            if self.transform:\n                image = self.transform(image=image)['image']\n        \n        # Get tabular features\n        if self.config.use_tabular_features:\n            tabular = row[self.config.tabular_features].values.astype(np.float32)\n            \n            # Scale if scaler exists\n            if self.preprocessor.tabular_scaler is not None:\n                tabular = self.preprocessor.tabular_scaler.transform(\n                    tabular.reshape(1, -1)\n                ).flatten()\n            \n            tabular = torch.tensor(tabular, dtype=torch.float32)\n        else:\n            tabular = torch.zeros(0, dtype=torch.float32)\n        \n        # Get targets\n        targets = row[self.target_cols].values.astype(np.float32)\n        \n        # Scale targets if scaler exists\n        if self.preprocessor.target_scaler is not None:\n            targets = self.preprocessor.target_scaler.transform(\n                targets.reshape(1, -1)\n            ).flatten()\n        \n        targets = torch.tensor(targets, dtype=torch.float32)\n        \n        return {\n            'image': image,\n            'tabular': tabular,\n            'target': targets\n        }\n\n\n# ============================================================================\n# MODEL\n# ============================================================================\n\nclass BiomassModel(nn.Module):\n    \"\"\"Multi-modal model for biomass prediction\"\"\"\n    \n    def __init__(self, config: Config):\n        super().__init__()\n        self.config = config\n        \n        # Image encoder (backbone)\n        self.backbone = timm.create_model(\n            config.model_name,\n            pretrained=config.pretrained,\n            num_classes=0,  # Remove classification head\n            global_pool='avg'\n        )\n        \n        # Get feature dimension\n        with torch.no_grad():\n            dummy_input = torch.randn(1, 3, config.img_size, config.img_size)\n            backbone_features = self.backbone(dummy_input)\n            self.backbone_dim = backbone_features.shape[1]\n        \n        # Calculate total feature dimension\n        if config.use_two_stream:\n            total_features = self.backbone_dim * 2  # Two streams concatenated\n        else:\n            total_features = self.backbone_dim\n        \n        # Add tabular features\n        if config.use_tabular_features:\n            self.tabular_encoder = nn.Sequential(\n                nn.Linear(len(config.tabular_features), 64),\n                nn.BatchNorm1d(64),\n                nn.ReLU(),\n                nn.Dropout(config.dropout),\n                nn.Linear(64, 128),\n                nn.BatchNorm1d(128),\n                nn.ReLU(),\n                nn.Dropout(config.dropout)\n            )\n            total_features += 128\n        \n        # Fusion layer\n        self.fusion = nn.Sequential(\n            nn.Linear(total_features, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(config.dropout),\n            nn.Linear(512, 256),\n            nn.BatchNorm1d(256),\n            nn.ReLU(),\n            nn.Dropout(config.dropout / 2)\n        )\n        \n        # Output heads\n        num_outputs = len(config.primary_targets) if config.predict_primary_only else len(config.target_cols)\n        \n        self.heads = nn.ModuleList([\n            nn.Linear(256, 1) for _ in range(num_outputs)\n        ])\n    \n    def forward(self, image, tabular=None):\n        # Process image\n        if self.config.use_two_stream:\n            # image shape: (B, 2, C, H, W)\n            batch_size = image.shape[0]\n            \n            # Process left and right separately\n            img_left = image[:, 0, :, :, :]  # (B, C, H, W)\n            img_right = image[:, 1, :, :, :]  # (B, C, H, W)\n            \n            feat_left = self.backbone(img_left)\n            feat_right = self.backbone(img_right)\n            \n            # Concatenate\n            image_features = torch.cat([feat_left, feat_right], dim=1)\n        else:\n            image_features = self.backbone(image)\n        \n        # Combine with tabular features\n        if self.config.use_tabular_features and tabular is not None and tabular.numel() > 0:\n            tabular_features = self.tabular_encoder(tabular)\n            combined = torch.cat([image_features, tabular_features], dim=1)\n        else:\n            combined = image_features\n        \n        # Fusion\n        fused = self.fusion(combined)\n        \n        # Multi-head output\n        outputs = [head(fused) for head in self.heads]\n        outputs = torch.cat(outputs, dim=1)  # (B, num_targets)\n        \n        return outputs\n\n\n# ============================================================================\n# LOSS FUNCTIONS\n# ============================================================================\n\ndef get_loss_function(config: Config):\n    \"\"\"Get loss function based on config\"\"\"\n    \n    if config.loss_fn == 'mse':\n        base_loss = nn.MSELoss(reduction='none')\n    elif config.loss_fn == 'mae':\n        base_loss = nn.L1Loss(reduction='none')\n    elif config.loss_fn == 'smooth_l1':\n        base_loss = nn.SmoothL1Loss(reduction='none')\n    elif config.loss_fn == 'huber':\n        base_loss = nn.HuberLoss(reduction='none', delta=1.0)\n    else:\n        raise ValueError(f\"Unknown loss function: {config.loss_fn}\")\n    \n    return base_loss\n\n\ndef calculate_loss(pred, target, config: Config, base_loss_fn):\n    \"\"\"Calculate weighted loss with optional physics constraints\"\"\"\n    \n    # Base loss per target\n    losses = base_loss_fn(pred, target)  # (B, num_targets)\n    \n    # Weight losses\n    if config.use_weighted_loss:\n        # Get weights for predicted targets\n        if config.predict_primary_only:\n            target_list = config.primary_targets\n        else:\n            target_list = config.target_cols\n        \n        weights = torch.tensor(\n            [config.target_weights[t] for t in target_list],\n            device=pred.device\n        ).unsqueeze(0)  # (1, num_targets)\n        \n        weighted_losses = losses * weights\n        loss = weighted_losses.sum(dim=1).mean()\n    else:\n        loss = losses.mean()\n    \n    # Add physics constraint loss if applicable\n    if config.add_physics_loss and not config.predict_primary_only:\n        # Total should equal sum of components\n        # Assuming order: Green, Dead, Clover, GDM, Total\n        pred_green = pred[:, 0]\n        pred_dead = pred[:, 1]\n        pred_clover = pred[:, 2]\n        pred_total = pred[:, 4]\n        \n        # Constraint: Total = Green + Dead + Clover\n        components_sum = pred_green + pred_dead + pred_clover\n        physics_loss = F.mse_loss(pred_total, components_sum)\n        \n        loss = loss + config.physics_loss_weight * physics_loss\n    \n    return loss\n\n\n# ============================================================================\n# TRAINING\n# ============================================================================\n\nclass Trainer:\n    \"\"\"Handle model training and validation\"\"\"\n    \n    def __init__(self, config: Config):\n        self.config = config\n        self.device = torch.device(config.device)\n        \n        # For mixed precision training\n        if config.fp16:\n            self.scaler = torch.cuda.amp.GradScaler()\n        else:\n            self.scaler = None\n    \n    def train_one_epoch(self, model, dataloader, optimizer, scheduler, epoch):\n        \"\"\"Train for one epoch\"\"\"\n        \n        model.train()\n        loss_meter = AverageMeter()\n        \n        if self.config.verbose > 0:\n            pbar = tqdm(dataloader, desc=f\"Epoch {epoch+1}/{self.config.epochs} [Train]\")\n        else:\n            pbar = dataloader\n        \n        base_loss_fn = get_loss_function(self.config)\n        optimizer.zero_grad()\n        \n        for batch_idx, batch in enumerate(pbar):\n            image = batch['image'].to(self.device)\n            tabular = batch['tabular'].to(self.device)\n            target = batch['target'].to(self.device)\n            \n            # Forward pass\n            if self.config.fp16:\n                with torch.cuda.amp.autocast():\n                    pred = model(image, tabular)\n                    loss = calculate_loss(pred, target, self.config, base_loss_fn)\n                    loss = loss / self.config.accumulation_steps\n                \n                # Backward pass\n                self.scaler.scale(loss).backward()\n                \n                # Update weights\n                if (batch_idx + 1) % self.config.accumulation_steps == 0:\n                    self.scaler.step(optimizer)\n                    self.scaler.update()\n                    optimizer.zero_grad()\n            else:\n                pred = model(image, tabular)\n                loss = calculate_loss(pred, target, self.config, base_loss_fn)\n                loss = loss / self.config.accumulation_steps\n                \n                loss.backward()\n                \n                if (batch_idx + 1) % self.config.accumulation_steps == 0:\n                    optimizer.step()\n                    optimizer.zero_grad()\n            \n            loss_meter.update(loss.item() * self.config.accumulation_steps, image.size(0))\n            \n            if self.config.verbose > 0:\n                pbar.set_postfix({'loss': loss_meter.avg})\n        \n        # Step scheduler\n        if scheduler is not None and self.config.scheduler != 'reduce_on_plateau':\n            scheduler.step()\n        \n        return loss_meter.avg\n    \n    def validate(self, model, dataloader, preprocessor):\n        \"\"\"Validate model\"\"\"\n        \n        model.eval()\n        loss_meter = AverageMeter()\n        \n        all_preds = []\n        all_targets = []\n        \n        if self.config.verbose > 0:\n            pbar = tqdm(dataloader, desc=\"Validating\")\n        else:\n            pbar = dataloader\n        \n        base_loss_fn = get_loss_function(self.config)\n        \n        with torch.no_grad():\n            for batch in pbar:\n                image = batch['image'].to(self.device)\n                tabular = batch['tabular'].to(self.device)\n                target = batch['target'].to(self.device)\n                \n                # Forward pass\n                pred = model(image, tabular)\n                loss = calculate_loss(pred, target, self.config, base_loss_fn)\n                \n                loss_meter.update(loss.item(), image.size(0))\n                \n                # Store predictions and targets\n                all_preds.append(pred.cpu().numpy())\n                all_targets.append(target.cpu().numpy())\n                \n                if self.config.verbose > 0:\n                    pbar.set_postfix({'loss': loss_meter.avg})\n        \n        # Concatenate all predictions\n        all_preds = np.concatenate(all_preds, axis=0)\n        all_targets = np.concatenate(all_targets, axis=0)\n        \n        # Inverse transform if scaling was used\n        all_preds = preprocessor.inverse_transform_targets(all_preds)\n        all_targets = preprocessor.inverse_transform_targets(all_targets)\n        \n        # Calculate metrics\n        metrics = self._calculate_metrics(all_targets, all_preds)\n        \n        return loss_meter.avg, metrics, all_preds, all_targets\n    \n    def _calculate_metrics(self, y_true, y_pred):\n        \"\"\"Calculate evaluation metrics\"\"\"\n        \n        metrics = {}\n        \n        # Determine target columns\n        if self.config.predict_primary_only:\n            target_list = self.config.primary_targets\n        else:\n            target_list = self.config.target_cols\n        \n        # Per-target R²\n        for idx, target in enumerate(target_list):\n            r2 = r2_score(y_true[:, idx], y_pred[:, idx])\n            metrics[f'r2_{target}'] = r2\n        \n        # Weighted R² (competition metric)\n        weighted_r2 = calculate_weighted_r2(\n            y_true, y_pred, \n            self.config.target_weights, \n            target_list\n        )\n        metrics['weighted_r2'] = weighted_r2\n        \n        return metrics\n\n\n# ============================================================================\n# CROSS-VALIDATION\n# ============================================================================\n\ndef create_folds(df: pd.DataFrame, config: Config):\n    \"\"\"Create cross-validation folds\"\"\"\n    \n    if config.fold_strategy == 'stratified':\n        folder = StratifiedKFold(\n            n_splits=config.n_folds, \n            shuffle=True, \n            random_state=config.seed\n        )\n        folds = list(folder.split(df, df['strat_bin']))\n    \n    elif config.fold_strategy == 'group':\n        # Group by date to avoid leakage\n        folder = GroupKFold(n_splits=config.n_folds)\n        groups = df[config.group_col] if config.group_col in df.columns else df.index\n        folds = list(folder.split(df, groups=groups))\n    \n    else:  # simple\n        from sklearn.model_selection import KFold\n        folder = KFold(\n            n_splits=config.n_folds, \n            shuffle=True, \n            random_state=config.seed\n        )\n        folds = list(folder.split(df))\n    \n    return folds\n\n\ndef train_fold(fold: int, \n               train_idx: np.ndarray, \n               val_idx: np.ndarray,\n               df: pd.DataFrame,\n               config: Config):\n    \"\"\"Train a single fold\"\"\"\n    \n    print(f\"\\n{'='*70}\")\n    print(f\"FOLD {fold + 1}/{config.n_folds}\")\n    print(f\"{'='*70}\")\n    \n    # Split data\n    train_df = df.iloc[train_idx].reset_index(drop=True)\n    val_df = df.iloc[val_idx].reset_index(drop=True)\n    \n    print(f\"Train samples: {len(train_df)}\")\n    print(f\"Valid samples: {len(val_df)}\")\n    \n    # Initialize preprocessor and fit scalers\n    preprocessor = DataPreprocessor(config)\n    preprocessor.fit_scalers(train_df)\n    \n    # Create datasets\n    train_dataset = BiomassDataset(\n        train_df, config.train_image_dir, config, preprocessor,\n        transform=get_train_transforms(config), is_training=True\n    )\n    \n    val_dataset = BiomassDataset(\n        val_df, config.train_image_dir, config, preprocessor,\n        transform=get_valid_transforms(config), is_training=False\n    )\n    \n    # Create dataloaders\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config.batch_size,\n        shuffle=True,\n        num_workers=config.num_workers,\n        pin_memory=config.pin_memory,\n        drop_last=True\n    )\n    \n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=config.batch_size,\n        shuffle=False,\n        num_workers=config.num_workers,\n        pin_memory=config.pin_memory\n    )\n    \n    # Create model\n    model = BiomassModel(config).to(config.device)\n    \n    # Create optimizer\n    if config.optimizer == 'adam':\n        optimizer = torch.optim.Adam(\n            model.parameters(), \n            lr=config.lr, \n            weight_decay=config.weight_decay\n        )\n    elif config.optimizer == 'adamw':\n        optimizer = torch.optim.AdamW(\n            model.parameters(), \n            lr=config.lr, \n            weight_decay=config.weight_decay\n        )\n    else:  # sgd\n        optimizer = torch.optim.SGD(\n            model.parameters(), \n            lr=config.lr, \n            weight_decay=config.weight_decay,\n            momentum=0.9\n        )\n    \n    # Create scheduler\n    if config.scheduler == 'cosine':\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            optimizer, \n            T_max=config.epochs, \n            eta_min=config.min_lr\n        )\n    elif config.scheduler == 'step':\n        scheduler = torch.optim.lr_scheduler.StepLR(\n            optimizer, \n            step_size=config.epochs // 3, \n            gamma=0.1\n        )\n    elif config.scheduler == 'reduce_on_plateau':\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, \n            mode='max',  # Maximize weighted R²\n            factor=0.5, \n            patience=3\n        )\n    else:\n        scheduler = None\n    \n    # Training loop\n    trainer = Trainer(config)\n    best_metric = -np.inf\n    patience_counter = 0\n    \n    history = {\n        'train_loss': [],\n        'val_loss': [],\n        'val_metrics': []\n    }\n    \n    for epoch in range(config.epochs):\n        # Train\n        train_loss = trainer.train_one_epoch(\n            model, train_loader, optimizer, scheduler, epoch\n        )\n        \n        # Validate\n        val_loss, val_metrics, val_preds, val_targets = trainer.validate(\n            model, val_loader, preprocessor\n        )\n        \n        # Step reduce_on_plateau scheduler\n        if config.scheduler == 'reduce_on_plateau':\n            scheduler.step(val_metrics['weighted_r2'])\n        \n        # Log\n        history['train_loss'].append(train_loss)\n        history['val_loss'].append(val_loss)\n        history['val_metrics'].append(val_metrics)\n        \n        print(f\"\\nEpoch {epoch+1}/{config.epochs}\")\n        print(f\"  Train Loss: {train_loss:.4f}\")\n        print(f\"  Val Loss: {val_loss:.4f}\")\n        print(f\"  Val Weighted R²: {val_metrics['weighted_r2']:.4f}\")\n        \n        if config.verbose > 1:\n            for key, value in val_metrics.items():\n                if key != 'weighted_r2':\n                    print(f\"  {key}: {value:.4f}\")\n        \n        # Save best model\n        current_metric = val_metrics['weighted_r2']\n        \n        if current_metric > best_metric:\n            best_metric = current_metric\n            patience_counter = 0\n            \n            # Save checkpoint\n            checkpoint = {\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'best_metric': best_metric,\n                'config': asdict(config),\n                'tabular_scaler': preprocessor.tabular_scaler,\n                'target_scaler': preprocessor.target_scaler\n            }\n            \n            checkpoint_path = os.path.join(\n                config.output_dir, \n                f'best_model_fold{fold}.pth'\n            )\n            torch.save(checkpoint, checkpoint_path)\n            \n            print(f\"  💾 Saved best model (R²: {best_metric:.4f})\")\n        else:\n            patience_counter += 1\n        \n        # Early stopping\n        if config.early_stopping and patience_counter >= config.patience:\n            print(f\"\\n⚠️ Early stopping triggered after {epoch+1} epochs\")\n            break\n    \n    print(f\"\\n✅ Fold {fold+1} complete. Best R²: {best_metric:.4f}\")\n    \n    return best_metric, history\n\n\n# ============================================================================\n# MAIN PIPELINE\n# ============================================================================\n\ndef main():\n    \"\"\"Main training pipeline\"\"\"\n    \n    # Initialize config\n    config = Config()\n    set_seed(config.seed)\n    \n    print(\"=\"*70)\n    print(\"CSIRO BIOMASS PREDICTION - TRAINING PIPELINE\")\n    print(\"=\"*70)\n    print(f\"\\n📋 Configuration:\")\n    print(f\"  Model: {config.model_name}\")\n    print(f\"  Image Size: {config.img_size}\")\n    print(f\"  Two-Stream: {config.use_two_stream}\")\n    print(f\"  Predict Primary Only: {config.predict_primary_only}\")\n    print(f\"  Folds: {config.n_folds}\")\n    print(f\"  Epochs: {config.epochs}\")\n    print(f\"  Batch Size: {config.batch_size}\")\n    print(f\"  Device: {config.device}\")\n    \n    # Save config\n    config.save(os.path.join(config.output_dir, 'config.json'))\n    \n    # Load and prepare data\n    print(f\"\\n📊 Loading data...\")\n    preprocessor = DataPreprocessor(config)\n    df = preprocessor.load_and_prepare_data()\n    \n    # Create folds\n    print(f\"\\n🔀 Creating {config.n_folds} folds ({config.fold_strategy})...\")\n    folds = create_folds(df, config)\n    \n    # Train each fold\n    fold_scores = []\n    \n    for fold, (train_idx, val_idx) in enumerate(folds):\n        fold_score, history = train_fold(fold, train_idx, val_idx, df, config)\n        fold_scores.append(fold_score)\n    \n    # Summary\n    print(f\"\\n{'='*70}\")\n    print(\"TRAINING COMPLETE\")\n    print(f\"{'='*70}\")\n    print(f\"\\nFold Scores (Weighted R²):\")\n    for fold, score in enumerate(fold_scores):\n        print(f\"  Fold {fold+1}: {score:.4f}\")\n    print(f\"\\nMean: {np.mean(fold_scores):.4f} ± {np.std(fold_scores):.4f}\")\n    print(f\"\\n💾 Models saved to: {config.output_dir}\")\n\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}