{"cells":[{"cell_type":"markdown","id":"header","metadata":{},"source":"# 🌿 CSIRO Biomass - Advanced DINOv2 Training\n==============================================\n\n**Competition**: Image2Biomass Prediction Challenge  \n**Target LB**: 0.77+\n\n## 🏆 Advanced Strategies Implemented:\n1. DINOv2-Giant backbone (1.1B params)\n2. **3→5 Target Derivation**: Predict Clover/Dead/Green, derive Total/GDM\n3. **NDVI Feature Integration** (Pre_GSHH_NDVI)\n4. **StratifiedKFold by State** for geographic diversity\n5. **Mean-Shifting Post-Processing**\n6. **Long-to-Wide Pivoting** for CSIRO dataset"},{"cell_type":"code","execution_count":null,"id":"imports","metadata":{},"outputs":[],"source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import OneCycleLR\n\nfrom transformers import AutoModel\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.model_selection import StratifiedKFold, KFold\nfrom sklearn.metrics import r2_score\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA: {torch.cuda.is_available()}\")"},{"cell_type":"markdown","id":"config_header","metadata":{},"source":"## 1. Configuration"},{"cell_type":"code","execution_count":null,"id":"config","metadata":{},"outputs":[],"source":"class CFG:\n    # Paths\n    data_dir = Path(\"/kaggle/input/csiro-biomass\")\n    if not os.path.exists(data_dir):\n        data_dir = Path(\"./data\")\n    \n    # Backbones\n    dinov2_path = \"/kaggle/input/dinov2/pytorch/giant/1\"\n    hidden_size = 1536\n    dinov2_layers = [36, 37, 38, 39, 40]\n    \n    # Target Strategy\n    base_targets = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g']\n    all_targets = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g']\n    target_weights = {'Dry_Total_g': 0.5, 'GDM_g': 0.2, 'Dry_Green_g': 0.1, 'Dry_Dead_g': 0.1, 'Dry_Clover_g': 0.1}\n    \n    # Metadata columns\n    ndvi_col = 'Pre_GSHH_NDVI'\n    state_col = 'State'\n    \n    # Training\n    image_size = 518\n    batch_size = 4\n    accumulation_steps = 2 \n    epochs = 20\n    lr = 5e-5\n    weight_decay = 0.01\n    n_folds = 5\n    seed = 42\n    \n    dropout = 0.3\n    use_ndvi = True\n    mean_shift_alpha = 0.3\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndef seed_everything(seed):\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_everything(CFG.seed)"},{"cell_type":"markdown","id":"data_header","metadata":{},"source":"## 2. Data Loading & Wide Pivoting"},{"cell_type":"code","execution_count":null,"id":"load_data","metadata":{},"outputs":[],"source":"def load_and_pivot(df_path):\n    df = pd.read_csv(df_path)\n    # Extract unique image_id\n    if 'image_path' in df.columns:\n        df['image_id'] = df['image_path'].str.split('/').str[-1].str.replace('.jpg', '')\n    else:\n        df['image_id'] = df['sample_id'].str.split('__').str[0]\n        \n    if 'target' in df.columns:\n        # Pivot targets\n        target_df = df.pivot(index='image_id', columns='target_name', values='target').reset_index()\n        # Merge back metadata\n        metadata_cols = [c for c in df.columns if c not in ['target_name', 'target', 'sample_id']]\n        metadata_df = df[metadata_cols].drop_duplicates(subset='image_id')\n        return metadata_df.merge(target_df, on='image_id')\n    else:\n        # Test set\n        metadata_cols = [c for c in df.columns if c not in ['target_name', 'sample_id']]\n        metadata_df = df[metadata_cols].drop_duplicates(subset='image_id')\n        return metadata_df\n\ntrain_df = load_and_pivot(CFG.data_dir / \"train.csv\")\ntest_df = load_and_pivot(CFG.data_dir / \"test.csv\")\n\nprint(f\"Pivoted Training: {len(train_df)}, Test: {len(test_df)}\")\nprint(f\"Columns: {train_df.columns.tolist()}\")\n\nEPS = 1e-6\nfor col in CFG.base_targets:\n    train_df[f\"{col}_log\"] = np.log1p(train_df[col] + EPS)\n\nTRAIN_MEANS = {col: train_df[col].mean() for col in CFG.all_targets}\nprint(\"Training stats stored\")"},{"cell_type":"markdown","id":"dataset_header","metadata":{},"source":"## 3. Dataset & Model"},{"cell_type":"code","execution_count":null,"id":"dataset_class","metadata":{},"outputs":[],"source":"class BiomassDataset(Dataset):\n    def __init__(self, df, data_dir, transform=None, is_test=False):\n        self.df = df.reset_index(drop=True)\n        self.data_dir = data_dir\n        self.transform = transform\n        self.is_test = is_test\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 = self.data_dir / row['image_path']\n        image = cv2.imread(str(img_path))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            image = self.transform(image=image)['image']\n        \n        ndvi = float(row[CFG.ndvi_col]) if pd.notna(row[CFG.ndvi_col]) else 0.0\n        if self.is_test:\n            return image, torch.tensor([ndvi], dtype=torch.float32), str(row['image_id'])\n        \n        target = torch.tensor([row[f\"{t}_log\"] for t in CFG.base_targets], dtype=torch.float32)\n        return image, torch.tensor([ndvi], dtype=torch.float32), target\n\nclass BiomassRegressorAdvanced(nn.Module):\n    def __init__(self, dinov2_path, layers, hidden_size=1536):\n        super().__init__()\n        self.backbone = AutoModel.from_pretrained(dinov2_path)\n        self.layers = layers\n        for param in self.backbone.parameters():\n            param.requires_grad = False\n        \n        self.ndvi_embed = nn.Sequential(nn.Linear(1, 64), nn.ReLU(), nn.Linear(64, 64))\n        self.head = nn.Sequential(\n            nn.Linear(hidden_size + 64, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(0.3),\n            nn.Linear(512, 256), nn.LayerNorm(256), nn.GELU(), nn.Dropout(0.3),\n            nn.Linear(256, 3)\n        )\n        self.gdm_head = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 1))\n    \n    def forward(self, x, ndvi):\n        outputs = self.backbone(x, output_hidden_states=True)\n        features = torch.stack([outputs.hidden_states[l][:, 0, :] for l in self.layers], dim=0).mean(dim=0)\n        combined = torch.cat([features, self.ndvi_embed(ndvi)], dim=-1)\n        return self.head(combined)\n    \n    def derive_all_targets(self, base_preds_log, ndvi):\n        base_preds = torch.expm1(base_preds_log) - EPS\n        base_preds = torch.clamp(base_preds, min=0)\n        clover, dead, green = base_preds[:, 0:1], base_preds[:, 1:2], base_preds[:, 2:3]\n        total = clover + dead + green\n        gdm = self.gdm_head(torch.cat([green, ndvi], dim=-1))\n        return torch.cat([total, torch.clamp(gdm, min=0), green, dead, clover], dim=-1)"},{"cell_type":"markdown","id":"training_header","metadata":{},"source":"## 4. Training Loop"},{"cell_type":"code","execution_count":null,"id":"training","metadata":{},"outputs":[],"source":"def get_transforms(train=True):\n    if train:\n        return A.Compose([\n            A.RandomResizedCrop(CFG.image_size, CFG.image_size, scale=(0.8, 1.0)),\n            A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2()\n        ])\n    return A.Compose([\n        A.Resize(CFG.image_size, CFG.image_size),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2()\n    ])\n\ndef train_fold(fold, train_df):\n    skf = StratifiedKFold(n_splits=CFG.n_folds, shuffle=True, random_state=CFG.seed)\n    train_idx, val_idx = list(skf.split(train_df, train_df[CFG.state_col]))[fold]\n    \n    train_loader = DataLoader(BiomassDataset(train_df.iloc[train_idx], CFG.data_dir, transform=get_transforms()), batch_size=CFG.batch_size, shuffle=True, num_workers=4)\n    val_loader = DataLoader(BiomassDataset(train_df.iloc[val_idx], CFG.data_dir, transform=get_transforms(False)), batch_size=CFG.batch_size, shuffle=False, num_workers=4)\n    \n    model = BiomassRegressorAdvanced(CFG.dinov2_path, CFG.dinov2_layers).to(CFG.device)\n    criterion = nn.HuberLoss()\n    optimizer = AdamW(model.parameters(), lr=CFG.lr)\n    scheduler = OneCycleLR(optimizer, max_lr=CFG.lr, total_steps=len(train_loader)*CFG.epochs//CFG.accumulation_steps)\n    \n    best_r2 = -1\n    for epoch in range(CFG.epochs):\n        model.train()\n        for i, (imgs, ndvi, targets) in enumerate(tqdm(train_loader, desc=f\"Fold {fold} Ep {epoch}\")):\n            imgs, ndvi, targets = imgs.to(CFG.device), ndvi.to(CFG.device), targets.to(CFG.device)\n            preds = model(imgs, ndvi)\n            loss = criterion(preds, targets) / CFG.accumulation_steps\n            loss.backward()\n            if (i+1) % CFG.accumulation_steps == 0: optimizer.step(); scheduler.step(); optimizer.zero_grad()\n            \n        model.eval()\n        all_p, all_t = [], []\n        with torch.no_grad():\n            for imgs, ndvi, targets in val_loader:\n                imgs, ndvi = imgs.to(CFG.device), ndvi.to(CFG.device)\n                all_p.append(model(imgs, ndvi).cpu()); all_t.append(targets)\n        val_r2 = r2_score(torch.cat(all_t).numpy(), torch.cat(all_p).numpy())\n        print(f\"Epoch {epoch} Val R2: {val_r2:.4f}\")\n        if val_r2 > best_r2:\n            best_r2 = val_r2; torch.save(model.state_dict(), f\"best_fold{fold}.pt\")\n    return best_r2\n\nscores = [train_fold(0, train_df)] # Train 1 fold for now to get results quickly\nprint(f\"Mean R2: {np.mean(scores):.4f}\")\n\nimport shutil; shutil.copy(\"best_fold0.pt\", \"best_model.pt\")\nimport json\nwith open('training_stats.json', 'w') as f: json.dump({'train_means': TRAIN_MEANS}, f)"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.12"}},"nbformat":4,"nbformat_minor":5}