{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":292960064,"sourceType":"kernelVersion"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\n\nimport os\nimport gc\nimport json\nimport random\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoModel\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.metrics import r2_score\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.ensemble import RandomForestRegressor\nfrom sklearn.decomposition import PCA\nfrom xgboost import XGBRegressor\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport pickle\n\n# HuggingFace login for DINOv3\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nHF_TOKEN = user_secrets.get_secret(\"HF_TOKEN\")\nfrom huggingface_hub import login\nlogin(token=HF_TOKEN)\n\nsns.set_style('whitegrid')\nplt.rcParams['figure.figsize'] = (14, 10)\nprint(\"+\")\n\n# ====================================================\n# CONFIG\n# ====================================================\nclass Config:\n    TRAIN_CSV = '/kaggle/input/image-segmentation/train.csv'\n    IMG_DIR = '/kaggle/input/image-segmentation'\n    \n    # Model specifications\n    SWIN_MODEL = 'swin_base_patch4_window12_384'\n    SWIN_IMG_SIZE = 384\n    \n    DINOV3_MODEL = 'facebook/dinov3-vitb16-pretrain-lvd1689m'\n    DINOV3_IMG_SIZE = 224\n    \n    SIGLIP_MODEL = 'vit_so400m_patch14_siglip_384'\n    SIGLIP_IMG_SIZE = 384\n    \n    # Training\n    SEED = 42\n    N_FOLDS = 5\n    EPOCHS_STAGE1 = 40\n    \n    BATCH_SIZE = 4\n    ACCUMULATION_STEPS = 4\n    \n    # Discriminative Learning Rates\n    BACKBONE_LR = 1e-5  # Lower LR for pretrained backbone\n    HEAD_LR = 3e-5      # Higher LR for regression head\n    WEIGHT_DECAY = 1e-4\n    \n    # OneCycleLR params\n    MAX_LR_BACKBONE = 1e-5\n    MAX_LR_HEAD = 3e-5\n    PCT_START = 0.3  # 30% warmup\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # PCA settings\n    N_PCA_COMPONENTS = 70\n    \n    # Ensemble strategy\n    ENSEMBLE_METHOD = 'learnable'\n    \n    # MixUp/CutMix\n    MIXUP_ALPHA = 0.2\n    CUTMIX_ALPHA = 1.0\n    MIXUP_PROB = 0.5  # Probability of applying mixup/cutmix\n    \n    # These will be set per fold\n    TARGET_MEAN = None\n    TARGET_STD = None\n\ndef seed_everything(seed=42):\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\nseed_everything(Config.SEED)\n\ndef clear_gpu_memory():\n    \"\"\"Aggressive memory cleanup\"\"\"\n    gc.collect()\n    torch.cuda.empty_cache()\n    if torch.cuda.is_available():\n        torch.cuda.synchronize()\n        torch.cuda.ipc_collect()\n\n# ====================================================\n# MIXUP / CUTMIX FUNCTIONS\n# ====================================================\ndef rand_bbox(size, lam):\n    \"\"\"Generate random bounding box for CutMix\"\"\"\n    W = size[2]\n    H = size[3]\n    cut_rat = np.sqrt(1. - lam)\n    cut_w = int(W * cut_rat)\n    cut_h = int(H * cut_rat)\n\n    cx = np.random.randint(W)\n    cy = np.random.randint(H)\n\n    bbx1 = np.clip(cx - cut_w // 2, 0, W)\n    bby1 = np.clip(cy - cut_h // 2, 0, H)\n    bbx2 = np.clip(cx + cut_w // 2, 0, W)\n    bby2 = np.clip(cy + cut_h // 2, 0, H)\n\n    return bbx1, bby1, bbx2, bby2\n\ndef mixup_data(x, y, alpha=0.2):\n    \"\"\"Apply MixUp augmentation\"\"\"\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size()[0]\n    index = torch.randperm(batch_size).to(x.device)\n\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef cutmix_data(x, y, alpha=1.0):\n    \"\"\"Apply CutMix augmentation\"\"\"\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size()[0]\n    index = torch.randperm(batch_size).to(x.device)\n\n    bbx1, bby1, bbx2, bby2 = rand_bbox(x.size(), lam)\n    x[:, :, bbx1:bbx2, bby1:bby2] = x[index, :, bbx1:bbx2, bby1:bby2]\n    \n    # Adjust lambda to exactly match pixel ratio\n    lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2]))\n    \n    y_a, y_b = y, y[index]\n    return x, y_a, y_b, lam\n\ndef mixup_criterion(criterion, pred, y_a, y_b, lam):\n    \"\"\"Calculate mixed loss\"\"\"\n    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)\n\n# ====================================================\n# PLOTTING FUNCTION\n# ====================================================\ndef plot_training_metrics(history, model_name, fold):\n    \"\"\"Plot training/validation loss and R² in one figure\"\"\"\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))\n    \n    epochs = range(1, len(history['train_loss']) + 1)\n    \n    ax1.plot(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2.5, marker='o', markersize=6)\n    ax1.plot(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2.5, marker='s', markersize=6)\n    ax1.set_xlabel('Epoch', fontsize=13, fontweight='bold')\n    ax1.set_ylabel('Loss', fontsize=13, fontweight='bold')\n    ax1.set_title(f'{model_name.upper()} - Fold {fold+1} | Loss', fontsize=15, fontweight='bold')\n    ax1.legend(fontsize=12, loc='upper right')\n    ax1.grid(True, alpha=0.3)\n    \n    ax2.plot(epochs, history['val_r2'], 'g-', label='Val R²', linewidth=3, marker='D', markersize=7)\n    best_r2 = max(history['val_r2'])\n    best_epoch = history['val_r2'].index(best_r2) + 1\n    ax2.axhline(y=best_r2, color='red', linestyle='--', linewidth=2, alpha=0.7)\n    ax2.axvline(x=best_epoch, color='orange', linestyle='--', linewidth=2, alpha=0.7)\n    ax2.scatter([best_epoch], [best_r2], color='red', s=200, zorder=5, edgecolors='black', linewidths=2)\n    ax2.text(best_epoch, best_r2, f'  Best: {best_r2:.4f}\\n  Epoch: {best_epoch}', \n             fontsize=11, fontweight='bold', va='bottom')\n    ax2.set_xlabel('Epoch', fontsize=13, fontweight='bold')\n    ax2.set_ylabel('R² Score', fontsize=13, fontweight='bold')\n    ax2.set_title(f'{model_name.upper()} - Fold {fold+1} | R²', fontsize=15, fontweight='bold')\n    ax2.legend(fontsize=12, loc='lower right')\n    ax2.grid(True, alpha=0.3)\n    \n    plt.tight_layout()\n    plt.savefig(f'{model_name}_fold{fold}_metrics.png', dpi=150, bbox_inches='tight')\n    plt.show()\n    plt.close()\n\n# ====================================================\n# DATA PREPARATION\n# ====================================================\ndef get_data():\n    df = pd.read_csv(Config.TRAIN_CSV)\n    \n    df_pivot = df.pivot(index='image_path', columns='target_name', values='target').reset_index()\n    \n    meta_df = df.groupby('image_path').agg({\n        'State': 'first',\n        'Species': 'first',\n        'Pre_GSHH_NDVI': 'first',\n        'Height_Ave_cm': 'first',\n        'Sampling_Date': 'first'\n    }).reset_index()\n    \n    train_df = df_pivot.merge(meta_df, on='image_path', how='left').fillna(0)\n    \n    state_le = LabelEncoder()\n    train_df['state_idx'] = state_le.fit_transform(train_df['State'].astype(str))\n    \n    target_cols = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    \n    return train_df, target_cols, state_le\n\n# ====================================================\n# DATASET\n# ====================================================\nclass BiomassDataset(Dataset):\n    def __init__(self, df, target_cols, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.target_cols = target_cols\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(Config.IMG_DIR, row['image_path'])\n        \n        image = cv2.imread(img_path)\n        if image is None:\n            image = np.zeros((384, 384, 3), dtype=np.uint8)\n        else:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            image = self.transform(image=image)['image']\n        \n        targets = row[self.target_cols].values.astype(np.float32)\n        targets_norm = (targets - Config.TARGET_MEAN) / (Config.TARGET_STD + 1e-6)\n        \n        return image, torch.tensor(targets_norm), row['image_path']\n\n# ====================================================\n# TRANSFORMS\n# ====================================================\ndef get_transforms(img_size, data='train'):\n    if data == 'train':\n        return A.Compose([\n            A.Resize(img_size, img_size),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.7),\n            A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),\n            A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.5),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2()\n        ])\n    else:\n        return A.Compose([\n            A.Resize(img_size, 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# MODELS\n# ====================================================\nclass SwinRegressor(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model(Config.SWIN_MODEL, pretrained=True, num_classes=0, global_pool='avg')\n        feat_dim = self.backbone.num_features\n        print(f\"Swin feature dimension: {feat_dim}\")\n        \n        self.fc1 = nn.Linear(feat_dim, 512)\n        self.ln1 = nn.LayerNorm(512)\n        self.gelu1 = nn.GELU()\n        self.drop1 = nn.Dropout(0.5)\n        \n        self.fc2 = nn.Linear(512, 256)\n        self.ln2 = nn.LayerNorm(256)\n        self.gelu2 = nn.GELU()\n        self.drop2 = nn.Dropout(0.4)\n        \n        self.fc3 = nn.Linear(256, 5)\n    \n    def forward(self, x, return_layers=False):\n        emb = self.backbone(x)\n        \n        x1 = self.fc1(emb)\n        x1 = self.ln1(x1)\n        x1 = self.gelu1(x1)\n        x1 = self.drop1(x1)\n        \n        x2 = self.fc2(x1)\n        x2 = self.ln2(x2)\n        x2 = self.gelu2(x2)\n        last_but_one = x2\n        x2 = self.drop2(x2)\n        \n        out = self.fc3(x2)\n        \n        if return_layers:\n            return out, emb, last_but_one\n        return out\n    \n    def get_param_groups(self):\n        \"\"\"Return backbone and head parameters separately for discriminative LR\"\"\"\n        backbone_params = []\n        head_params = []\n        \n        for name, param in self.named_parameters():\n            if 'backbone' in name:\n                backbone_params.append(param)\n            else:\n                head_params.append(param)\n        \n        return [\n            {'params': backbone_params, 'lr': Config.BACKBONE_LR},\n            {'params': head_params, 'lr': Config.HEAD_LR}\n        ]\n\nclass DINOv3Regressor(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = AutoModel.from_pretrained(Config.DINOV3_MODEL, trust_remote_code=True)\n        feat_dim = self.backbone.config.hidden_size\n        print(f\"DINOv3 feature dimension: {feat_dim}\")\n        \n        self.fc1 = nn.Linear(feat_dim, 512)\n        self.ln1 = nn.LayerNorm(512)\n        self.gelu1 = nn.GELU()\n        self.drop1 = nn.Dropout(0.5)\n        \n        self.fc2 = nn.Linear(512, 256)\n        self.ln2 = nn.LayerNorm(256)\n        self.gelu2 = nn.GELU()\n        self.drop2 = nn.Dropout(0.4)\n        \n        self.fc3 = nn.Linear(256, 5)\n    \n    def forward(self, x, return_layers=False):\n        emb = self.backbone(pixel_values=x).pooler_output\n        \n        x1 = self.fc1(emb)\n        x1 = self.ln1(x1)\n        x1 = self.gelu1(x1)\n        x1 = self.drop1(x1)\n        \n        x2 = self.fc2(x1)\n        x2 = self.ln2(x2)\n        x2 = self.gelu2(x2)\n        last_but_one = x2\n        x2 = self.drop2(x2)\n        \n        out = self.fc3(x2)\n        \n        if return_layers:\n            return out, emb, last_but_one\n        return out\n    \n    def get_param_groups(self):\n        \"\"\"Return backbone and head parameters separately for discriminative LR\"\"\"\n        backbone_params = []\n        head_params = []\n        \n        for name, param in self.named_parameters():\n            if 'backbone' in name:\n                backbone_params.append(param)\n            else:\n                head_params.append(param)\n        \n        return [\n            {'params': backbone_params, 'lr': Config.BACKBONE_LR},\n            {'params': head_params, 'lr': Config.HEAD_LR}\n        ]\n\nclass SigLIPRegressor(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model(Config.SIGLIP_MODEL, pretrained=True, num_classes=0, global_pool='avg')\n        feat_dim = self.backbone.num_features\n        print(f\"SigLIP feature dimension: {feat_dim}\")\n        \n        self.fc1 = nn.Linear(feat_dim, 512)\n        self.ln1 = nn.LayerNorm(512)\n        self.gelu1 = nn.GELU()\n        self.drop1 = nn.Dropout(0.5)\n        \n        self.fc2 = nn.Linear(512, 256)\n        self.ln2 = nn.LayerNorm(256)\n        self.gelu2 = nn.GELU()\n        self.drop2 = nn.Dropout(0.4)\n        \n        self.fc3 = nn.Linear(256, 5)\n    \n    def forward(self, x, return_layers=False):\n        emb = self.backbone(x)\n        \n        x1 = self.fc1(emb)\n        x1 = self.ln1(x1)\n        x1 = self.gelu1(x1)\n        x1 = self.drop1(x1)\n        \n        x2 = self.fc2(x1)\n        x2 = self.ln2(x2)\n        x2 = self.gelu2(x2)\n        last_but_one = x2\n        x2 = self.drop2(x2)\n        \n        out = self.fc3(x2)\n        \n        if return_layers:\n            return out, emb, last_but_one\n        return out\n    \n    def get_param_groups(self):\n        \"\"\"Return backbone and head parameters separately for discriminative LR\"\"\"\n        backbone_params = []\n        head_params = []\n        \n        for name, param in self.named_parameters():\n            if 'backbone' in name:\n                backbone_params.append(param)\n            else:\n                head_params.append(param)\n        \n        return [\n            {'params': backbone_params, 'lr': Config.BACKBONE_LR},\n            {'params': head_params, 'lr': Config.HEAD_LR}\n        ]\n\nclass LearnableEnsemble(nn.Module):\n    def __init__(self, num_models=5, num_targets=5):\n        super().__init__()\n        self.weights = nn.Parameter(torch.ones(num_models, num_targets) / num_models)\n    \n    def forward(self, predictions_list):\n        stacked = torch.stack(predictions_list, dim=0)\n        normalized_weights = F.softmax(self.weights, dim=0)\n        ensemble_pred = torch.einsum('mbt,mt->bt', stacked, normalized_weights)\n        return ensemble_pred, normalized_weights\n\n# ====================================================\n# EVALUATION\n# ====================================================\ndef evaluate_r2(preds, targets):\n    weights = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\n    w_j = np.tile(weights, (len(targets), 1))\n    y_bar_w = np.sum(w_j * targets) / np.sum(w_j)\n    ss_res = np.sum(w_j * (targets - preds)**2)\n    ss_tot = np.sum(w_j * (targets - y_bar_w)**2)\n    weighted_r2 = 1 - (ss_res / ss_tot)\n    per_target_r2 = r2_score(targets, preds, multioutput='raw_values')\n    return weighted_r2, per_target_r2\n\n# ====================================================\n# TRAINING (WITH ONECYCLELR + DISCRIMINATIVE LR + MIXUP/CUTMIX)\n# ====================================================\ndef train_single_model(model, train_loader, val_loader, model_name, fold):\n    print(f\"\\n{'='*70}\")\n    print(f\"Training {model_name} - Fold {fold}\")\n    print(f\"{'='*70}\")\n    \n    # Use discriminative learning rates\n    param_groups = model.get_param_groups()\n    optimizer = optim.AdamW(param_groups, weight_decay=Config.WEIGHT_DECAY)\n    \n    # OneCycleLR scheduler\n    steps_per_epoch = len(train_loader) // Config.ACCUMULATION_STEPS\n    total_steps = Config.EPOCHS_STAGE1 * steps_per_epoch\n    \n    scheduler = optim.lr_scheduler.OneCycleLR(\n        optimizer,\n        max_lr=[Config.MAX_LR_BACKBONE, Config.MAX_LR_HEAD],\n        total_steps=total_steps,\n        pct_start=Config.PCT_START,\n        anneal_strategy='cos',\n        div_factor=25.0,\n        final_div_factor=10000.0\n    )\n    \n    criterion = nn.SmoothL1Loss()\n    scaler = torch.cuda.amp.GradScaler()\n    \n    best_r2 = -np.inf\n    patience = 0\n    max_patience = 5  # Increased patience for OneCycleLR\n    \n    history = {'train_loss': [], 'val_loss': [], 'val_r2': []}\n    \n    for epoch in range(Config.EPOCHS_STAGE1):\n        # Train\n        model.train()\n        train_loss = 0.0\n        optimizer.zero_grad()\n        \n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{Config.EPOCHS_STAGE1}\")\n        for batch_idx, (imgs, targets_norm, _) in enumerate(pbar):\n            imgs = imgs.to(Config.DEVICE)\n            targets_norm = targets_norm.to(Config.DEVICE)\n            \n            # Apply MixUp or CutMix randomly\n            use_mixup = random.random() < Config.MIXUP_PROB\n            if use_mixup:\n                if random.random() < 0.5:\n                    # MixUp\n                    imgs, targets_a, targets_b, lam = mixup_data(imgs, targets_norm, Config.MIXUP_ALPHA)\n                else:\n                    # CutMix\n                    imgs, targets_a, targets_b, lam = cutmix_data(imgs, targets_norm, Config.CUTMIX_ALPHA)\n                \n                with torch.cuda.amp.autocast():\n                    preds = model(imgs)\n                    loss = mixup_criterion(criterion, preds, targets_a, targets_b, lam)\n                    loss = loss / Config.ACCUMULATION_STEPS\n            else:\n                # Standard training\n                with torch.cuda.amp.autocast():\n                    preds = model(imgs)\n                    loss = criterion(preds, targets_norm)\n                    loss = loss / Config.ACCUMULATION_STEPS\n            \n            scaler.scale(loss).backward()\n            \n            if (batch_idx + 1) % Config.ACCUMULATION_STEPS == 0:\n                scaler.step(optimizer)\n                scaler.update()\n                scheduler.step()  # Step OneCycleLR after each optimizer step\n                optimizer.zero_grad()\n            \n            train_loss += loss.item() * Config.ACCUMULATION_STEPS\n            \n            # Show current LR in progress bar\n            current_lr_backbone = optimizer.param_groups[0]['lr']\n            current_lr_head = optimizer.param_groups[1]['lr']\n            pbar.set_postfix({\n                'loss': loss.item() * Config.ACCUMULATION_STEPS,\n                'lr_bb': f'{current_lr_backbone:.2e}',\n                'lr_head': f'{current_lr_head:.2e}'\n            })\n            \n            if (batch_idx + 1) % 10 == 0:\n                torch.cuda.empty_cache()\n        \n        train_loss /= len(train_loader)\n        \n        # Validate\n        model.eval()\n        val_loss = 0.0\n        val_preds = []\n        val_targets = []\n        \n        with torch.no_grad():\n            for imgs, targets_norm, _ in val_loader:\n                imgs = imgs.to(Config.DEVICE)\n                targets_norm_gpu = targets_norm.to(Config.DEVICE)\n                \n                with torch.cuda.amp.autocast():\n                    preds = model(imgs)\n                    loss = criterion(preds, targets_norm_gpu)\n                    val_loss += loss.item()\n                    preds = preds.cpu().numpy()\n                \n                val_preds.append(preds)\n                val_targets.append(targets_norm.numpy())\n        \n        val_loss /= len(val_loader)\n        val_preds = np.vstack(val_preds)\n        val_targets = np.vstack(val_targets)\n        \n        val_preds_denorm = val_preds * Config.TARGET_STD + Config.TARGET_MEAN\n        val_targets_denorm = val_targets * Config.TARGET_STD + Config.TARGET_MEAN\n        val_preds_denorm = np.maximum(0, val_preds_denorm)\n        \n        val_r2, per_target = evaluate_r2(val_preds_denorm, val_targets_denorm)\n        \n        history['train_loss'].append(train_loss)\n        history['val_loss'].append(val_loss)\n        history['val_r2'].append(val_r2)\n        \n        print(f\"Epoch {epoch+1}: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}, Val R²={val_r2:.4f}\")\n        \n        if val_r2 > best_r2:\n            best_r2 = val_r2\n            patience = 0\n            torch.save(model.state_dict(), f'{model_name}_fold{fold}_best.pth')\n            print(f\"  ✓ Saved best model (R²={best_r2:.4f})\")\n        else:\n            patience += 1\n            if patience >= max_patience:\n                print(f\"  Early stopping at epoch {epoch+1}\")\n                break\n        \n        clear_gpu_memory()\n    \n    plot_training_metrics(history, model_name, fold)\n    model.load_state_dict(torch.load(f'{model_name}_fold{fold}_best.pth'))\n    return model, best_r2\n\n# ====================================================\n# EXTRACT LAYERS\n# ====================================================\ndef extract_layers_and_predictions(model, loader, model_name):\n    \"\"\"Extract embedding, last-but-one layer, and predictions\"\"\"\n    model.eval()\n    all_preds = []\n    all_embeddings = []\n    all_last_but_one = []\n    all_paths = []\n    \n    with torch.no_grad():\n        for imgs, _, paths in tqdm(loader, desc=f\"Extracting {model_name}\"):\n            imgs = imgs.to(Config.DEVICE)\n            with torch.cuda.amp.autocast():\n                preds, embeddings, last_but_one = model(imgs, return_layers=True)\n            all_preds.append(preds.cpu().numpy())\n            all_embeddings.append(embeddings.cpu().numpy())\n            all_last_but_one.append(last_but_one.cpu().numpy())\n            all_paths.extend(paths)\n            \n            del preds, embeddings, last_but_one, imgs\n            torch.cuda.empty_cache()\n    \n    all_preds = np.vstack(all_preds)\n    all_embeddings = np.vstack(all_embeddings)\n    all_last_but_one = np.vstack(all_last_but_one)\n    return all_preds, all_embeddings, all_last_but_one, all_paths\n\n# ====================================================\n# TRAIN RF AND XGBOOST WITH PCA (OOF)\n# ====================================================\ndef train_tree_models_with_pca(\n    train_swin_emb, train_swin_l1, train_dinov3_emb, train_dinov3_l1, train_siglip_emb, train_siglip_l1,\n    train_swin_preds, train_dinov3_preds, train_siglip_preds, train_targets,\n    val_swin_emb, val_swin_l1, val_dinov3_emb, val_dinov3_l1, val_siglip_emb, val_siglip_l1,\n    val_swin_preds, val_dinov3_preds, val_siglip_preds, val_targets,\n    fold\n):\n    \"\"\"Train RF and XGBoost on TRAIN data, evaluate on VAL data\"\"\"\n    print(f\"\\n{'='*70}\")\n    print(f\"Training Tree Models with PCA (OOF) - Fold {fold}\")\n    print(f\"{'='*70}\")\n    \n    X_train = np.hstack([\n        train_swin_emb, train_swin_l1,\n        train_dinov3_emb, train_dinov3_l1,\n        train_siglip_emb, train_siglip_l1,\n        train_swin_preds, train_dinov3_preds, train_siglip_preds\n    ])\n    print(f\"Train feature shape before PCA: {X_train.shape}\")\n    \n    X_val = np.hstack([\n        val_swin_emb, val_swin_l1,\n        val_dinov3_emb, val_dinov3_l1,\n        val_siglip_emb, val_siglip_l1,\n        val_swin_preds, val_dinov3_preds, val_siglip_preds\n    ])\n    print(f\"Val feature shape before PCA: {X_val.shape}\")\n    \n    pca = PCA(n_components=Config.N_PCA_COMPONENTS, random_state=Config.SEED)\n    X_train_pca = pca.fit_transform(X_train)\n    X_val_pca = pca.transform(X_val)\n    \n    explained_var = pca.explained_variance_ratio_.sum()\n    print(f\"PCA reduced to {Config.N_PCA_COMPONENTS} components\")\n    print(f\"Explained variance: {explained_var:.4f}\")\n    print(f\"Train PCA shape: {X_train_pca.shape}\")\n    print(f\"Val PCA shape: {X_val_pca.shape}\")\n    \n    with open(f'pca_fold{fold}.pkl', 'wb') as f:\n        pickle.dump(pca, f)\n    \n    del X_train, X_val\n    gc.collect()\n    \n    print(\"\\nTraining Random Forest on TRAIN PCA features...\")\n    rf_models = []\n    rf_val_preds = np.zeros_like(val_targets)\n    \n    for target_idx in range(5):\n        rf = RandomForestRegressor(\n            n_estimators=200,\n            max_depth=10,\n            min_samples_split=5,\n            min_samples_leaf=2,\n            random_state=Config.SEED,\n            n_jobs=-1\n        )\n        rf.fit(X_train_pca, train_targets[:, target_idx])\n        rf_val_preds[:, target_idx] = rf.predict(X_val_pca)\n        rf_models.append(rf)\n    \n    rf_val_preds = np.maximum(0, rf_val_preds)\n    rf_r2, _ = evaluate_r2(rf_val_preds, val_targets)\n    print(f\"Random Forest VAL R²: {rf_r2:.4f}\")\n    \n    print(\"\\nTraining XGBoost on TRAIN PCA features...\")\n    xgb_models = []\n    xgb_val_preds = np.zeros_like(val_targets)\n    \n    for target_idx in range(5):\n        xgb = XGBRegressor(\n            n_estimators=200,\n            max_depth=6,\n            learning_rate=0.05,\n            subsample=0.8,\n            colsample_bytree=0.8,\n            random_state=Config.SEED,\n            n_jobs=-1,\n            tree_method='hist'\n        )\n        xgb.fit(X_train_pca, train_targets[:, target_idx])\n        xgb_val_preds[:, target_idx] = xgb.predict(X_val_pca)\n        xgb_models.append(xgb)\n    \n    xgb_val_preds = np.maximum(0, xgb_val_preds)\n    xgb_r2, _ = evaluate_r2(xgb_val_preds, val_targets)\n    print(f\"XGBoost VAL R²: {xgb_r2:.4f}\")\n    \n    with open(f'rf_models_fold{fold}.pkl', 'wb') as f:\n        pickle.dump(rf_models, f)\n    with open(f'xgb_models_fold{fold}.pkl', 'wb') as f:\n        pickle.dump(xgb_models, f)\n    \n    return rf_models, xgb_models, rf_val_preds, xgb_val_preds, rf_r2, xgb_r2, pca\n\n# ====================================================\n# ENSEMBLE WEIGHTS\n# ====================================================\ndef train_ensemble_weights(swin_preds, dinov3_preds, siglip_preds, rf_preds, xgb_preds, targets, fold, method='learnable'):\n    print(f\"\\n{'='*70}\")\n    print(f\"Training Ensemble Weights ({method})\")\n    print(f\"{'='*70}\")\n    \n    swin_preds = swin_preds.astype(np.float32)\n    dinov3_preds = dinov3_preds.astype(np.float32)\n    siglip_preds = siglip_preds.astype(np.float32)\n    rf_preds = rf_preds.astype(np.float32)\n    xgb_preds = xgb_preds.astype(np.float32)\n    targets = targets.astype(np.float32)\n    \n    if method == 'learnable':\n        dataset = torch.utils.data.TensorDataset(\n            torch.tensor(swin_preds, dtype=torch.float32),\n            torch.tensor(dinov3_preds, dtype=torch.float32),\n            torch.tensor(siglip_preds, dtype=torch.float32),\n            torch.tensor(rf_preds, dtype=torch.float32),\n            torch.tensor(xgb_preds, dtype=torch.float32),\n            torch.tensor(targets, dtype=torch.float32)\n        )\n        loader = DataLoader(dataset, batch_size=32, shuffle=True)\n        \n        ensemble_model = LearnableEnsemble(num_models=5).to(Config.DEVICE)\n        optimizer = optim.Adam(ensemble_model.parameters(), lr=1e-3)\n        criterion = nn.MSELoss()\n        \n        for epoch in range(100):\n            ensemble_model.train()\n            total_loss = 0\n            \n            for swin_b, dinov3_b, siglip_b, rf_b, xgb_b, targets_b in loader:\n                swin_b = swin_b.to(Config.DEVICE)\n                dinov3_b = dinov3_b.to(Config.DEVICE)\n                siglip_b = siglip_b.to(Config.DEVICE)\n                rf_b = rf_b.to(Config.DEVICE)\n                xgb_b = xgb_b.to(Config.DEVICE)\n                targets_b = targets_b.to(Config.DEVICE)\n                \n                optimizer.zero_grad()\n                ensemble_pred, _ = ensemble_model([swin_b, dinov3_b, siglip_b, rf_b, xgb_b])\n                loss = criterion(ensemble_pred, targets_b)\n                loss.backward()\n                optimizer.step()\n                total_loss += loss.item()\n            \n            if (epoch + 1) % 20 == 0:\n                print(f\"Epoch {epoch+1}: Loss={total_loss/len(loader):.4f}\")\n        \n        ensemble_model.eval()\n        with torch.no_grad():\n            ensemble_preds, final_weights = ensemble_model([\n                torch.tensor(swin_preds, dtype=torch.float32).to(Config.DEVICE),\n                torch.tensor(dinov3_preds, dtype=torch.float32).to(Config.DEVICE),\n                torch.tensor(siglip_preds, dtype=torch.float32).to(Config.DEVICE),\n                torch.tensor(rf_preds, dtype=torch.float32).to(Config.DEVICE),\n                torch.tensor(xgb_preds, dtype=torch.float32).to(Config.DEVICE)\n            ])\n            ensemble_preds = ensemble_preds.cpu().numpy()\n            final_weights = final_weights.cpu().numpy()\n        \n        ensemble_preds = np.maximum(0, ensemble_preds)\n        ensemble_r2, _ = evaluate_r2(ensemble_preds, targets)\n        \n        print(f\"\\nLearnable Ensemble R²: {ensemble_r2:.4f}\")\n        avg_weights = final_weights.mean(axis=1)\n        print(f\"Avg Weights: Swin={avg_weights[0]:.4f}, DINOv3={avg_weights[1]:.4f}, SigLIP={avg_weights[2]:.4f}, RF={avg_weights[3]:.4f}, XGB={avg_weights[4]:.4f}\")\n        \n        torch.save(ensemble_model.state_dict(), f'learnable_ensemble_fold{fold}.pth')\n        return avg_weights, ensemble_preds\n    \n    else:\n        ensemble_preds = (swin_preds + dinov3_preds + siglip_preds + rf_preds + xgb_preds) / 5\n        ensemble_r2, _ = evaluate_r2(ensemble_preds, targets)\n        print(f\"Simple Average R²: {ensemble_r2:.4f}\")\n        return np.array([0.2, 0.2, 0.2, 0.2, 0.2]), ensemble_preds\n\n# ====================================================\n# MAIN\n# ====================================================\ndef main():\n    print(\"=\"*70)\n    print(\"ENSEMBLE: ONECYCLELR + DISCRIMINATIVE LR + MIXUP/CUTMIX\")\n    print(\"=\"*70)\n    \n    train_df, target_cols, state_le = get_data()\n    train_df = train_df[~train_df['image_path'].str.contains('ID230058600')].reset_index(drop=True)\n    print(f\"\\n✓ Training samples: {len(train_df)}\")\n    \n    sgkf = StratifiedGroupKFold(n_splits=Config.N_FOLDS, shuffle=True, random_state=Config.SEED)\n    splits = list(sgkf.split(train_df, y=train_df['state_idx'], groups=train_df['Sampling_Date']))\n    \n    fold_results = []\n    \n    for fold, (train_idx, val_idx) in enumerate(splits):\n        print(f\"\\n\\n{'#'*70}\")\n        print(f\"FOLD {fold+1}/{Config.N_FOLDS}\")\n        print(f\"{'#'*70}\\n\")\n        \n        df_train = train_df.iloc[train_idx].reset_index(drop=True)\n        df_val = train_df.iloc[val_idx].reset_index(drop=True)\n        \n        Config.TARGET_MEAN = df_train[target_cols].mean().values.astype(np.float32)\n        Config.TARGET_STD = df_train[target_cols].std().values.astype(np.float32)\n        print(f\"Fold {fold+1} TARGET_MEAN: {Config.TARGET_MEAN}\")\n        print(f\"Fold {fold+1} TARGET_STD: {Config.TARGET_STD}\")\n        \n        train_targets = df_train[target_cols].values.astype(np.float32)\n        val_targets = df_val[target_cols].values.astype(np.float32)\n        \n        # ===== SWIN =====\n        print(f\"\\n{'='*70}\\nSTAGE 1: SWIN\\n{'='*70}\")\n        swin_train_ds = BiomassDataset(df_train, target_cols, get_transforms(Config.SWIN_IMG_SIZE, 'train'))\n        swin_val_ds = BiomassDataset(df_val, target_cols, get_transforms(Config.SWIN_IMG_SIZE, 'val'))\n        swin_train_loader = DataLoader(swin_train_ds, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n        swin_val_loader = DataLoader(swin_val_ds, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n        \n        swin_model = SwinRegressor().to(Config.DEVICE)\n        swin_model, swin_r2 = train_single_model(swin_model, swin_train_loader, swin_val_loader, 'swin', fold)\n        \n        print(\"Extracting TRAIN features...\")\n        swin_train_preds, swin_train_emb, swin_train_l1, _ = extract_layers_and_predictions(swin_model, swin_train_loader, 'Swin-Train')\n        swin_train_preds = np.maximum(0, swin_train_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        print(\"Extracting VAL features...\")\n        swin_val_preds, swin_val_emb, swin_val_l1, _ = extract_layers_and_predictions(swin_model, swin_val_loader, 'Swin-Val')\n        swin_val_preds = np.maximum(0, swin_val_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        del swin_model, swin_train_loader, swin_val_loader, swin_train_ds, swin_val_ds\n        clear_gpu_memory()\n        \n        # ===== DINOV3 =====\n        print(f\"\\n{'='*70}\\nSTAGE 2: DINOV3\\n{'='*70}\")\n        dinov3_train_ds = BiomassDataset(df_train, target_cols, get_transforms(Config.DINOV3_IMG_SIZE, 'train'))\n        dinov3_val_ds = BiomassDataset(df_val, target_cols, get_transforms(Config.DINOV3_IMG_SIZE, 'val'))\n        dinov3_train_loader = DataLoader(dinov3_train_ds, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n        dinov3_val_loader = DataLoader(dinov3_val_ds, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n        \n        dinov3_model = DINOv3Regressor().to(Config.DEVICE)\n        dinov3_model, dinov3_r2 = train_single_model(dinov3_model, dinov3_train_loader, dinov3_val_loader, 'dinov3', fold)\n        \n        print(\"Extracting TRAIN features...\")\n        dinov3_train_preds, dinov3_train_emb, dinov3_train_l1, _ = extract_layers_and_predictions(dinov3_model, dinov3_train_loader, 'DINOv3-Train')\n        dinov3_train_preds = np.maximum(0, dinov3_train_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        print(\"Extracting VAL features...\")\n        dinov3_val_preds, dinov3_val_emb, dinov3_val_l1, _ = extract_layers_and_predictions(dinov3_model, dinov3_val_loader, 'DINOv3-Val')\n        dinov3_val_preds = np.maximum(0, dinov3_val_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        del dinov3_model, dinov3_train_loader, dinov3_val_loader, dinov3_train_ds, dinov3_val_ds\n        clear_gpu_memory()\n        \n        # ===== SIGLIP =====\n        print(f\"\\n{'='*70}\\nSTAGE 3: SIGLIP\\n{'='*70}\")\n        siglip_train_ds = BiomassDataset(df_train, target_cols, get_transforms(Config.SIGLIP_IMG_SIZE, 'train'))\n        siglip_val_ds = BiomassDataset(df_val, target_cols, get_transforms(Config.SIGLIP_IMG_SIZE, 'val'))\n        siglip_train_loader = DataLoader(siglip_train_ds, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n        siglip_val_loader = DataLoader(siglip_val_ds, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n        \n        siglip_model = SigLIPRegressor().to(Config.DEVICE)\n        siglip_model, siglip_r2 = train_single_model(siglip_model, siglip_train_loader, siglip_val_loader, 'siglip', fold)\n        \n        print(\"Extracting TRAIN features...\")\n        siglip_train_preds, siglip_train_emb, siglip_train_l1, _ = extract_layers_and_predictions(siglip_model, siglip_train_loader, 'SigLIP-Train')\n        siglip_train_preds = np.maximum(0, siglip_train_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        print(\"Extracting VAL features...\")\n        siglip_val_preds, siglip_val_emb, siglip_val_l1, _ = extract_layers_and_predictions(siglip_model, siglip_val_loader, 'SigLIP-Val')\n        siglip_val_preds = np.maximum(0, siglip_val_preds * Config.TARGET_STD + Config.TARGET_MEAN)\n        \n        del siglip_model, siglip_train_loader, siglip_val_loader, siglip_train_ds, siglip_val_ds\n        clear_gpu_memory()\n        \n        # ===== TREE MODELS WITH PCA (OOF) =====\n        print(f\"\\n{'='*70}\\nSTAGE 4: TREE MODELS WITH PCA (OOF)\\n{'='*70}\")\n        rf_models, xgb_models, rf_val_preds, xgb_val_preds, rf_r2, xgb_r2, pca = train_tree_models_with_pca(\n            swin_train_emb, swin_train_l1, dinov3_train_emb, dinov3_train_l1, siglip_train_emb, siglip_train_l1,\n            swin_train_preds, dinov3_train_preds, siglip_train_preds, train_targets,\n            swin_val_emb, swin_val_l1, dinov3_val_emb, dinov3_val_l1, siglip_val_emb, siglip_val_l1,\n            swin_val_preds, dinov3_val_preds, siglip_val_preds, val_targets,\n            fold\n        )\n        \n        del (swin_train_emb, swin_train_l1, swin_train_preds,\n             dinov3_train_emb, dinov3_train_l1, dinov3_train_preds,\n             siglip_train_emb, siglip_train_l1, siglip_train_preds)\n        del (swin_val_emb, swin_val_l1, dinov3_val_emb, dinov3_val_l1, siglip_val_emb, siglip_val_l1)\n        gc.collect()\n        \n        # ===== ENSEMBLE =====\n        ensemble_weights, ensemble_preds = train_ensemble_weights(\n            swin_val_preds, dinov3_val_preds, siglip_val_preds, rf_val_preds, xgb_val_preds,\n            val_targets, fold, method=Config.ENSEMBLE_METHOD\n        )\n        ensemble_r2, per_target_r2 = evaluate_r2(ensemble_preds, val_targets)\n        \n        fold_results.append({\n            'fold': fold,\n            'swin_r2': swin_r2,\n            'dinov3_r2': dinov3_r2,\n            'siglip_r2': siglip_r2,\n            'rf_r2': rf_r2,\n            'xgb_r2': xgb_r2,\n            'ensemble_r2': ensemble_r2,\n            'ensemble_weights': ensemble_weights,\n            'per_target_r2': per_target_r2\n        })\n        \n        print(f\"\\n{'='*70}\")\n        print(f\"FOLD {fold+1} SUMMARY\")\n        print(f\"{'='*70}\")\n        print(f\"Swin: {swin_r2:.4f} | DINOv3: {dinov3_r2:.4f} | SigLIP: {siglip_r2:.4f} | RF: {rf_r2:.4f} | XGB: {xgb_r2:.4f} | Ensemble: {ensemble_r2:.4f}\")\n        \n        clear_gpu_memory()\n    \n    # Final summary\n    print(f\"\\n\\n{'#'*70}\\nFINAL RESULTS\\n{'#'*70}\\n\")\n    avg_ensemble = np.mean([r['ensemble_r2'] for r in fold_results])\n    std_ensemble = np.std([r['ensemble_r2'] for r in fold_results])\n    \n    print(f\"Average Ensemble R²: {avg_ensemble:.4f} ± {std_ensemble:.4f}\")\n    print(f\"\\nPer-model averages:\")\n    print(f\"  Swin:   {np.mean([r['swin_r2'] for r in fold_results]):.4f} ± {np.std([r['swin_r2'] for r in fold_results]):.4f}\")\n    print(f\"  DINOv3: {np.mean([r['dinov3_r2'] for r in fold_results]):.4f} ± {np.std([r['dinov3_r2'] for r in fold_results]):.4f}\")\n    print(f\"  SigLIP: {np.mean([r['siglip_r2'] for r in fold_results]):.4f} ± {np.std([r['siglip_r2'] for r in fold_results]):.4f}\")\n    print(f\"  RF:     {np.mean([r['rf_r2'] for r in fold_results]):.4f} ± {np.std([r['rf_r2'] for r in fold_results]):.4f}\")\n    print(f\"  XGB:    {np.mean([r['xgb_r2'] for r in fold_results]):.4f} ± {np.std([r['xgb_r2'] for r in fold_results]):.4f}\")\n    \n    config_dict = {\n        'n_folds': Config.N_FOLDS,\n        'n_pca_components': Config.N_PCA_COMPONENTS,\n        'avg_ensemble_r2': float(avg_ensemble),\n        'std_ensemble_r2': float(std_ensemble),\n        'fold_results': [\n            {\n                'fold': r['fold'],\n                'swin_r2': float(r['swin_r2']),\n                'dinov3_r2': float(r['dinov3_r2']),\n                'siglip_r2': float(r['siglip_r2']),\n                'rf_r2': float(r['rf_r2']),\n                'xgb_r2': float(r['xgb_r2']),\n                'ensemble_r2': float(r['ensemble_r2'])\n            } for r in fold_results\n        ]\n    }\n    \n    with open('ensemble_config.json', 'w') as f:\n        json.dump(config_dict, f, indent=4)\n    \n    print(f\"\\n✓ Saved ensemble_config.json\")\n    print(f\"✓ Training complete!\")\n\nif __name__ == '__main__':\n    main()\n","metadata":{"execution":{"iopub.status.busy":"2026-01-20T17:00:51.715167Z","iopub.execute_input":"2026-01-20T17:00:51.715453Z","execution_failed":"2026-01-20T17:02:20.862Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}