{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ===============================================================\n# CSIRO Image2Biomass – Training (Weighted R² Validation)\n# ===============================================================\nimport os, gc, cv2, numpy as np, pandas as pd\nfrom tqdm import tqdm\nimport torch, torch.nn as nn, torch.optim as optim\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom sklearn.model_selection import KFold\n\n# ---------------------------------------------------------------\n# 1. CONFIG (memory-safe + R² metric)\n# ---------------------------------------------------------------\nclass CFG:\n    BASE_PATH       = '/kaggle/input/csiro-biomass'\n    TRAIN_CSV       = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    MODEL_DIR       = '/kaggle/working/'\n    N_FOLDS         = 5\n\n    MODEL_NAME = 'convnext_tiny'      # safe & matches inference\n    IMG_SIZE   = 512                  # fits easily\n    PRETRAINED = True\n\n    BATCH_SIZE   = 2\n    GRAD_ACC     = 4                  # effective batch = 8\n    NUM_WORKERS  = 1\n    EPOCHS       = 25\n    LR           = 1e-4\n    WD           = 1e-2\n    PATIENCE     = 5\n\n    TARGET_COLS    = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS   = ['Dry_Clover_g', 'Dry_Dead_g']\n    ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    R2_WEIGHTS     = np.array([0.1, 0.1, 0.1, 0.2, 0.5])  # matches metric\n\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Device : {CFG.DEVICE}\")\nprint(f\"Backbone: {CFG.MODEL_NAME} | Size: {CFG.IMG_SIZE}\")\n\n# ---------------------------------------------------------------\n# 2. WEIGHTED R² METRIC (your function)\n# ---------------------------------------------------------------\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    \"\"\"\n    y_true, y_pred: shape (N, 5)\n    weights: [0.1, 0.1, 0.1, 0.2, 0.5]\n    \"\"\"\n    weights = CFG.R2_WEIGHTS\n    r2_scores = []\n    for i in range(5):\n        y_t = y_true[:, i]\n        y_p = y_pred[:, i]\n        ss_res = np.sum((y_t - y_p) ** 2)\n        ss_tot = np.sum((y_t - np.mean(y_t)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n    r2_scores = np.array(r2_scores)\n    weighted_r2 = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted_r2, r2_scores\n\n# ---------------------------------------------------------------\n# 3. AUGMENTATIONS\n# ---------------------------------------------------------------\ndef get_train_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.3),\n        A.Rotate(limit=(-10, 10), p=0.3,\n                 interpolation=cv2.INTER_LINEAR,\n                 border_mode=cv2.BORDER_REFLECT_101),\n        A.ColorJitter(brightness=0.1, contrast=0.1, p=0.3),\n        A.Normalize(mean=[0.485, 0.456, 0.406],\n                    std =[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef get_val_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.Normalize(mean=[0.485, 0.456, 0.406],\n                    std =[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\n# ---------------------------------------------------------------\n# 4. DATASET\n# ---------------------------------------------------------------\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df        = df\n        self.transform = transform\n        self.img_dir   = img_dir\n        self.paths     = df['image_path'].values\n        self.labels    = df[CFG.ALL_TARGET_COLS].values.astype(np.float32)\n\n    def __len__(self): return len(self.df)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.img_dir, os.path.basename(self.paths[idx]))\n        img  = cv2.imread(path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        h, w, _ = img.shape\n        mid = w // 2\n        left  = img[:, :mid]\n        right = img[:, mid:]\n\n        left  = self.transform(image=left)['image']\n        right = self.transform(image=right)['image']\n\n        label = torch.from_numpy(self.labels[idx])\n        return left, right, label\n\n# ---------------------------------------------------------------\n# 5. MODEL (safe pretrained loading)\n# ---------------------------------------------------------------\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(\n            model_name, pretrained=False, num_classes=0, global_pool='avg')\n        nf = self.backbone.num_features\n        comb = nf * 2\n\n        def head():\n            return nn.Sequential(\n                nn.Linear(comb, comb // 2),\n                nn.ReLU(inplace=True),\n                nn.Dropout(0.3),\n                nn.Linear(comb // 2, 1)\n            )\n        self.head_total = head()\n        self.head_gdm   = head()\n        self.head_green = head()\n\n        if pretrained:\n            self.load_pretrained()\n\n    def load_pretrained(self):\n        try:\n            state_dict = timm.create_model(CFG.MODEL_NAME, pretrained=True, num_classes=0).state_dict()\n            self.backbone.load_state_dict(state_dict, strict=False)\n            print(\"Pretrained weights loaded (CPU)\")\n        except Exception as e:\n            print(f\"Warning: Pretrained load failed: {e}\")\n\n    def forward(self, left, right):\n        fl = self.backbone(left)\n        fr = self.backbone(right)\n        x  = torch.cat([fl, fr], dim=1)\n        return (self.head_total(x), self.head_gdm(x), self.head_green(x))\n\n# ---------------------------------------------------------------\n# 6. LOSS (MSE on all 5)\n# ---------------------------------------------------------------\ndef biomass_loss(p_total, p_gdm, p_green, labels):\n    mse = nn.MSELoss()\n    loss_total = mse(p_total.squeeze(), labels[:, 4])\n    loss_gdm   = mse(p_gdm.squeeze(),   labels[:, 3])\n    loss_green = mse(p_green.squeeze(), labels[:, 0])\n\n    p_clover = torch.clamp(p_gdm - p_green, min=0)\n    p_dead   = torch.clamp(p_total - p_gdm, min=0)\n\n    loss_clover = mse(p_clover.squeeze(), labels[:, 2])\n    loss_dead   = mse(p_dead.squeeze(),   labels[:, 1])\n\n    return (loss_total + loss_gdm + loss_green + loss_clover + loss_dead) / 5\n\n# ---------------------------------------------------------------\n# 7. VALIDATION WITH WEIGHTED R²\n# ---------------------------------------------------------------\n@torch.no_grad()\ndef valid_epoch(model, loader, device):\n    model.eval()\n    running_loss = 0.0\n    preds = {'total':[], 'gdm':[], 'green':[]}\n    all_labels = []\n\n    for l, r, lab in tqdm(loader, desc='valid', leave=False):\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        p_tot, p_gdm, p_green = model(l, r)\n        loss = biomass_loss(p_tot, p_gdm, p_green, lab)\n        running_loss += loss.item() * l.size(0)\n\n        preds['total'].extend(p_tot.cpu().numpy().ravel())\n        preds['gdm'].extend(p_gdm.cpu().numpy().ravel())\n        preds['green'].extend(p_green.cpu().numpy().ravel())\n        all_labels.extend(lab.cpu().numpy())\n\n    # Convert to numpy\n    pred_total = np.array(preds['total'])\n    pred_gdm   = np.array(preds['gdm'])\n    pred_green = np.array(preds['green'])\n    true_labels = np.stack(all_labels)  # (N, 5)\n\n    # Compute derived\n    pred_clover = np.clip(pred_gdm - pred_green, 0, None)\n    pred_dead   = np.clip(pred_total - pred_gdm, 0, None)\n\n    # Stack predictions in correct order\n    pred_all = np.stack([\n        pred_green,      # Dry_Green_g\n        pred_dead,       # Dry_Dead_g\n        pred_clover,     # Dry_Clover_g\n        pred_gdm,        # GDM_g\n        pred_total       # Dry_Total_g\n    ], axis=1)\n\n    # Compute weighted R²\n    weighted_r2, per_target_r2 = weighted_r2_score(true_labels, pred_all)\n\n    return running_loss / len(loader.dataset), weighted_r2, per_target_r2, pred_all, true_labels\n\n# ---------------------------------------------------------------\n# 8. TRAINING LOOP\n# ---------------------------------------------------------------\ndef train_epoch(model, loader, opt, scheduler, device):\n    model.train()\n    running = 0.0\n    opt.zero_grad()\n    for i, (l, r, lab) in enumerate(tqdm(loader, desc='train', leave=False)):\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        p_tot, p_gdm, p_green = model(l, r)\n        loss = biomass_loss(p_tot, p_gdm, p_green, lab) / CFG.GRAD_ACC\n        loss.backward()\n        running += loss.item() * l.size(0) * CFG.GRAD_ACC\n\n        if (i + 1) % CFG.GRAD_ACC == 0 or (i + 1) == len(loader):\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            opt.step()\n            opt.zero_grad()\n\n    scheduler.step()\n    return running / len(loader.dataset)\n\n# ---------------------------------------------------------------\n# 9. MAIN – 5-FOLD WITH R² TRACKING\n# ---------------------------------------------------------------\nif __name__ == '__main__':\n    print(\"Loading data...\")\n    df_long = pd.read_csv(CFG.TRAIN_CSV)\n    df_wide = df_long.pivot(index='image_path', columns='target_name', values='target').reset_index()\n    df_wide = df_wide[['image_path'] + CFG.ALL_TARGET_COLS]\n    print(f\"{len(df_wide)} training images\")\n\n    kfold = KFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=42)\n\n    for fold, (tr_idx, val_idx) in enumerate(kfold.split(df_wide)):\n        print('\\n' + '='*70)\n        print(f'   FOLD {fold+1}/{CFG.N_FOLDS}   |   {len(tr_idx)} train / {len(val_idx)} val')\n        print('='*70)\n\n        torch.cuda.empty_cache()\n        gc.collect()\n\n        tr_df  = df_wide.iloc[tr_idx].reset_index(drop=True)\n        val_df = df_wide.iloc[val_idx].reset_index(drop=True)\n\n        tr_set = BiomassDataset(tr_df,  get_train_transforms(), CFG.TRAIN_IMAGE_DIR)\n        val_set= BiomassDataset(val_df,get_val_transforms(),   CFG.TRAIN_IMAGE_DIR)\n\n        tr_loader  = DataLoader(tr_set,  batch_size=CFG.BATCH_SIZE, shuffle=True,\n                               num_workers=CFG.NUM_WORKERS, pin_memory=True, drop_last=True)\n        val_loader = DataLoader(val_set, batch_size=CFG.BATCH_SIZE, shuffle=False,\n                               num_workers=CFG.NUM_WORKERS, pin_memory=True)\n\n        print(\"Building model...\")\n        model = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED)\n        model = model.to(CFG.DEVICE)\n        model = nn.DataParallel(model)\n\n        optimizer = optim.AdamW(model.parameters(), lr=CFG.LR, weight_decay=CFG.WD)\n        scheduler = CosineAnnealingLR(optimizer, T_max=CFG.EPOCHS)\n\n        best_r2 = -np.inf\n        patience = 0\n\n        for epoch in range(1, CFG.EPOCHS+1):\n            tr_loss = train_epoch(model, tr_loader, optimizer, scheduler, CFG.DEVICE)\n            val_loss, val_r2, per_r2, _, _ = valid_epoch(model, val_loader, CFG.DEVICE)\n\n            per_r2_str = \" | \".join([f\"{CFG.ALL_TARGET_COLS[i][:5]}: {r2:.3f}\" for i, r2 in enumerate(per_r2)])\n\n            print(f'Epoch {epoch:02d} | '\n                  f'TrainLoss {tr_loss:.5f} | '\n                  f'ValLoss {val_loss:.5f} | '\n                  f'ValR² {val_r2:.4f} {\"(BEST)\" if val_r2 > best_r2 else \"\"}')\n            print(f'     → {per_r2_str}')\n\n            if val_r2 > best_r2:\n                best_r2 = val_r2\n                save_path = os.path.join(CFG.MODEL_DIR, f'best_model_fold{fold}.pth')\n                torch.save(model.module.state_dict() if hasattr(model, 'module') else model.state_dict(), save_path)\n                print(f'   → SAVED (R²: {best_r2:.4f})')\n                patience = 0\n            else:\n                patience += 1\n                if patience >= CFG.PATIENCE:\n                    print(f'   → EARLY STOP (no improvement in {CFG.PATIENCE} epochs)')\n                    break\n\n        # Cleanup\n        del model, tr_loader, val_loader, optimizer, scheduler\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    print('\\nTraining complete! Best models saved in:', CFG.MODEL_DIR)\n    print('Use these in inference with:')\n    print('   MODEL_NAME = \"convnext_tiny\"')\n    print('   IMG_SIZE = 512')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-31T11:53:22.750056Z","iopub.execute_input":"2025-10-31T11:53:22.750364Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}