{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport time\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.metrics import r2_score\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\ntry:\n    from torch.amp import autocast, GradScaler\nexcept ImportError:\n    from torch.cuda.amp import autocast, GradScaler\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\n\n# =============================== Configuration =================================\nclass CFG:\n    # --- Paths ---\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    OUTPUT_DIR = \"./output_weights\"\n    \n    # --- Model Architecture ---\n    MODEL_NAME = \"vit_base_patch14_dinov2\"\n    GRID = (2, 2)\n    VARIANT = \"tiled_film\"\n    \n    # --- YOUR SPECIFIC SETTINGS ---\n    IMAGE_SIZE = 518      # Native DINOv2 resolution\n    EPOCHS = 15           # Reduced to 15 as requested\n    UNFREEZE_EPOCH = 5    # Backbone unfreezes at Epoch 5 (0-4 are frozen)\n    \n    # --- Memory & Optimization ---\n    BATCH_SIZE = 2        # Small batch to fit 518px on GPU\n    ACCUM_ITER = 8        # Effective Batch Size = 16\n    \n    LR_HEAD = 1e-3        # High LR for heads\n    LR_BACKBONE = 8e-6    # Slightly higher than before since epochs are fewer\n    \n    NUM_WORKERS = 4\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    \n    TARGET_COLS = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n    TARGET_WEIGHTS = torch.tensor([0.1, 0.1, 0.1, 0.2, 0.5])\n\nif not os.path.exists(CFG.OUTPUT_DIR): \n    os.makedirs(CFG.OUTPUT_DIR)\n\n# =============================== Dataset =================================\nclass TrainBiomassDataset(Dataset):\n    def __init__(self, df, transform, image_dir):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.image_dir = image_dir\n        self.paths = self.df[\"image_path\"].values\n        self.targets = self.df[CFG.TARGET_COLS].values.astype(np.float32)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.image_dir, os.path.basename(str(self.paths[idx])))\n        img = cv2.imread(path)\n        \n        if img is None: \n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else: \n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # Split Wide Image -> Left & Right\n        h, w, _ = img.shape\n        mid = w // 2\n        left_t = self.transform(image=img[:, :mid])[\"image\"]\n        right_t = self.transform(image=img[:, mid:])[\"image\"]\n        return left_t, right_t, torch.tensor(self.targets[idx])\n\ndef get_transforms(data='train'):\n    ops = [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]),\n           ToTensorV2()]\n    if data == 'train':\n        ops = [\n            A.HorizontalFlip(p=0.5), \n            A.VerticalFlip(p=0.5), \n            A.Rotate(limit=15, p=0.5),\n            A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, p=0.5),\n            A.CoarseDropout(max_holes=4, max_height=32, max_width=32, p=0.3)\n        ] + ops\n    return A.Compose(ops)\n\n# =============================== Model =================================\nclass FiLM(nn.Module):\n    def __init__(self, in_dim):\n        super().__init__()\n        self.mlp = nn.Sequential(\n            nn.Linear(in_dim, in_dim//2), nn.ReLU(True), nn.Linear(in_dim//2, in_dim*2)\n        )\n    def forward(self, x):\n        g, b = torch.chunk(self.mlp(x), 2, dim=1)\n        return g, b\n\nclass HybridTwoStreamModel(nn.Module):\n    def __init__(self, backbone_name, grid=(2,2), variant=\"tiled_film\"):\n        super().__init__()\n        print(f\"Loading {backbone_name}...\")\n        self.backbone = timm.create_model(backbone_name, pretrained=True, num_classes=0)\n        \n        # Disable checkpointing initially\n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(False)\n        \n        self.input_res = 518 if \"dinov2\" in backbone_name else 224\n        feat_dim = self.backbone.num_features\n        self.grid = grid\n        \n        # Freeze Backbone\n        for param in self.backbone.parameters():\n            param.requires_grad = False\n            \n        if \"film\" in variant:\n            self.film_left = FiLM(feat_dim)\n            self.film_right = FiLM(feat_dim)\n\n        hid = int((feat_dim * 2) * 0.25)\n        def head(): \n            return nn.Sequential(\n                nn.Linear(feat_dim * 2, hid), nn.ReLU(True), nn.Dropout(0.3), nn.Linear(hid, 1)\n            )\n        \n        self.head_green, self.head_clover, self.head_dead = head(), head(), head()\n        self.softplus = nn.Softplus()\n\n    def set_backbone_trainable(self, trainable=True):\n        if trainable:\n            if hasattr(self.backbone, 'set_grad_checkpointing'):\n                self.backbone.set_grad_checkpointing(True)\n                print(\"--> Gradient Checkpointing ENABLED\")\n        \n        for param in self.backbone.parameters():\n            param.requires_grad = trainable\n\n    def _batch_tiles(self, x):\n        B, C, H, W = x.shape\n        r, c = self.grid\n        crops = []\n        h_step, w_step = H // r, W // c\n        \n        for i in range(r):\n            for j in range(c):\n                tile = x[:, :, i*h_step:(i+1)*h_step, j*w_step:(j+1)*w_step]\n                if tile.shape[-1] != self.input_res:\n                    tile = F.interpolate(tile, size=(self.input_res, self.input_res), mode='bilinear', align_corners=False)\n                crops.append(tile)\n        \n        batch_input = torch.cat(crops, dim=0) \n        feats = self.backbone(batch_input)    \n        return feats.view(r*c, B, -1).permute(1, 0, 2)\n\n    def forward(self, xl, xr):\n        # Left\n        tiles_l = self._batch_tiles(xl)\n        if hasattr(self, 'film_left'):\n            g, b = self.film_left(tiles_l.mean(1))\n            tiles_l = tiles_l * (1 + g.unsqueeze(1)) + b.unsqueeze(1)\n        fl = tiles_l.mean(1)\n\n        # Right\n        tiles_r = self._batch_tiles(xr)\n        if hasattr(self, 'film_right'):\n            g, b = self.film_right(tiles_r.mean(1))\n            tiles_r = tiles_r * (1 + g.unsqueeze(1)) + b.unsqueeze(1)\n        fr = tiles_r.mean(1)\n\n        f = torch.cat([fl, fr], dim=1)\n        \n        gr = self.softplus(self.head_green(f))\n        cl = self.softplus(self.head_clover(f))\n        de = self.softplus(self.head_dead(f))\n        \n        return torch.cat([gr, de, cl, gr+cl, gr+cl+de], dim=1)\n\n# =============================== Engine =================================\ndef train_epoch(model, loader, opt, scaler, scheduler):\n    model.train()\n    losses = []\n    amp_device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n    pbar = tqdm(loader, desc=\"Train\", leave=False)\n    opt.zero_grad()\n    \n    for step, (xl, xr, y) in enumerate(pbar):\n        xl, xr, y = xl.to(CFG.DEVICE), xr.to(CFG.DEVICE), y.to(CFG.DEVICE)\n        \n        with autocast(device_type=amp_device, enabled=True):\n            preds = model(xl, xr)\n            loss = ((preds - y)**2 * CFG.TARGET_WEIGHTS.to(CFG.DEVICE)).mean()\n            loss = loss / CFG.ACCUM_ITER\n        \n        scaler.scale(loss).backward()\n        \n        if (step + 1) % CFG.ACCUM_ITER == 0:\n            scaler.step(opt)\n            scaler.update()\n            opt.zero_grad()\n            if scheduler is not None:\n                scheduler.step()\n        \n        losses.append(loss.item() * CFG.ACCUM_ITER)\n        pbar.set_postfix(loss=f\"{np.mean(losses[-15:]):.4f}\")\n    \n    # Process remaining gradients\n    if (step + 1) % CFG.ACCUM_ITER != 0:\n        scaler.step(opt)\n        scaler.update()\n        opt.zero_grad()\n\n    return np.mean(losses)\n\n@torch.no_grad()\ndef valid_epoch(model, loader):\n    model.eval()\n    preds, targets = [], []\n    for xl, xr, y in tqdm(loader, desc=\"Valid\", leave=False):\n        xl, xr = xl.to(CFG.DEVICE), xr.to(CFG.DEVICE)\n        preds.append(model(xl, xr).float().cpu().numpy())\n        targets.append(y.float().numpy())\n        \n    P, T = np.concatenate(preds), np.concatenate(targets)\n    scores = [r2_score(T[:, i], P[:, i]) for i in range(5)]\n    return np.average(scores, weights=CFG.TARGET_WEIGHTS.numpy())\n\n# =============================== Main Runner =================================\ndef run_fold(fold_idx):\n    print(f\"Reading Data from {CFG.TRAIN_CSV}\")\n    raw = pd.read_csv(CFG.TRAIN_CSV)\n    df = raw.pivot(index='image_path', columns='target_name', values='target').reset_index()\n    meta = raw[['image_path', 'Sampling_Date']].drop_duplicates(subset=['image_path'])\n    df = pd.merge(df, meta, on='image_path', how='left').fillna(0.0)\n    \n    gkf = GroupKFold(n_splits=5)\n    splits = list(gkf.split(df, df[CFG.TARGET_COLS], df['Sampling_Date']))\n    \n    if fold_idx >= 5:\n        print(f\"Error: Fold index {fold_idx} invalid.\")\n        return\n\n    t_idx, v_idx = splits[fold_idx]\n    \n    print(f\"\\n{'='*40}\")\n    print(f\" Training Fold {fold_idx + 1} / 5 \")\n    print(f\"{'='*40}\")\n    \n    # 1. Setup\n    train_ds = TrainBiomassDataset(df.iloc[t_idx], get_transforms('train'), CFG.TRAIN_IMAGE_DIR)\n    valid_ds = TrainBiomassDataset(df.iloc[v_idx], get_transforms('valid'), CFG.TRAIN_IMAGE_DIR)\n    \n    t_loader = DataLoader(train_ds, batch_size=CFG.BATCH_SIZE, shuffle=True, \n                          num_workers=CFG.NUM_WORKERS, pin_memory=True, drop_last=True)\n    v_loader = DataLoader(valid_ds, batch_size=CFG.BATCH_SIZE, shuffle=False, \n                          num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    \n    model = HybridTwoStreamModel(CFG.MODEL_NAME, CFG.GRID, CFG.VARIANT)\n    model.to(CFG.DEVICE)\n    \n    # 2. Initial Optimizer\n    params = [p for p in model.parameters() if p.requires_grad]\n    opt = optim.AdamW(params, lr=CFG.LR_HEAD)\n    scaler = GradScaler()\n    \n    steps_per_epoch = len(t_loader) // CFG.ACCUM_ITER\n    scheduler = CosineAnnealingWarmRestarts(opt, T_0=steps_per_epoch * 5)\n    \n    best_r2 = -np.inf\n    \n    # 3. Training Loop\n    for ep in range(CFG.EPOCHS):\n        # UNFREEZE CHECK\n        if ep == CFG.UNFREEZE_EPOCH:\n            print(f\"\\n!!! Unfreezing Backbone at Epoch {ep} !!!\")\n            del scheduler, scaler, opt\n            gc.collect()\n            torch.cuda.empty_cache()\n            time.sleep(1)\n            \n            model.set_backbone_trainable(True)\n            \n            opt = optim.AdamW([\n                {'params': model.backbone.parameters(), 'lr': CFG.LR_BACKBONE},\n                {'params': model.head_green.parameters(), 'lr': CFG.LR_HEAD},\n                {'params': model.head_clover.parameters(), 'lr': CFG.LR_HEAD},\n                {'params': model.head_dead.parameters(), 'lr': CFG.LR_HEAD},\n                {'params': model.film_left.parameters(), 'lr': CFG.LR_HEAD},\n                {'params': model.film_right.parameters(), 'lr': CFG.LR_HEAD},\n            ], weight_decay=1e-4)\n            \n            scaler = GradScaler()\n            scheduler = CosineAnnealingWarmRestarts(opt, T_0=steps_per_epoch * (CFG.EPOCHS - ep))\n\n        t0 = time.time()\n        loss = train_epoch(model, t_loader, opt, scaler, scheduler)\n        r2 = valid_epoch(model, v_loader)\n        \n        if r2 > best_r2:\n            best_r2 = r2\n            save_path = os.path.join(CFG.OUTPUT_DIR, f\"best_fold{fold_idx+1}.pth\")\n            torch.save(model.state_dict(), save_path)\n            print(f\"  --> Saved (R2: {r2:.4f})\")\n        \n        status = \"FROZEN\" if ep < CFG.UNFREEZE_EPOCH else \"UNFROZEN\"\n        print(f\"Ep {ep+1:02d} [{status}] | Loss: {loss:.3f} | R2: {r2:.4f} (Best: {best_r2:.4f})\")\n\nif __name__ == \"__main__\":\n    # !!! CONFIGURATION FOR KAGGLE EXECUTION !!!\n    # Change FOLD_TO_TRAIN to 0, 1, 2, 3, or 4\n    # Run 5 separate notebook sessions, each with a different number\n    FOLD_TO_TRAIN = 4 \n    run_fold(FOLD_TO_TRAIN)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}