{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13693738,"sourceType":"datasetVersion","datasetId":8709989},{"sourceId":13852231,"sourceType":"datasetVersion","datasetId":8823915},{"sourceId":13969945,"sourceType":"datasetVersion","datasetId":8905857},{"sourceId":14549047,"sourceType":"datasetVersion","datasetId":9292546},{"sourceId":14556286,"sourceType":"datasetVersion","datasetId":9297345},{"sourceId":14698350,"sourceType":"datasetVersion","datasetId":9389806}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import glob\nimport os\nimport random\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset, TensorDataset\n\nimport torchvision.models as tv\nimport torchvision.transforms as T\nimport timm\n\nfrom sklearn.model_selection import StratifiedKFold, train_test_split","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:24.512384Z","iopub.execute_input":"2026-02-02T11:47:24.512983Z","iopub.status.idle":"2026-02-02T11:47:35.479456Z","shell.execute_reply.started":"2026-02-02T11:47:24.512955Z","shell.execute_reply":"2026-02-02T11:47:35.478652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# reproducibility package\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed(SEED)\ntorch.cuda.manual_seed_all(SEED)\n\ntorch.backends.cudnn.deterministic = True\ntorch.backends.cudnn.benchmark = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:40.467107Z","iopub.execute_input":"2026-02-02T11:47:40.467311Z","iopub.status.idle":"2026-02-02T11:47:40.479584Z","shell.execute_reply.started":"2026-02-02T11:47:40.467294Z","shell.execute_reply":"2026-02-02T11:47:40.478956Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset creation","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/csiro-biomass/train.csv')\n\ntrain","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:40.480900Z","iopub.execute_input":"2026-02-02T11:47:40.481114Z","iopub.status.idle":"2026-02-02T11:47:40.733401Z","shell.execute_reply.started":"2026-02-02T11:47:40.481097Z","shell.execute_reply":"2026-02-02T11:47:40.732669Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train.info()\nprint('\\n'f'Missing values: {train.isna().sum().sum()}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:47.104009Z","iopub.execute_input":"2026-02-02T11:47:47.104288Z","iopub.status.idle":"2026-02-02T11:47:47.129672Z","shell.execute_reply.started":"2026-02-02T11:47:47.104266Z","shell.execute_reply":"2026-02-02T11:47:47.128952Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train.describe()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:49.123567Z","iopub.execute_input":"2026-02-02T11:47:49.123840Z","iopub.status.idle":"2026-02-02T11:47:49.142450Z","shell.execute_reply.started":"2026-02-02T11:47:49.123819Z","shell.execute_reply":"2026-02-02T11:47:49.141855Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train.describe(include='object')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T07:51:23.530952Z","iopub.execute_input":"2026-01-26T07:51:23.531626Z","iopub.status.idle":"2026-01-26T07:51:23.556691Z","shell.execute_reply.started":"2026-01-26T07:51:23.531591Z","shell.execute_reply":"2026-01-26T07:51:23.555835Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"*Checked the data: no missing values or constant features.*","metadata":{}},{"cell_type":"code","source":"# building a dataset with unique images and their five target values\ndf_image = (\n    train.pivot_table(index='image_path',\n                   columns='target_name',\n                   values='target')\n    .reset_index()\n    .rename_axis(None, axis=1)  \n)\n\ndf_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:51.280919Z","iopub.execute_input":"2026-02-02T11:47:51.281442Z","iopub.status.idle":"2026-02-02T11:47:51.313776Z","shell.execute_reply.started":"2026-02-02T11:47:51.281415Z","shell.execute_reply":"2026-02-02T11:47:51.313037Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"meta = train[[\"image_path\", \"State\", \"Sampling_Date\"]].drop_duplicates()\ndf_image = df_image.merge(meta, on=\"image_path\", how=\"left\")\n\ndf_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:53.391395Z","iopub.execute_input":"2026-02-02T11:47:53.392352Z","iopub.status.idle":"2026-02-02T11:47:53.409936Z","shell.execute_reply.started":"2026-02-02T11:47:53.392301Z","shell.execute_reply":"2026-02-02T11:47:53.409230Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test = pd.read_csv('/kaggle/input/csiro-biomass/test.csv')\ntest['target'] = 0\n\ndf_image_test = (\n    test.pivot_table(index='image_path',\n                   columns='target_name',\n                   values='target')\n    .reset_index()\n    .rename_axis(None, axis=1)  \n)\n\ndf_image_test","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:55.402281Z","iopub.execute_input":"2026-02-02T11:47:55.402577Z","iopub.status.idle":"2026-02-02T11:47:55.425036Z","shell.execute_reply.started":"2026-02-02T11:47:55.402555Z","shell.execute_reply":"2026-02-02T11:47:55.424423Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualization of images before and after transformations","metadata":{}},{"cell_type":"code","source":"X = df_image[\"image_path\"]\ny = df_image.drop([\"image_path\", \"State\", \"Sampling_Date\"], axis=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:47:57.104868Z","iopub.execute_input":"2026-02-02T11:47:57.105685Z","iopub.status.idle":"2026-02-02T11:47:57.109926Z","shell.execute_reply.started":"2026-02-02T11:47:57.105654Z","shell.execute_reply":"2026-02-02T11:47:57.109289Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"*I create two transformations for further ensemble training on two models with two different resolutions.*","metadata":{}},{"cell_type":"markdown","source":"# Building the network architecture","metadata":{}},{"cell_type":"code","source":"IMG_SIZE_288 = 384\n\ntrain_tf_288 = T.Compose([\n    T.RandomHorizontalFlip(p=0.5),\n    T.RandomVerticalFlip(p=0.5),\n    T.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.15, hue=0.1),  \n    T.Resize((IMG_SIZE_288, IMG_SIZE_288)),\n    T.ToTensor(),  \n    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])\n])\n\nval_tf_288 = T.Compose([\n    T.Resize((IMG_SIZE_288, IMG_SIZE_288)),\n    T.ToTensor(),\n    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T12:08:38.489117Z","iopub.execute_input":"2026-02-02T12:08:38.489354Z","iopub.status.idle":"2026-02-02T12:08:38.494998Z","shell.execute_reply.started":"2026-02-02T12:08:38.489339Z","shell.execute_reply":"2026-02-02T12:08:38.494089Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# set the dataset root\nDATA_ROOT = Path(\"/kaggle/input/csiro-biomass/train\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:02.093442Z","iopub.execute_input":"2026-02-02T11:48:02.093964Z","iopub.status.idle":"2026-02-02T11:48:02.097180Z","shell.execute_reply.started":"2026-02-02T11:48:02.093943Z","shell.execute_reply":"2026-02-02T11:48:02.096382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PastureDataset(Dataset):\n    def __init__(self, paths, targets=None, transform=None, root=None):\n        \"\"\"\n        paths: pd.Series with relative or absolute image paths\n        targets: pd.DataFrame or None (for test); columns = 5 targets\n        transform: torchvision.transforms.Compose\n        root: optional common root for relative paths\n        \"\"\"\n        # pandas.Series of length N containing image paths\n        self.paths = paths.reset_index(drop=True)\n        self.has_targets = targets is not None\n        self.transform = transform\n        self.root = Path(root) if root is not None else None\n\n        if self.has_targets:\n            # keep the original column order for consistency\n            self.target_cols = list(targets.columns)\n            # torch.FloatTensor of shape [N, 5] (N rows, 5 target values)\n            self.targets = torch.tensor(targets.values, dtype=torch.float32)\n        else:\n            self.target_cols = None\n            self.targets = None\n\n    def __len__(self):\n        return len(self.paths)\n\n    def _resolve_path(self, p):\n        # return absolute path if root is provided and path is relative\n        if self.root is not None and not os.path.isabs(p):\n            return os.path.join(self.root, p)\n        return p\n\n    def __getitem__(self, idx):\n        # image file name, e.g. \"ID1011485656.jpg\" or \"train/ID1011485656.jpg\"\n        fname = Path(self.paths.iloc[idx]).name\n        # absolute path = root + file name\n        img_path = self.root / fname\n\n        # open the image and convert to rgb\n        img = Image.open(img_path).convert(\"RGB\")\n\n        # apply transforms\n        if self.transform:\n            img = self.transform(img)\n\n        if self.has_targets:\n            # target vector [5]\n            target = self.targets[idx]\n            return img, target\n        else:\n            # return image path for test mode\n            return img, str(img_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:04.382002Z","iopub.execute_input":"2026-02-02T11:48:04.382265Z","iopub.status.idle":"2026-02-02T11:48:04.389195Z","shell.execute_reply.started":"2026-02-02T11:48:04.382249Z","shell.execute_reply":"2026-02-02T11:48:04.388503Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"*Create the competition metric for evaluating the model architecture*","metadata":{}},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# metric weights in the required order\nWEIGHTS_MAP = {\n    'Dry_Green_g': 0.1,\n    'Dry_Dead_g':  0.1,\n    'Dry_Clover_g':0.1,\n    'GDM_g':       0.2,\n    'Dry_Total_g': 0.5,\n}\n\n# rearrange weights to match y columns order\nWEIGHTS = torch.tensor([WEIGHTS_MAP[c] for c in y.columns],\n                       dtype=torch.float32, device=device)\n\n# r2 for each target\ndef r2_per_target(y_true: torch.Tensor, y_pred: torch.Tensor, eps=1e-9):\n    # compute mean value of each target: y_true [n,5] → mean [1,5]\n    mean = y_true.mean(dim=0, keepdim=True)\n    # mse residual sum for each target\n    ss_res = ((y_true - y_pred) ** 2).sum(dim=0)\n    # total variance (mse vs mean) for each target\n    ss_tot = ((y_true - mean) ** 2).sum(dim=0) + eps\n    # r2 value for each target\n    return 1.0 - ss_res / ss_tot\n\n# weighted r2\ndef weighted_r2(y_true: torch.Tensor, y_pred: torch.Tensor, weights: torch.Tensor):\n    r2_vec = r2_per_target(y_true, y_pred)\n    w = weights.to(r2_vec.device, dtype=r2_vec.dtype)\n    w = w / (w.sum() + 1e-9)\n    return (r2_vec * w).sum(), r2_vec","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:06.705193Z","iopub.execute_input":"2026-02-02T11:48:06.705466Z","iopub.status.idle":"2026-02-02T11:48:06.925336Z","shell.execute_reply.started":"2026-02-02T11:48:06.705447Z","shell.execute_reply":"2026-02-02T11:48:06.924781Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BATCH_SIZE = 32\nVAL_BATCH_SIZE = 32\nNUM_FOLDS = 5\nNUM_TARGETS = 5\nskf = StratifiedKFold(n_splits=NUM_FOLDS, shuffle=True, random_state=839425)\n\noof_preds = torch.zeros((len(y), y.shape[1]), dtype=torch.float32)\noof_index = np.arange(len(y))\n\nfold_wr2_list = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:13.817000Z","iopub.execute_input":"2026-02-02T11:48:13.817575Z","iopub.status.idle":"2026-02-02T11:48:13.824472Z","shell.execute_reply.started":"2026-02-02T11:48:13.817546Z","shell.execute_reply":"2026-02-02T11:48:13.823651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DINO_WEIGHTS_PATH = \"/kaggle/input/vit-small-patch16-224-dino/vit_small_patch16_224_dino.pth\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:15.773056Z","iopub.execute_input":"2026-02-02T11:48:15.773618Z","iopub.status.idle":"2026-02-02T11:48:15.776684Z","shell.execute_reply.started":"2026-02-02T11:48:15.773598Z","shell.execute_reply":"2026-02-02T11:48:15.776042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DinoViTRegressor(nn.Module):\n    def __init__(self, model_name: str, num_targets: int, img_size: int, drop: float = 0.1, hidden: int = 512):\n        super().__init__()\n\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=False,\n            num_classes=0,\n            img_size=img_size,\n        )\n\n        state_dict = torch.load(DINO_WEIGHTS_PATH, map_location=\"cpu\")\n        state_dict.pop(\"pos_embed\", None)\n        self.backbone.load_state_dict(state_dict, strict=False)\n\n        feat_dim = self.backbone.num_features\n\n        def make_head():\n            return nn.Sequential(\n                nn.Linear(feat_dim, hidden),\n                nn.GELU(),\n                nn.Dropout(drop),\n                nn.Linear(hidden, hidden),\n                nn.GELU(),\n                nn.Dropout(drop),\n                nn.Linear(hidden, 1),\n            )\n            \n        self.heads = nn.ModuleList([make_head() for _ in range(num_targets)])\n\n    def forward(self, x):\n        feats = self.backbone(x)                      \n        outs = [h(feats) for h in self.heads]        \n        return torch.cat(outs, dim=1)                 ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:17.562185Z","iopub.execute_input":"2026-02-02T11:48:17.562761Z","iopub.status.idle":"2026-02-02T11:48:17.568647Z","shell.execute_reply.started":"2026-02-02T11:48:17.562738Z","shell.execute_reply":"2026-02-02T11:48:17.568036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DINO_MODEL_NAME = \"vit_small_patch16_224.dino\"\n\ndef build_dino_vit_regressor(\n    num_targets: int,\n    model_name: str = DINO_MODEL_NAME,\n    img_size=IMG_SIZE_288,\n    drop: float = 0.1,\n):\n    return DinoViTRegressor(\n        model_name=model_name,\n        num_targets=num_targets,\n        img_size=img_size,\n        drop=drop,\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:19.485002Z","iopub.execute_input":"2026-02-02T11:48:19.485567Z","iopub.status.idle":"2026-02-02T11:48:19.489362Z","shell.execute_reply.started":"2026-02-02T11:48:19.485545Z","shell.execute_reply":"2026-02-02T11:48:19.488746Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# create directory for saving checkpoints if it does not exist yet\nCKPT_DIR = \"checkpoints\"\nos.makedirs(CKPT_DIR, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:21.669210Z","iopub.execute_input":"2026-02-02T11:48:21.669784Z","iopub.status.idle":"2026-02-02T11:48:21.673143Z","shell.execute_reply.started":"2026-02-02T11:48:21.669764Z","shell.execute_reply":"2026-02-02T11:48:21.672514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WeightedSmoothL1(nn.Module):\n    def __init__(self, w: torch.Tensor, beta: float = 1.0):\n        super().__init__()\n        self.register_buffer(\"w\", w / (w.sum() + 1e-9))\n        self.beta = beta\n\n    def forward(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:\n        # pred/true: (B,5)\n        diff = torch.abs(pred - true)\n        loss = torch.where(diff < self.beta, 0.5 * (diff ** 2) / self.beta, diff - 0.5 * self.beta)  # (B,5)\n        return (loss * self.w).mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:48:22.875413Z","iopub.execute_input":"2026-02-02T11:48:22.875969Z","iopub.status.idle":"2026-02-02T11:48:22.880681Z","shell.execute_reply.started":"2026-02-02T11:48:22.875948Z","shell.execute_reply":"2026-02-02T11:48:22.880031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"w = WEIGHTS.to(device)  \ncriterion = WeightedSmoothL1(w, beta=1.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T11:49:01.617041Z","iopub.execute_input":"2026-02-02T11:49:01.617771Z","iopub.status.idle":"2026-02-02T11:49:01.622146Z","shell.execute_reply.started":"2026-02-02T11:49:01.617746Z","shell.execute_reply":"2026-02-02T11:49:01.621434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\nfor fold, (tr_idx, va_idx) in enumerate(skf.split(oof_index, df_image[\"State\"])):\n    print(f\"\\n===== FOLD {fold + 1}/{NUM_FOLDS} =====\")\n\n    X_tr, X_va = X.iloc[tr_idx], X.iloc[va_idx]\n    y_tr, y_va = y.iloc[tr_idx], y.iloc[va_idx]\n\n    w = WEIGHTS.to(device)\n    w = w / (w.sum() + 1e-9)\n    \n    criterion = WeightedSmoothL1(w, beta=1.0)\n\n    train_ds = PastureDataset(X_tr, y_tr, transform=train_tf_288, root=DATA_ROOT)\n    val_ds   = PastureDataset(X_va, y_va, transform=val_tf_288, root=DATA_ROOT)\n\n    g = torch.Generator()\n    g.manual_seed(SEED)\n\n    train_loader = DataLoader(\n        train_ds,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=True,\n        prefetch_factor=2,\n        generator=g\n    )\n    val_loader = DataLoader(\n        val_ds,\n        batch_size=VAL_BATCH_SIZE,\n        shuffle=False,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=True\n    )\n\n    # ----- model -----\n    model = build_dino_vit_regressor(num_targets=y.shape[1], img_size=IMG_SIZE_288).to(device)\n\n    # ----- optimizer: different lr for backbone/head -----\n    backbone_params = [p for n, p in model.named_parameters() if not n.startswith(\"heads.\")]\n    head_params     = [p for n, p in model.named_parameters() if n.startswith(\"heads.\")]\n\n\n    optimizer = torch.optim.AdamW(\n        [\n            {\"params\": backbone_params, \"lr\": 1e-5, \"weight_decay\": 1e-4},  \n            {\"params\": head_params,     \"lr\": 5e-4, \"weight_decay\": 1e-4},  \n        ]\n    )\n\n    epochs = 35\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n\n    # ----- early stopping -----\n    patience = 8\n    best_wr2 = -1e9\n    bad = 0\n    best_state = None\n\n    for epoch in range(1, epochs + 1):\n        # ---- train ----\n        model.train()\n        tr_running, tr_batches = 0.0, 0\n\n        for xb, yb in train_loader:\n            xb = xb.to(device, non_blocking=True)\n            yb = yb.to(device, non_blocking=True)\n\n            optimizer.zero_grad(set_to_none=True)\n            preds = model(xb)\n\n            loss = criterion(preds, yb)\n            \n            loss.backward()\n            optimizer.step()\n\n            tr_running += loss.item()\n            tr_batches += 1\n\n        scheduler.step()\n\n        # ---- val ----\n        model.eval()\n        val_running, val_batches = 0.0, 0\n        preds_all, ys_all = [], []\n\n        with torch.no_grad():\n            for xb, yb in val_loader:\n                xb = xb.to(device, non_blocking=True)\n                yb = yb.to(device, non_blocking=True)\n\n                preds = model(xb)\n                \n                vloss = criterion(preds, yb)\n\n                val_running += vloss.item()\n                val_batches += 1\n                preds_all.append(preds.detach().cpu())\n                ys_all.append(yb.detach().cpu())\n\n        val_loss = val_running / max(val_batches, 1)\n\n        y_true = torch.cat(ys_all, dim=0)\n        y_pred = torch.cat(preds_all, dim=0)\n\n        wr2, r2_vec = weighted_r2(y_true, y_pred, WEIGHTS)\n\n        print(\n            f\"epoch {epoch:02d} | train_loss {tr_running/max(tr_batches,1):.4f} \"\n            f\"| val_loss {val_loss:.4f} | WR2 {wr2.item():.4f} \"\n            f\"| R2 {np.round(r2_vec.detach().cpu().numpy(), 4)}\"\n        )\n\n        if wr2.item() > best_wr2:\n            best_wr2 = wr2.item()\n            bad = 0\n            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}\n        else:\n            bad += 1\n            if bad >= patience:\n                print(f\"early stop on epoch {epoch}\")\n                break\n\n    fold_wr2_list.append(best_wr2)\n\n    # ---- restore best + save checkpoint ----\n    if best_state is not None:\n        model.load_state_dict(best_state, strict=True)\n\n        torch.save(\n            {\n                \"state_dict\": best_state,\n                \"fold\": fold,\n                \"wr2\": best_wr2,\n                \"target_cols\": list(y.columns),\n                \"model_name\": DINO_MODEL_NAME,\n            },\n            f\"{CKPT_DIR}/dino_vit_fold{fold}_wr2_{best_wr2:.4f}.pt\"\n        )\n\n    # ---- OOF predictions ----\n    model.eval()\n    fold_preds = []\n    with torch.no_grad():\n        for xb, _ in val_loader:\n            xb = xb.to(device, non_blocking=True)\n            fold_preds.append(model(xb).detach().cpu())\n    oof_preds[va_idx] = torch.cat(fold_preds, dim=0)[:len(va_idx)]\n\n    torch.save(best_state, f\"best_fold{fold}.pt\")\n","metadata":{"execution":{"iopub.status.busy":"2026-02-02T12:08:44.838052Z","iopub.execute_input":"2026-02-02T12:08:44.838741Z","iopub.status.idle":"2026-02-02T13:46:59.893420Z","shell.execute_reply.started":"2026-02-02T12:08:44.838707Z","shell.execute_reply":"2026-02-02T13:46:59.892300Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# final oof evaluation (on the full train set)\ny_true_all = torch.tensor(y.values, dtype=torch.float32)\nwr2_all, r2_vec_all = weighted_r2(y_true_all.to(device), oof_preds.to(device), WEIGHTS)\nprint(\"\\nOOF WR2:\", wr2_all.item(), \"| OOF R2 per target:\", r2_vec_all.detach().cpu().numpy())\n\n# keep OOF predictions for VIT (for ensemble OOF check)\noof_preds_convext = oof_preds.clone().cpu()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T13:47:13.756172Z","iopub.execute_input":"2026-02-02T13:47:13.757128Z","iopub.status.idle":"2026-02-02T13:47:13.855760Z","shell.execute_reply.started":"2026-02-02T13:47:13.757093Z","shell.execute_reply":"2026-02-02T13:47:13.854950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"target_cols = list(y.columns)\n\n# SAVE OOF FOR POSTPROCESS\nnp.save(\"oof_targets.npy\", y_true_all)\nnp.save(\"target_cols.npy\", np.array(target_cols))\n\nnp.save(\"oof_preds_VIT.npy\", oof_preds)\n\nprint(\"OOF predictions and targets saved\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T10:40:13.932017Z","iopub.execute_input":"2026-02-01T10:40:13.932317Z","iopub.status.idle":"2026-02-01T10:40:13.939543Z","shell.execute_reply.started":"2026-02-01T10:40:13.932291Z","shell.execute_reply":"2026-02-01T10:40:13.938845Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Final prediction and submission","metadata":{}},{"cell_type":"code","source":"# creating a data loader for testing\nDATA_ROOT_TEST = Path(\"/kaggle/input/csiro-biomass/test\")\n\n# build three test loaders with the correct validation transforms\ntest_ds_288 = PastureDataset(df_image_test['image_path'], targets=None, transform=val_tf_288, root=DATA_ROOT_TEST)\n\nTEST_BS = 32\n\ntest_loader_288 = DataLoader(\n    test_ds_288,\n    batch_size=TEST_BS,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=True,\n    persistent_workers=True,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T10:40:32.570605Z","iopub.execute_input":"2026-02-01T10:40:32.571088Z","iopub.status.idle":"2026-02-01T10:40:32.575916Z","shell.execute_reply.started":"2026-02-01T10:40:32.571063Z","shell.execute_reply":"2026-02-01T10:40:32.575065Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# target order is the same as during training\ntarget_cols = list(y.columns)   \n\ndef _load_state_dict(ckpt_path: str):\n    state = torch.load(ckpt_path, map_location=\"cpu\")\n    if isinstance(state, dict) and \"state_dict\" in state:\n        return state[\"state_dict\"]\n    return state\n\n@torch.no_grad()\ndef predict_from_checkpoints(ckpt_paths, build_fn, loader, device):\n    \"\"\"Return [n_test, n_targets] averaged over folds for a given model family.\"\"\"\n    assert len(ckpt_paths) > 0, \"No checkpoints found.\"\n    fold_preds = []\n\n    for ckpt in ckpt_paths:\n        model = build_fn(num_targets=len(target_cols)).to(device)\n        sd = _load_state_dict(ckpt)\n        model.load_state_dict(sd, strict=True)\n        model.eval()\n\n        preds_parts = []\n        for xb in loader:\n            # PastureDataset(test) may return imgs or (imgs, meta). Keep only imgs.\n            if isinstance(xb, (list, tuple)):\n                xb = xb[0]\n            xb = xb.to(device, non_blocking=True)\n            out = model(xb).detach().cpu()\n            preds_parts.append(out)\n\n        fold_preds.append(torch.cat(preds_parts, dim=0))\n\n    # [num_folds, n_test, n_targets] -> [n_test, n_targets]\n    return torch.stack(fold_preds, dim=0).mean(dim=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T10:40:34.791535Z","iopub.execute_input":"2026-02-01T10:40:34.792052Z","iopub.status.idle":"2026-02-01T10:40:34.798850Z","shell.execute_reply.started":"2026-02-01T10:40:34.792025Z","shell.execute_reply":"2026-02-01T10:40:34.797950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# VIT predictions \nckpt_paths_VIT = sorted(glob.glob(\"checkpoints/dino_vit_fold*.pt\"))\nassert len(ckpt_paths_VIT) > 0, \"No VIT checkpoints found in checkpoints/ (expected dino_vit_fold*.pt)\"\n\npredsVIT = predict_from_checkpoints(\n    ckpt_paths_VIT,\n    build_fn=build_dino_vit_regressor,\n    loader=test_loader_288,\n    device=device,\n)\n\n# Ensemble blend \npreds_ens = predsVIT\npreds_ens.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T10:40:46.554889Z","iopub.execute_input":"2026-02-01T10:40:46.555615Z","iopub.status.idle":"2026-02-01T10:40:49.698386Z","shell.execute_reply.started":"2026-02-01T10:40:46.555587Z","shell.execute_reply":"2026-02-01T10:40:49.697608Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# we form a submission in the format sample_id,target\nrows = []\nimg_ids = [Path(p).stem for p in df_image_test['image_path']]\n\npreds_np = preds_ens.numpy()\n\nfor i, img_id in enumerate(img_ids):\n    for j, tname in enumerate(target_cols):\n        rows.append({\n            \"sample_id\": f\"{img_id}__{tname}\",\n            \"target\": float(preds_np[i, j]),\n        })\n\nsubmission = pd.DataFrame(rows, columns=[\"sample_id\", \"target\"])\nsubmission.to_csv(\"submission.csv\", index=False)\n\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T10:40:52.764183Z","iopub.execute_input":"2026-02-01T10:40:52.764532Z","iopub.status.idle":"2026-02-01T10:40:52.786465Z","shell.execute_reply.started":"2026-02-01T10:40:52.764499Z","shell.execute_reply":"2026-02-01T10:40:52.785800Z"}},"outputs":[],"execution_count":null}]}