Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,1185 @@
+#!/usr/bin/env python3
+"""
+CSIRO Image2Biomass Prediction - Training Pipeline
+====================================================
+Multi-output regression: predicting 5 biomass targets from pasture images.
+
+Targets: Dry_Green_g, Dry_Dead_g, Dry_Clover_g, GDM_g, Dry_Total_g
+Metric: Weighted R² (weights: 0.1, 0.1, 0.1, 0.2, 0.5)
+
+Architecture:
+- Backbone: DINOv2 / ConvNeXt / EfficientNet (via timm)
+- Head: MLP with LayerNorm, GELU, Dropout
+- Loss: SmoothL1 + optional weighted R² + consistency regularizer
+- Training: Mixed precision, gradient checkpointing, differential LR
+
+Usage:
+    python train.py --data_dir /path/to/competition/data --backbone dinov2_base --epochs 50
+"""
+
+import os
+import sys
+import json
+import time
+import random
+import argparse
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
+from torch.cuda.amp import GradScaler, autocast
+import timm
+
+try:
+    import albumentations as A
+    from albumentations.pytorch import ToTensorV2
+    HAS_ALBUMENTATIONS = True
+except ImportError:
+    HAS_ALBUMENTATIONS = False
+    from torchvision import transforms
+
+from PIL import Image
+from sklearn.model_selection import KFold, StratifiedKFold
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# ============================================================
+# Constants
+# ============================================================
+TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']
+TARGET_WEIGHTS = [0.1, 0.1, 0.1, 0.2, 0.5]
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+# Backbone configurations
+BACKBONE_CONFIGS = {
+    'dinov2_small': {
+        'name': 'vit_small_patch14_dinov2.lvd142m',
+        'feat_dim': 384, 'native_size': 518, 'default_size': 224,
+    },
+    'dinov2_base': {
+        'name': 'vit_base_patch14_dinov2.lvd142m',
+        'feat_dim': 768, 'native_size': 518, 'default_size': 224,
+    },
+    'dinov2_large': {
+        'name': 'vit_large_patch14_dinov2.lvd142m',
+        'feat_dim': 1024, 'native_size': 518, 'default_size': 224,
+    },
+    'dinov2_base_reg': {
+        'name': 'vit_base_patch14_reg4_dinov2.lvd142m',
+        'feat_dim': 768, 'native_size': 518, 'default_size': 224,
+    },
+    'convnext_large': {
+        'name': 'convnext_large.fb_in22k_ft_in1k',
+        'feat_dim': 1536, 'native_size': 224, 'default_size': 224,
+    },
+    'convnextv2_large': {
+        'name': 'convnextv2_large.fcmae_ft_in22k_in1k',
+        'feat_dim': 1536, 'native_size': 224, 'default_size': 224,
+    },
+    'efficientnet_b4': {
+        'name': 'efficientnet_b4.ra2_in1k',
+        'feat_dim': 1792, 'native_size': 380, 'default_size': 320,
+    },
+    'swin_large': {
+        'name': 'swin_large_patch4_window7_224.ms_in22k_ft_in1k',
+        'feat_dim': 1536, 'native_size': 224, 'default_size': 224,
+    },
+    'eva02_large': {
+        'name': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',
+        'feat_dim': 1024, 'native_size': 448, 'default_size': 448,
+    },
+}
+
+
+# ============================================================
+# Dataset
+# ============================================================
+class BiomassDataset(Dataset):
+    """Dataset for pasture biomass regression from images."""
+    
+    def __init__(
+        self,
+        image_dir: str,
+        df: pd.DataFrame,
+        targets: Optional[pd.DataFrame] = None,
+        transform=None,
+        img_size: int = 224,
+        use_ndvi: bool = False,
+        log_transform: bool = True,
+        is_test: bool = False,
+    ):
+        self.image_dir = Path(image_dir)
+        self.df = df.reset_index(drop=True)
+        self.targets = targets
+        self.transform = transform
+        self.img_size = img_size
+        self.use_ndvi = use_ndvi
+        self.log_transform = log_transform
+        self.is_test = is_test
+        
+        # Pre-compute image paths
+        self.image_ids = self.df['image_id'].values if 'image_id' in self.df.columns else self.df.index.values
+        
+    def __len__(self):
+        return len(self.df)
+    
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+        img_id = row['image_id'] if 'image_id' in row.index else row.name
+        
+        # Load image
+        img_path = self.image_dir / f"{img_id}.jpg"
+        if not img_path.exists():
+            img_path = self.image_dir / f"{img_id}.png"
+        if not img_path.exists():
+            # Try without extension - search
+            candidates = list(self.image_dir.glob(f"{img_id}.*"))
+            if candidates:
+                img_path = candidates[0]
+            else:
+                raise FileNotFoundError(f"Image not found: {img_id}")
+        
+        img = Image.open(img_path).convert('RGB')
+        img = np.array(img)
+        
+        # Apply transforms
+        if self.transform is not None:
+            if HAS_ALBUMENTATIONS:
+                augmented = self.transform(image=img)
+                img_tensor = augmented['image']
+            else:
+                img = Image.fromarray(img)
+                img_tensor = self.transform(img)
+        else:
+            img = Image.fromarray(img)
+            img_tensor = transforms.ToTensor()(img)
+        
+        result = {'image': img_tensor, 'image_id': str(img_id)}
+        
+        # Add NDVI if available
+        if self.use_ndvi and 'NDVI' in self.df.columns:
+            ndvi = torch.tensor(row['NDVI'], dtype=torch.float32)
+            result['ndvi'] = ndvi
+        
+        # Add targets if training
+        if self.targets is not None:
+            target_values = self.targets.iloc[idx][TARGET_COLS].values.astype(np.float32)
+            if self.log_transform:
+                target_values = np.log1p(target_values)
+            result['targets'] = torch.tensor(target_values, dtype=torch.float32)
+        
+        return result
+
+
+# ============================================================
+# Augmentations
+# ============================================================
+def get_train_transforms(img_size: int = 224, aug_strength: str = 'medium'):
+    """Get training augmentations."""
+    if HAS_ALBUMENTATIONS:
+        if aug_strength == 'light':
+            return A.Compose([
+                A.RandomResizedCrop(size=(img_size, img_size), scale=(0.7, 1.0)),
+                A.HorizontalFlip(p=0.5),
+                A.VerticalFlip(p=0.5),
+                A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+                ToTensorV2(),
+            ])
+        elif aug_strength == 'medium':
+            return A.Compose([
+                A.RandomResizedCrop(size=(img_size, img_size), scale=(0.5, 1.0)),
+                A.HorizontalFlip(p=0.5),
+                A.VerticalFlip(p=0.5),
+                A.RandomRotate90(p=0.5),
+                A.Transpose(p=0.5),
+                A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
+                A.HueSaturationValue(hue_shift_limit=15, sat_shift_limit=25, val_shift_limit=15, p=0.4),
+                A.OneOf([
+                    A.GaussianBlur(blur_limit=(3, 5)),
+                    A.MotionBlur(blur_limit=5),
+                ], p=0.15),
+                A.CoarseDropout(
+                    num_holes_range=(1, 4),
+                    hole_height_range=(int(img_size*0.05), int(img_size*0.15)),
+                    hole_width_range=(int(img_size*0.05), int(img_size*0.15)),
+                    p=0.2,
+                ),
+                A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+                ToTensorV2(),
+            ])
+        else:  # heavy
+            return A.Compose([
+                A.RandomResizedCrop(size=(img_size, img_size), scale=(0.4, 1.0)),
+                A.HorizontalFlip(p=0.5),
+                A.VerticalFlip(p=0.5),
+                A.RandomRotate90(p=0.5),
+                A.Transpose(p=0.5),
+                A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.7),
+                A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5),
+                A.RandomGamma(gamma_limit=(80, 120), p=0.3),
+                A.OneOf([
+                    A.GaussianBlur(blur_limit=(3, 7)),
+                    A.MotionBlur(blur_limit=7),
+                ], p=0.2),
+                A.OneOf([
+                    A.GaussNoise(p=1.0),
+                    A.ISONoise(p=1.0),
+                ], p=0.2),
+                A.CoarseDropout(
+                    num_holes_range=(1, 8),
+                    hole_height_range=(int(img_size*0.05), int(img_size*0.2)),
+                    hole_width_range=(int(img_size*0.05), int(img_size*0.2)),
+                    p=0.3,
+                ),
+                A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+                ToTensorV2(),
+            ])
+    else:
+        return transforms.Compose([
+            transforms.RandomResizedCrop(img_size, scale=(0.5, 1.0)),
+            transforms.RandomHorizontalFlip(0.5),
+            transforms.RandomVerticalFlip(0.5),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+        ])
+
+
+def get_val_transforms(img_size: int = 224):
+    """Get validation transforms."""
+    if HAS_ALBUMENTATIONS:
+        return A.Compose([
+            A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)),
+            A.CenterCrop(height=img_size, width=img_size),
+            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+            ToTensorV2(),
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize(int(img_size * 1.14)),
+            transforms.CenterCrop(img_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+        ])
+
+
+def get_tta_transforms(img_size: int = 224, n_augments: int = 8):
+    """Get TTA (test-time augmentation) transforms. Returns list of transforms."""
+    tta_list = [get_val_transforms(img_size)]  # Original
+    if HAS_ALBUMENTATIONS:
+        # Add flipped/rotated versions
+        tta_list.append(A.Compose([
+            A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)),
+            A.CenterCrop(height=img_size, width=img_size),
+            A.HorizontalFlip(p=1.0),
+            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+            ToTensorV2(),
+        ]))
+        tta_list.append(A.Compose([
+            A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)),
+            A.CenterCrop(height=img_size, width=img_size),
+            A.VerticalFlip(p=1.0),
+            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+            ToTensorV2(),
+        ]))
+        tta_list.append(A.Compose([
+            A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)),
+            A.CenterCrop(height=img_size, width=img_size),
+            A.HorizontalFlip(p=1.0),
+            A.VerticalFlip(p=1.0),
+            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+            ToTensorV2(),
+        ]))
+        # Slightly different crops
+        for scale in [0.9, 1.0, 1.2]:
+            tta_list.append(A.Compose([
+                A.Resize(height=int(img_size * scale * 1.14), width=int(img_size * scale * 1.14)),
+                A.CenterCrop(height=img_size, width=img_size),
+                A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+                ToTensorV2(),
+            ]))
+    return tta_list[:n_augments]
+
+
+# ============================================================
+# Model
+# ============================================================
+class BiomassModel(nn.Module):
+    """
+    Multi-output regression model for biomass prediction.
+    
+    Architecture:
+    - timm backbone (DINOv2, ConvNeXt, etc.)
+    - Optional auxiliary features (NDVI)
+    - MLP regression head with LayerNorm + GELU + Dropout
+    - Optional: separate heads per target for better specialization
+    """
+    
+    def __init__(
+        self,
+        backbone_name: str = 'vit_base_patch14_dinov2.lvd142m',
+        num_targets: int = 5,
+        hidden_dim: int = 512,
+        dropout: float = 0.3,
+        pretrained: bool = True,
+        img_size: int = 224,
+        use_ndvi: bool = False,
+        separate_heads: bool = False,
+        grad_checkpointing: bool = False,
+    ):
+        super().__init__()
+        self.use_ndvi = use_ndvi
+        self.separate_heads = separate_heads
+        self.num_targets = num_targets
+        
+        # Create backbone
+        kwargs = {'pretrained': pretrained, 'num_classes': 0}
+        if 'vit' in backbone_name or 'dinov2' in backbone_name:
+            kwargs['img_size'] = img_size
+        
+        self.backbone = timm.create_model(backbone_name, **kwargs)
+        feat_dim = self.backbone.num_features
+        
+        # Enable gradient checkpointing for memory efficiency
+        if grad_checkpointing:
+            if hasattr(self.backbone, 'set_grad_checkpointing'):
+                self.backbone.set_grad_checkpointing(True)
+                logger.info("Gradient checkpointing enabled")
+        
+        # NDVI embedding
+        if use_ndvi:
+            self.ndvi_embed = nn.Sequential(
+                nn.Linear(1, 32),
+                nn.GELU(),
+                nn.Linear(32, 64),
+            )
+            feat_dim += 64
+        
+        # Regression head(s)
+        if separate_heads:
+            # Separate MLP head per target - better specialization
+            self.heads = nn.ModuleList([
+                nn.Sequential(
+                    nn.LayerNorm(feat_dim),
+                    nn.Dropout(dropout),
+                    nn.Linear(feat_dim, hidden_dim),
+                    nn.GELU(),
+                    nn.Dropout(dropout * 0.5),
+                    nn.Linear(hidden_dim, 1),
+                )
+                for _ in range(num_targets)
+            ])
+        else:
+            # Shared head - better when data is limited
+            self.head = nn.Sequential(
+                nn.LayerNorm(feat_dim),
+                nn.Dropout(dropout),
+                nn.Linear(feat_dim, hidden_dim),
+                nn.GELU(),
+                nn.Dropout(dropout * 0.5),
+                nn.Linear(hidden_dim, hidden_dim // 2),
+                nn.GELU(),
+                nn.Dropout(dropout * 0.3),
+                nn.Linear(hidden_dim // 2, num_targets),
+            )
+    
+    def forward(self, x, ndvi=None):
+        features = self.backbone(x)
+        
+        if self.use_ndvi and ndvi is not None:
+            ndvi_feats = self.ndvi_embed(ndvi.unsqueeze(-1))
+            features = torch.cat([features, ndvi_feats], dim=-1)
+        
+        if self.separate_heads:
+            outputs = [head(features) for head in self.heads]
+            return torch.cat(outputs, dim=-1)
+        else:
+            return self.head(features)
+    
+    def get_param_groups(self, backbone_lr: float = 5e-5, head_lr: float = 1e-3):
+        """Get parameter groups with differential learning rates."""
+        backbone_params = list(self.backbone.parameters())
+        head_params = [p for n, p in self.named_parameters() if 'backbone' not in n]
+        
+        return [
+            {'params': backbone_params, 'lr': backbone_lr},
+            {'params': head_params, 'lr': head_lr},
+        ]
+
+
+# ============================================================
+# Loss Functions
+# ============================================================
+class WeightedSmoothL1Loss(nn.Module):
+    """SmoothL1 loss weighted by target importance."""
+    
+    def __init__(self, target_weights=None, beta=1.0):
+        super().__init__()
+        self.beta = beta
+        if target_weights is None:
+            target_weights = TARGET_WEIGHTS
+        self.register_buffer('weights', torch.tensor(target_weights, dtype=torch.float32))
+    
+    def forward(self, pred, target):
+        loss = F.smooth_l1_loss(pred, target, beta=self.beta, reduction='none')  # [B, 5]
+        weighted = loss * self.weights.unsqueeze(0)
+        return weighted.mean()
+
+
+class WeightedMSELoss(nn.Module):
+    """MSE loss weighted by target importance."""
+    
+    def __init__(self, target_weights=None):
+        super().__init__()
+        if target_weights is None:
+            target_weights = TARGET_WEIGHTS
+        self.register_buffer('weights', torch.tensor(target_weights, dtype=torch.float32))
+    
+    def forward(self, pred, target):
+        loss = (pred - target) ** 2  # [B, 5]
+        weighted = loss * self.weights.unsqueeze(0)
+        return weighted.mean()
+
+
+class ConsistencyLoss(nn.Module):
+    """
+    Enforce structural constraint: Dry_Total_g ≈ Dry_Green_g + Dry_Dead_g + Dry_Clover_g
+    Only approximate because GDM includes all dry matter components.
+    """
+    
+    def __init__(self, weight=0.1):
+        super().__init__()
+        self.weight = weight
+    
+    def forward(self, pred):
+        # pred columns: [Green, Dead, Clover, GDM, Total]
+        component_sum = pred[:, 0] + pred[:, 1] + pred[:, 2]
+        total = pred[:, 4]
+        return self.weight * F.mse_loss(component_sum, total)
+
+
+class CombinedLoss(nn.Module):
+    """Combined loss with SmoothL1 + consistency regularization."""
+    
+    def __init__(self, smoothl1_weight=1.0, mse_weight=0.0, consistency_weight=0.1,
+                 target_weights=None):
+        super().__init__()
+        self.smoothl1 = WeightedSmoothL1Loss(target_weights)
+        self.mse = WeightedMSELoss(target_weights) if mse_weight > 0 else None
+        self.consistency = ConsistencyLoss(consistency_weight) if consistency_weight > 0 else None
+        self.smoothl1_weight = smoothl1_weight
+        self.mse_weight = mse_weight
+    
+    def forward(self, pred, target):
+        loss = self.smoothl1_weight * self.smoothl1(pred, target)
+        if self.mse is not None:
+            loss += self.mse_weight * self.mse(pred, target)
+        if self.consistency is not None:
+            loss += self.consistency(pred)
+        return loss
+
+
+# ============================================================
+# Label Distribution Smoothing (LDS)
+# ============================================================
+def get_lds_weights(labels: np.ndarray, bins: int = 100, kernel_size: int = 5, sigma: float = 2.0):
+    """
+    Compute Label Distribution Smoothing (LDS) weights.
+    From "Delving into Deep Imbalanced Regression" (ICML 2021).
+    """
+    from scipy.ndimage import convolve1d
+    
+    # Use the most important target (Dry_Total_g) for weighting
+    if labels.ndim > 1:
+        labels = labels[:, -1]  # Last column = Dry_Total_g
+    
+    hist, bin_edges = np.histogram(labels, bins=bins)
+    kernel = np.exp(-np.linspace(-3, 3, kernel_size)**2 / (2 * sigma**2))
+    kernel /= kernel.sum()
+    smoothed = convolve1d(hist.astype(float), kernel, mode='reflect')
+    
+    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
+    weights = 1.0 / (np.interp(labels, bin_centers, smoothed) + 1e-8)
+    weights = weights / weights.mean()  # Normalize to mean=1
+    
+    return weights
+
+
+# ============================================================
+# Metrics
+# ============================================================
+def compute_weighted_r2(preds: np.ndarray, targets: np.ndarray, 
+                         target_weights: List[float] = None) -> float:
+    """
+    Compute the globally weighted R² (competition metric).
+    
+    Args:
+        preds: [N, 5] predictions
+        targets: [N, 5] ground truth
+        target_weights: per-target weights (default: competition weights)
+    
+    Returns:
+        Weighted R² score
+    """
+    if target_weights is None:
+        target_weights = TARGET_WEIGHTS
+    
+    n_samples = preds.shape[0]
+    n_targets = preds.shape[1]
+    
+    # Expand to long format with per-row weights
+    all_preds = []
+    all_targets = []
+    all_weights = []
+    
+    for j in range(n_targets):
+        all_preds.extend(preds[:, j])
+        all_targets.extend(targets[:, j])
+        all_weights.extend([target_weights[j]] * n_samples)
+    
+    all_preds = np.array(all_preds)
+    all_targets = np.array(all_targets)
+    all_weights = np.array(all_weights)
+    
+    # Weighted mean
+    weighted_mean = np.sum(all_weights * all_targets) / np.sum(all_weights)
+    
+    # SS_res and SS_tot
+    ss_res = np.sum(all_weights * (all_targets - all_preds) ** 2)
+    ss_tot = np.sum(all_weights * (all_targets - weighted_mean) ** 2)
+    
+    r2 = 1.0 - ss_res / (ss_tot + 1e-8)
+    return r2
+
+
+def compute_per_target_r2(preds: np.ndarray, targets: np.ndarray) -> Dict[str, float]:
+    """Compute R² per target column."""
+    results = {}
+    for i, name in enumerate(TARGET_COLS):
+        ss_res = np.sum((targets[:, i] - preds[:, i]) ** 2)
+        ss_tot = np.sum((targets[:, i] - targets[:, i].mean()) ** 2)
+        r2 = 1.0 - ss_res / (ss_tot + 1e-8)
+        results[name] = r2
+    return results
+
+
+# ============================================================
+# Training Engine
+# ============================================================
+class Trainer:
+    """Training engine with mixed precision, gradient accumulation, and k-fold."""
+    
+    def __init__(self, args):
+        self.args = args
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.scaler = GradScaler() if self.device.type == 'cuda' else None
+        
+        logger.info(f"Device: {self.device}")
+        if self.device.type == 'cuda':
+            logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
+            logger.info(f"GPU Memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+    
+    def build_model(self):
+        """Build model from args."""
+        backbone_cfg = BACKBONE_CONFIGS[self.args.backbone]
+        img_size = self.args.img_size or backbone_cfg['default_size']
+        
+        model = BiomassModel(
+            backbone_name=backbone_cfg['name'],
+            num_targets=5,
+            hidden_dim=self.args.hidden_dim,
+            dropout=self.args.dropout,
+            pretrained=True,
+            img_size=img_size,
+            use_ndvi=self.args.use_ndvi,
+            separate_heads=self.args.separate_heads,
+            grad_checkpointing=self.args.grad_checkpointing,
+        )
+        return model.to(self.device)
+    
+    def build_optimizer(self, model):
+        """Build optimizer with differential learning rates."""
+        param_groups = model.get_param_groups(
+            backbone_lr=self.args.backbone_lr,
+            head_lr=self.args.head_lr,
+        )
+        
+        if self.args.optimizer == 'adamw':
+            optimizer = torch.optim.AdamW(param_groups, weight_decay=self.args.weight_decay)
+        elif self.args.optimizer == 'sgd':
+            optimizer = torch.optim.SGD(param_groups, momentum=0.9, weight_decay=self.args.weight_decay)
+        else:
+            raise ValueError(f"Unknown optimizer: {self.args.optimizer}")
+        
+        return optimizer
+    
+    def build_scheduler(self, optimizer, num_training_steps):
+        """Build learning rate scheduler."""
+        warmup_steps = int(num_training_steps * self.args.warmup_ratio)
+        
+        if self.args.scheduler == 'cosine':
+            from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
+            warmup = LinearLR(optimizer, start_factor=0.01, total_iters=warmup_steps)
+            cosine = CosineAnnealingLR(optimizer, T_max=num_training_steps - warmup_steps, 
+                                        eta_min=self.args.min_lr)
+            scheduler = SequentialLR(optimizer, [warmup, cosine], milestones=[warmup_steps])
+        elif self.args.scheduler == 'plateau':
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer, mode='max', factor=0.5, patience=5, verbose=True)
+        else:
+            scheduler = None
+        
+        return scheduler
+    
+    def train_one_epoch(self, model, loader, optimizer, scheduler, loss_fn, epoch):
+        """Train for one epoch."""
+        model.train()
+        running_loss = 0.0
+        num_samples = 0
+        
+        for batch_idx, batch in enumerate(loader):
+            images = batch['image'].to(self.device)
+            targets = batch['targets'].to(self.device)
+            ndvi = batch.get('ndvi', None)
+            if ndvi is not None:
+                ndvi = ndvi.to(self.device)
+            
+            # Forward pass with mixed precision
+            if self.scaler is not None:
+                with autocast(dtype=torch.float16):
+                    preds = model(images, ndvi)
+                    loss = loss_fn(preds, targets)
+                
+                # Gradient accumulation
+                loss = loss / self.args.grad_accum_steps
+                self.scaler.scale(loss).backward()
+                
+                if (batch_idx + 1) % self.args.grad_accum_steps == 0:
+                    self.scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)
+                    self.scaler.step(optimizer)
+                    self.scaler.update()
+                    optimizer.zero_grad()
+                    
+                    if scheduler is not None and not isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                        scheduler.step()
+            else:
+                preds = model(images, ndvi)
+                loss = loss_fn(preds, targets)
+                loss = loss / self.args.grad_accum_steps
+                loss.backward()
+                
+                if (batch_idx + 1) % self.args.grad_accum_steps == 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    
+                    if scheduler is not None and not isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                        scheduler.step()
+            
+            running_loss += loss.item() * self.args.grad_accum_steps * images.size(0)
+            num_samples += images.size(0)
+            
+            if (batch_idx + 1) % self.args.log_interval == 0:
+                avg_loss = running_loss / num_samples
+                lr = optimizer.param_groups[0]['lr']
+                logger.info(f"Epoch {epoch} [{batch_idx+1}/{len(loader)}] loss={avg_loss:.4f} lr={lr:.2e}")
+        
+        return running_loss / num_samples
+    
+    @torch.no_grad()
+    def validate(self, model, loader, loss_fn, log_transform=True):
+        """Validate and compute metrics."""
+        model.eval()
+        all_preds = []
+        all_targets = []
+        running_loss = 0.0
+        num_samples = 0
+        
+        for batch in loader:
+            images = batch['image'].to(self.device)
+            targets = batch['targets'].to(self.device)
+            ndvi = batch.get('ndvi', None)
+            if ndvi is not None:
+                ndvi = ndvi.to(self.device)
+            
+            if self.scaler is not None:
+                with autocast(dtype=torch.float16):
+                    preds = model(images, ndvi)
+                    loss = loss_fn(preds, targets)
+            else:
+                preds = model(images, ndvi)
+                loss = loss_fn(preds, targets)
+            
+            running_loss += loss.item() * images.size(0)
+            num_samples += images.size(0)
+            
+            all_preds.append(preds.cpu().numpy())
+            all_targets.append(targets.cpu().numpy())
+        
+        all_preds = np.concatenate(all_preds, axis=0)
+        all_targets = np.concatenate(all_targets, axis=0)
+        
+        # Inverse log transform for metric computation
+        if log_transform:
+            all_preds_orig = np.expm1(all_preds)
+            all_targets_orig = np.expm1(all_targets)
+        else:
+            all_preds_orig = all_preds
+            all_targets_orig = all_targets
+        
+        # Clip negative predictions
+        all_preds_orig = np.clip(all_preds_orig, 0, None)
+        
+        # Compute metrics
+        weighted_r2 = compute_weighted_r2(all_preds_orig, all_targets_orig)
+        per_target_r2 = compute_per_target_r2(all_preds_orig, all_targets_orig)
+        
+        avg_loss = running_loss / num_samples
+        
+        return {
+            'loss': avg_loss,
+            'weighted_r2': weighted_r2,
+            'per_target_r2': per_target_r2,
+            'preds': all_preds_orig,
+            'targets': all_targets_orig,
+        }
+    
+    @torch.no_grad()
+    def predict(self, model, loader, log_transform=True, tta_transforms=None):
+        """Generate predictions (inference)."""
+        model.eval()
+        all_preds = []
+        all_ids = []
+        
+        for batch in loader:
+            images = batch['image'].to(self.device)
+            ndvi = batch.get('ndvi', None)
+            if ndvi is not None:
+                ndvi = ndvi.to(self.device)
+            
+            if self.scaler is not None:
+                with autocast(dtype=torch.float16):
+                    preds = model(images, ndvi)
+            else:
+                preds = model(images, ndvi)
+            
+            all_preds.append(preds.cpu().numpy())
+            all_ids.extend(batch['image_id'])
+        
+        all_preds = np.concatenate(all_preds, axis=0)
+        
+        if log_transform:
+            all_preds = np.expm1(all_preds)
+        
+        all_preds = np.clip(all_preds, 0, None)
+        
+        return all_preds, all_ids
+    
+    def train_fold(self, fold: int, train_df: pd.DataFrame, val_df: pd.DataFrame,
+                   train_targets: pd.DataFrame, val_targets: pd.DataFrame,
+                   image_dir: str):
+        """Train a single fold."""
+        backbone_cfg = BACKBONE_CONFIGS[self.args.backbone]
+        img_size = self.args.img_size or backbone_cfg['default_size']
+        
+        # Datasets
+        train_dataset = BiomassDataset(
+            image_dir=image_dir,
+            df=train_df,
+            targets=train_targets,
+            transform=get_train_transforms(img_size, self.args.aug_strength),
+            img_size=img_size,
+            use_ndvi=self.args.use_ndvi,
+            log_transform=self.args.log_transform,
+        )
+        val_dataset = BiomassDataset(
+            image_dir=image_dir,
+            df=val_df,
+            targets=val_targets,
+            transform=get_val_transforms(img_size),
+            img_size=img_size,
+            use_ndvi=self.args.use_ndvi,
+            log_transform=self.args.log_transform,
+        )
+        
+        # Optional: LDS sample weights
+        if self.args.use_lds:
+            sample_weights = get_lds_weights(
+                train_targets[TARGET_COLS].values,
+                bins=self.args.lds_bins,
+                kernel_size=self.args.lds_kernel_size,
+                sigma=self.args.lds_sigma,
+            )
+            sampler = WeightedRandomSampler(
+                weights=sample_weights,
+                num_samples=len(train_dataset),
+                replacement=True,
+            )
+            train_loader = DataLoader(
+                train_dataset, batch_size=self.args.batch_size,
+                sampler=sampler, num_workers=self.args.num_workers,
+                pin_memory=True, drop_last=True,
+            )
+        else:
+            train_loader = DataLoader(
+                train_dataset, batch_size=self.args.batch_size,
+                shuffle=True, num_workers=self.args.num_workers,
+                pin_memory=True, drop_last=True,
+            )
+        
+        val_loader = DataLoader(
+            val_dataset, batch_size=self.args.batch_size * 2,
+            shuffle=False, num_workers=self.args.num_workers,
+            pin_memory=True,
+        )
+        
+        # Model, optimizer, scheduler
+        model = self.build_model()
+        optimizer = self.build_optimizer(model)
+        
+        num_training_steps = len(train_loader) * self.args.epochs // self.args.grad_accum_steps
+        scheduler = self.build_scheduler(optimizer, num_training_steps)
+        
+        # Loss
+        loss_fn = CombinedLoss(
+            smoothl1_weight=1.0,
+            mse_weight=self.args.mse_weight,
+            consistency_weight=self.args.consistency_weight,
+            target_weights=TARGET_WEIGHTS,
+        ).to(self.device)
+        
+        # Training loop
+        best_r2 = -float('inf')
+        best_epoch = 0
+        patience_counter = 0
+        save_dir = Path(self.args.output_dir) / f"fold_{fold}"
+        save_dir.mkdir(parents=True, exist_ok=True)
+        
+        logger.info(f"\n{'='*60}")
+        logger.info(f"FOLD {fold}")
+        logger.info(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}")
+        logger.info(f"Backbone: {backbone_cfg['name']}, img_size: {img_size}")
+        logger.info(f"{'='*60}")
+        
+        for epoch in range(1, self.args.epochs + 1):
+            t0 = time.time()
+            
+            # Train
+            train_loss = self.train_one_epoch(model, train_loader, optimizer, scheduler, loss_fn, epoch)
+            
+            # Validate
+            val_metrics = self.validate(model, val_loader, loss_fn, self.args.log_transform)
+            
+            # LR scheduler step (for ReduceLROnPlateau)
+            if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                scheduler.step(val_metrics['weighted_r2'])
+            
+            elapsed = time.time() - t0
+            
+            # Logging
+            logger.info(
+                f"Epoch {epoch}/{self.args.epochs} | "
+                f"train_loss={train_loss:.4f} | "
+                f"val_loss={val_metrics['loss']:.4f} | "
+                f"val_R²={val_metrics['weighted_r2']:.4f} | "
+                f"time={elapsed:.1f}s"
+            )
+            for name, r2 in val_metrics['per_target_r2'].items():
+                logger.info(f"  {name}: R²={r2:.4f}")
+            
+            # Save best model
+            if val_metrics['weighted_r2'] > best_r2:
+                best_r2 = val_metrics['weighted_r2']
+                best_epoch = epoch
+                patience_counter = 0
+                
+                torch.save({
+                    'epoch': epoch,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'weighted_r2': best_r2,
+                    'per_target_r2': val_metrics['per_target_r2'],
+                    'args': vars(self.args),
+                }, save_dir / 'best_model.pth')
+                
+                logger.info(f"  *** New best R²={best_r2:.4f} (epoch {epoch}) ***")
+            else:
+                patience_counter += 1
+            
+            # Early stopping
+            if patience_counter >= self.args.patience:
+                logger.info(f"Early stopping at epoch {epoch}. Best R²={best_r2:.4f} (epoch {best_epoch})")
+                break
+        
+        # Load best model for final predictions
+        checkpoint = torch.load(save_dir / 'best_model.pth', map_location=self.device, weights_only=False)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        
+        # OOF predictions
+        val_metrics = self.validate(model, val_loader, loss_fn, self.args.log_transform)
+        
+        logger.info(f"\nFold {fold} Final: R²={val_metrics['weighted_r2']:.4f}")
+        
+        return model, val_metrics
+    
+    def train_kfold(self, df: pd.DataFrame, targets: pd.DataFrame, image_dir: str):
+        """Train with K-Fold cross-validation."""
+        n_folds = self.args.n_folds
+        
+        # Stratified bins based on Dry_Total_g
+        bins = pd.qcut(targets['Dry_Total_g'], q=min(10, n_folds), labels=False, duplicates='drop')
+        
+        kf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=self.args.seed)
+        
+        oof_preds = np.zeros((len(df), 5))
+        fold_scores = []
+        
+        for fold, (train_idx, val_idx) in enumerate(kf.split(df, bins)):
+            train_df = df.iloc[train_idx]
+            val_df = df.iloc[val_idx]
+            train_targets = targets.iloc[train_idx]
+            val_targets = targets.iloc[val_idx]
+            
+            model, val_metrics = self.train_fold(
+                fold, train_df, val_df, train_targets, val_targets, image_dir
+            )
+            
+            oof_preds[val_idx] = val_metrics['preds']
+            fold_scores.append(val_metrics['weighted_r2'])
+            
+            logger.info(f"Fold {fold} R²: {val_metrics['weighted_r2']:.4f}")
+        
+        # Overall OOF score
+        targets_arr = targets[TARGET_COLS].values
+        overall_r2 = compute_weighted_r2(oof_preds, targets_arr)
+        
+        logger.info(f"\n{'='*60}")
+        logger.info(f"Overall OOF R²: {overall_r2:.4f}")
+        logger.info(f"Per-fold R²: {[f'{s:.4f}' for s in fold_scores]}")
+        logger.info(f"Mean fold R²: {np.mean(fold_scores):.4f} ± {np.std(fold_scores):.4f}")
+        logger.info(f"{'='*60}")
+        
+        # Save OOF predictions
+        oof_df = df[['image_id']].copy()
+        for i, col in enumerate(TARGET_COLS):
+            oof_df[col] = oof_preds[:, i]
+        oof_df.to_csv(Path(self.args.output_dir) / 'oof_predictions.csv', index=False)
+        
+        return overall_r2, fold_scores
+
+
+# ============================================================
+# Data Loading Utilities
+# ============================================================
+def load_competition_data(data_dir: str):
+    """
+    Load competition data. Expected structure:
+    data_dir/
+        train.csv
+        test.csv
+        train_images/
+        test_images/
+        sample_submission.csv
+    """
+    data_dir = Path(data_dir)
+    
+    # Load CSVs
+    train_df = pd.read_csv(data_dir / 'train.csv')
+    test_df = pd.read_csv(data_dir / 'test.csv')
+    
+    if (data_dir / 'sample_submission.csv').exists():
+        sample_sub = pd.read_csv(data_dir / 'sample_submission.csv')
+    else:
+        sample_sub = None
+    
+    # Determine image directories
+    train_img_dir = data_dir / 'train_images'
+    test_img_dir = data_dir / 'test_images'
+    
+    if not train_img_dir.exists():
+        train_img_dir = data_dir / 'train'
+    if not test_img_dir.exists():
+        test_img_dir = data_dir / 'test'
+    
+    logger.info(f"Train samples: {len(train_df)}")
+    logger.info(f"Test samples: {len(test_df)}")
+    logger.info(f"Train columns: {list(train_df.columns)}")
+    logger.info(f"Test columns: {list(test_df.columns)}")
+    
+    # Check for target columns
+    available_targets = [c for c in TARGET_COLS if c in train_df.columns]
+    logger.info(f"Available targets: {available_targets}")
+    
+    # Print target statistics
+    if available_targets:
+        logger.info("\nTarget statistics:")
+        for col in available_targets:
+            logger.info(f"  {col}: mean={train_df[col].mean():.2f}, "
+                       f"median={train_df[col].median():.2f}, "
+                       f"std={train_df[col].std():.2f}, "
+                       f"min={train_df[col].min():.2f}, "
+                       f"max={train_df[col].max():.2f}")
+    
+    return train_df, test_df, sample_sub, str(train_img_dir), str(test_img_dir)
+
+
+def create_submission(preds: np.ndarray, image_ids: List[str], output_path: str):
+    """
+    Create submission CSV in required format.
+    
+    Args:
+        preds: [N, 5] predictions
+        image_ids: list of image IDs
+        output_path: path to save CSV
+    """
+    rows = []
+    for i, img_id in enumerate(image_ids):
+        for j, target_name in enumerate(TARGET_COLS):
+            rows.append({
+                'sample_id': f"{img_id}__{target_name}",
+                'target': max(0, preds[i, j]),  # Ensure non-negative
+            })
+    
+    sub_df = pd.DataFrame(rows)
+    sub_df.to_csv(output_path, index=False)
+    logger.info(f"Submission saved to {output_path} ({len(sub_df)} rows)")
+    return sub_df
+
+
+# ============================================================
+# Seed and Reproducibility
+# ============================================================
+def set_seed(seed: int):
+    """Set random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+# ============================================================
+# Main
+# ============================================================
+def get_args():
+    parser = argparse.ArgumentParser(description='CSIRO Image2Biomass Training')
+    
+    # Data
+    parser.add_argument('--data_dir', type=str, required=True, help='Competition data directory')
+    parser.add_argument('--output_dir', type=str, default='./output', help='Output directory')
+    
+    # Model
+    parser.add_argument('--backbone', type=str, default='dinov2_base',
+                        choices=list(BACKBONE_CONFIGS.keys()), help='Backbone architecture')
+    parser.add_argument('--img_size', type=int, default=None, help='Image size (default: backbone native)')
+    parser.add_argument('--hidden_dim', type=int, default=512, help='Hidden dim in MLP head')
+    parser.add_argument('--dropout', type=float, default=0.3, help='Dropout rate')
+    parser.add_argument('--separate_heads', action='store_true', help='Use separate heads per target')
+    parser.add_argument('--grad_checkpointing', action='store_true', help='Enable gradient checkpointing')
+    parser.add_argument('--use_ndvi', action='store_true', help='Use NDVI features')
+    
+    # Training
+    parser.add_argument('--epochs', type=int, default=50, help='Max epochs')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+    parser.add_argument('--backbone_lr', type=float, default=5e-5, help='Backbone learning rate')
+    parser.add_argument('--head_lr', type=float, default=1e-3, help='Head learning rate')
+    parser.add_argument('--min_lr', type=float, default=1e-7, help='Min learning rate')
+    parser.add_argument('--weight_decay', type=float, default=1e-2, help='Weight decay')
+    parser.add_argument('--optimizer', type=str, default='adamw', choices=['adamw', 'sgd'])
+    parser.add_argument('--scheduler', type=str, default='cosine', choices=['cosine', 'plateau', 'none'])
+    parser.add_argument('--warmup_ratio', type=float, default=0.05, help='Warmup ratio')
+    parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm')
+    parser.add_argument('--grad_accum_steps', type=int, default=1, help='Gradient accumulation steps')
+    parser.add_argument('--patience', type=int, default=10, help='Early stopping patience')
+    parser.add_argument('--log_interval', type=int, default=10, help='Log every N batches')
+    
+    # Augmentation
+    parser.add_argument('--aug_strength', type=str, default='medium', choices=['light', 'medium', 'heavy'])
+    parser.add_argument('--log_transform', action='store_true', default=True, help='Log-transform targets')
+    parser.add_argument('--no_log_transform', action='store_true', help='Disable log-transform')
+    
+    # LDS
+    parser.add_argument('--use_lds', action='store_true', help='Use Label Distribution Smoothing')
+    parser.add_argument('--lds_bins', type=int, default=100)
+    parser.add_argument('--lds_kernel_size', type=int, default=5)
+    parser.add_argument('--lds_sigma', type=float, default=2.0)
+    
+    # Loss
+    parser.add_argument('--mse_weight', type=float, default=0.0, help='MSE loss weight')
+    parser.add_argument('--consistency_weight', type=float, default=0.1, help='Consistency loss weight')
+    
+    # CV
+    parser.add_argument('--n_folds', type=int, default=5, help='Number of CV folds')
+    parser.add_argument('--fold', type=int, default=None, help='Train single fold (None=all)')
+    
+    # Misc
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--mixed_precision', action='store_true', default=True)
+    
+    args = parser.parse_args()
+    
+    if args.no_log_transform:
+        args.log_transform = False
+    
+    return args
+
+
+def main():
+    args = get_args()
+    set_seed(args.seed)
+    
+    # Load data
+    train_df, test_df, sample_sub, train_img_dir, test_img_dir = load_competition_data(args.data_dir)
+    
+    # Separate features and targets
+    targets = train_df[TARGET_COLS].copy()
+    
+    # Create output directory
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    
+    # Save args
+    with open(Path(args.output_dir) / 'args.json', 'w') as f:
+        json.dump(vars(args), f, indent=2)
+    
+    # Train
+    trainer = Trainer(args)
+    
+    if args.fold is not None:
+        # Single fold training
+        from sklearn.model_selection import StratifiedKFold
+        bins = pd.qcut(targets['Dry_Total_g'], q=min(10, args.n_folds), labels=False, duplicates='drop')
+        kf = StratifiedKFold(n_splits=args.n_folds, shuffle=True, random_state=args.seed)
+        
+        for fold_idx, (train_idx, val_idx) in enumerate(kf.split(train_df, bins)):
+            if fold_idx == args.fold:
+                train_fold_df = train_df.iloc[train_idx]
+                val_fold_df = train_df.iloc[val_idx]
+                train_targets = targets.iloc[train_idx]
+                val_targets = targets.iloc[val_idx]
+                
+                model, val_metrics = trainer.train_fold(
+                    args.fold, train_fold_df, val_fold_df, 
+                    train_targets, val_targets, train_img_dir
+                )
+                break
+    else:
+        # Full K-fold training
+        overall_r2, fold_scores = trainer.train_kfold(train_df, targets, train_img_dir)
+    
+    logger.info("Training complete!")
+
+
+if __name__ == '__main__':
+    main()