{"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":4535,"sourceType":"modelInstanceVersion","modelInstanceId":3327,"modelId":986}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# CSIRO - Image2Biomass Prediction\n# HF DINOv2 -> embeddings -> regressor -> weighted R² -> submission.csv\n# ============================================================\n\nimport os, random\nfrom dataclasses import dataclass\nfrom typing import Optional, Tuple, List, Dict, Any\n\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom PIL import Image\nimport torchvision.transforms as T\n\nfrom transformers import AutoImageProcessor, AutoModel\n\nfrom PIL import Image, ImageFile\nImageFile.LOAD_TRUNCATED_IMAGES = True\n\n\n\n# ----------------------------\n# 0) Config\n# ----------------------------\n@dataclass\nclass CFG:\n    COMP_ROOT: str = \"/kaggle/input/csiro-biomass\"\n    DINO_ROOT: str = \"/kaggle/input/dinov2\"  # top-level input that contains /pytorch/...\n\n    TRAIN_CSV: str = \"train.csv\"\n    TEST_CSV: str = \"test.csv\"\n    TRAIN_IMG_DIR: str = \"train\"\n    TEST_IMG_DIR: str = \"test\"\n\n    IMG_SIZE: int = 518\n    BATCH_SIZE: int = 8\n    # NUM_WORKERS: int = 2\n    NUM_WORKERS: int = 0\n    SEED: int = 42\n    DEVICE: str = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Predict 3 and derive the other 2\n    TARGETS_3: Tuple[str, str, str] = (\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\")\n\n    # Metric order (competition)\n    TARGET_ORDER_5: Tuple[str, str, str, str, str] = (\n        \"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"\n    )\n    TARGET_WEIGHTS_5: Tuple[float, float, float, float, float] = (0.1, 0.1, 0.1, 0.2, 0.5)\n\n    # From preview: GDM ~ Dry_Green (NOT total)\n    GDM_EQUALS_TOTAL: bool = False\n\ncfg = CFG()\n\ndef collate_pil(batch):\n    # batch is a list of dicts: {\"image\": PIL, \"fname\": str, optional \"y\": tensor}\n    out = {\n        \"image\": [b[\"image\"] for b in batch],\n        \"fname\": [b[\"fname\"] for b in batch],\n    }\n    if \"y\" in batch[0]:\n        out[\"y\"] = torch.stack([b[\"y\"] for b in batch], dim=0)\n    return out\n\n\n# ----------------------------\n# 1) Repro\n# ----------------------------\ndef seed_everything(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = False\n    torch.backends.cudnn.benchmark = True\n\nseed_everything(cfg.SEED)\n\n\n# ----------------------------\n# 2) Locate HF DINOv2 folder and load it ONCE\n# ----------------------------\ndef find_hf_model_dir(root: str):\n    need = {\"config.json\", \"pytorch_model.bin\"}\n    for r, _, files in os.walk(root):\n        if need.issubset(set(files)):\n            return r\n    raise FileNotFoundError(f\"Couldn't find {need} under {root}\")\n\nHF_DIR = find_hf_model_dir(cfg.DINO_ROOT)\nprint(\"[OK] Found HF model dir:\", HF_DIR)\nprint(\"Files:\", [f for f in os.listdir(HF_DIR) if f.endswith((\".json\", \".bin\"))])\n\nprocessor = AutoImageProcessor.from_pretrained(HF_DIR)\ndinov2 = AutoModel.from_pretrained(HF_DIR).to(cfg.DEVICE)\ndinov2.eval()\nfor p in dinov2.parameters():\n    p.requires_grad = False\n\n# infer embedding dim once\nwith torch.no_grad():\n    dummy = torch.zeros(1, 3, cfg.IMG_SIZE, cfg.IMG_SIZE, device=cfg.DEVICE)\n    out = dinov2(pixel_values=dummy)\n    EMBED_DIM = out.last_hidden_state.shape[-1]\nprint(\"[OK] HF DINOv2 loaded. EMBED_DIM =\", EMBED_DIM)\n\n\n# ----------------------------\n# 3) Load CSVs\n# ----------------------------\ntrain_csv_path = os.path.join(cfg.COMP_ROOT, cfg.TRAIN_CSV)\ntest_csv_path  = os.path.join(cfg.COMP_ROOT, cfg.TEST_CSV)\ntrain_img_dir = os.path.join(cfg.COMP_ROOT, cfg.TRAIN_IMG_DIR)\ntest_img_dir  = os.path.join(cfg.COMP_ROOT, cfg.TEST_IMG_DIR)\n\ntrain_df = pd.read_csv(train_csv_path)\ntest_df  = pd.read_csv(test_csv_path)\n\nprint(\"train_df:\", train_df.shape, \"test_df:\", test_df.shape)\n\n\n# ----------------------------\n# 4) Pivot train long -> wide (one row per image_path)\n# ----------------------------\nrequired_cols = {\"image_path\", \"target_name\", \"target\"}\nmissing_req = required_cols - set(train_df.columns)\nassert not missing_req, f\"train.csv missing required columns: {missing_req}\"\n\nwide_df = (\n    train_df.pivot_table(\n        index=[\"image_path\"],\n        columns=\"target_name\",\n        values=\"target\",\n        aggfunc=\"mean\",\n    )\n    .reset_index()\n)\nwide_df.columns = [str(c) for c in wide_df.columns]\n\nprint(\"wide_df:\", wide_df.shape)\nprint(\"wide_df cols:\", wide_df.columns.tolist())\n\ndef _norm(s: str) -> str:\n    return str(s).strip().lower().replace(\" \", \"_\").replace(\"-\", \"_\")\n\ncols_norm = {c: _norm(c) for c in wide_df.columns}\n\ndef find_col_contains(tokens: List[str]) -> Optional[str]:\n    hits = []\n    for c, cn in cols_norm.items():\n        if all(t in cn for t in tokens):\n            hits.append(c)\n    hits.sort(key=lambda x: (len(cols_norm[x]), x))\n    return hits[0] if hits else None\n\ndry_clover_col = find_col_contains([\"dry\", \"clover\"])\ndry_dead_col   = find_col_contains([\"dry\", \"dead\"])\ndry_green_col  = find_col_contains([\"dry\", \"green\"])\nassert all([dry_clover_col, dry_dead_col, dry_green_col]), \"Could not find all 3 dry targets.\"\n\ncfg.TARGETS_3 = (dry_clover_col, dry_dead_col, dry_green_col)\nIMG_COL = \"image_path\"\nprint(\"Using TARGETS_3:\", cfg.TARGETS_3)\n\n# Drop incomplete labels BEFORE split/embedding\ntrain_df_wide = wide_df.dropna(subset=list(cfg.TARGETS_3)).reset_index(drop=True)\nprint(\"After dropna:\", train_df_wide.shape)\nassert len(train_df_wide) > 0, \"No rows left after dropna.\"\n\n\n# ----------------------------\n# 5) Metric / loss\n# ----------------------------\nTARGET_WEIGHTS_5 = torch.tensor(cfg.TARGET_WEIGHTS_5, dtype=torch.float32)\n\ndef build_5_from_3(y3: torch.Tensor, gdm_equals_total: bool = False) -> torch.Tensor:\n    # y3 order: [clover, dead, green]\n    clover = y3[:, 0:1]\n    dead   = y3[:, 1:2]\n    green  = y3[:, 2:3]\n    total = clover + dead + green\n    gdm = total if gdm_equals_total else green\n    return torch.cat([green, dead, clover, gdm, total], dim=1)\n\ndef weighted_r2_score(y_true5: torch.Tensor, y_pred5: torch.Tensor, weights5: torch.Tensor) -> torch.Tensor:\n    y_true5 = y_true5.float()\n    y_pred5 = y_pred5.float()\n    w = weights5.view(1, -1).float().to(y_true5.device)\n\n    w_sum = (y_true5.shape[0] * torch.sum(w)).clamp_min(1e-12)\n    y_bar = torch.sum(w * y_true5) / w_sum\n\n    sse = torch.sum(w * (y_true5 - y_pred5) ** 2)\n    sst = torch.sum(w * (y_true5 - y_bar) ** 2)\n\n    eps = 1e-12\n    if float(sst) <= eps:\n        return torch.tensor(1.0 if float(sse) <= eps else 0.0, device=y_true5.device)\n\n    return 1.0 - sse / (sst + eps)\n\nclass WeightedMSEOn5(nn.Module):\n    def __init__(self, weights5: torch.Tensor, gdm_equals_total: bool = False):\n        super().__init__()\n        self.register_buffer(\"weights5\", weights5.float())\n        self.gdm_equals_total = gdm_equals_total\n\n    def forward(self, pred3: torch.Tensor, true3: torch.Tensor) -> torch.Tensor:\n        pred5 = build_5_from_3(pred3, self.gdm_equals_total)\n        true5 = build_5_from_3(true3, self.gdm_equals_total)\n        w = self.weights5.to(pred5.device).view(1, -1)\n        return torch.mean(w * (pred5 - true5) ** 2)\n\n\n# ----------------------------\n# 6) Transforms + Dataset (PIL -> optional aug -> HF processor)\n# ----------------------------\ntrain_pil_tfms = T.Compose([\n    T.Resize((cfg.IMG_SIZE, cfg.IMG_SIZE)),\n    T.RandomHorizontalFlip(p=0.5),\n    T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.05, hue=0.02),\n])\n\nvalid_pil_tfms = T.Compose([\n    T.Resize((cfg.IMG_SIZE, cfg.IMG_SIZE)),\n])\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, img_dir: str, img_col: str, targets=None, pil_tfms=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.img_col = img_col\n        self.targets = targets\n        self.pil_tfms = pil_tfms\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx: int) -> Dict[str, Any]:\n        row = self.df.iloc[idx]\n        fname = str(row[self.img_col]).replace(\"\\\\\", \"/\")\n        path = os.path.join(cfg.COMP_ROOT, fname) if \"/\" in fname else os.path.join(self.img_dir, fname)\n\n        try:\n            img = Image.open(path).convert(\"RGB\")\n        except Exception:\n            # Never crash on hidden test: use a blank fallback image\n            img = Image.new(\"RGB\", (cfg.IMG_SIZE, cfg.IMG_SIZE), (0, 0, 0))\n\n        if self.pil_tfms is not None:\n            img = self.pil_tfms(img)\n\n        out = {\"image\": img, \"fname\": fname}\n        if self.targets is not None:\n            out[\"y\"] = torch.tensor(row[list(self.targets)].values.astype(np.float32))\n        return out\n\n\n# ----------------------------\n# 7) HF embedding extraction (the ONLY extractor)\n# ----------------------------\n@torch.no_grad()\ndef extract_embeddings_hf(dino: nn.Module, loader: DataLoader, device: str):\n    dino.eval()\n    feats_list, y_list, fname_list = [], [], []\n\n    for batch in loader:\n        imgs = batch[\"image\"]  # list of PIL images\n        inputs = processor(images=imgs, return_tensors=\"pt\").to(device)\n\n        out = dino(**inputs)\n        feats = out.last_hidden_state[:, 0, :]  # CLS token\n\n        feats_list.append(feats.detach().cpu())\n        fname_list += list(batch[\"fname\"])\n        if \"y\" in batch:\n            y_list.append(batch[\"y\"].detach().cpu())\n\n    X = torch.cat(feats_list, dim=0).float().numpy()\n    y = torch.cat(y_list, dim=0).float().numpy() if len(y_list) else None\n    return X, y, fname_list\n\n\n# ----------------------------\n# 8) Train/Val split, embeddings, regressor\n# ----------------------------\nfrom sklearn.model_selection import train_test_split\n\nidx = np.arange(len(train_df_wide))\ntr_idx, va_idx = train_test_split(idx, test_size=0.2, random_state=cfg.SEED)\n\ntr_df = train_df_wide.iloc[tr_idx].reset_index(drop=True)\nva_df = train_df_wide.iloc[va_idx].reset_index(drop=True)\n\ntr_ds = BiomassDataset(tr_df, train_img_dir, IMG_COL, targets=cfg.TARGETS_3, pil_tfms=train_pil_tfms)\nva_ds = BiomassDataset(va_df, train_img_dir, IMG_COL, targets=cfg.TARGETS_3, pil_tfms=valid_pil_tfms)\n\ntr_dl = DataLoader(\n    tr_ds,\n    batch_size=cfg.BATCH_SIZE,\n    shuffle=False,\n    num_workers=cfg.NUM_WORKERS,\n    pin_memory=(cfg.DEVICE == \"cuda\"),\n    collate_fn=collate_pil,\n)\n\nb = next(iter(tr_dl))\nprint(type(b[\"image\"]), type(b[\"image\"][0]), b[\"y\"].shape)\n\nva_dl = DataLoader(va_ds, batch_size=cfg.BATCH_SIZE, shuffle=False, num_workers=cfg.NUM_WORKERS, pin_memory=(cfg.DEVICE == \"cuda\"),\n    collate_fn=collate_pil,)\n\nX_tr, y_tr, _ = extract_embeddings_hf(dinov2, tr_dl, cfg.DEVICE)\nX_va, y_va, _ = extract_embeddings_hf(dinov2, va_dl, cfg.DEVICE)\n\nimport inspect\nprint(inspect.signature(BiomassDataset))\n\nassert np.isfinite(X_tr).all() and np.isfinite(X_va).all()\nassert np.isfinite(y_tr).all() and np.isfinite(y_va).all()\n\nprint(\"Embeddings:\", X_tr.shape, X_va.shape, \"Targets:\", y_tr.shape, y_va.shape)\n\nclass MLPRegressor(nn.Module):\n    def __init__(self, in_dim: int, hidden=512, out_dim=3, p=0.1):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.LayerNorm(in_dim),\n            nn.Linear(in_dim, hidden),\n            nn.GELU(),\n            nn.Dropout(p),\n            nn.Linear(hidden, out_dim),\n        )\n    def forward(self, x): return self.net(x)\n\ndef make_tensor_loader(X, y, batch_size=256, shuffle=True):\n    X_t = torch.tensor(X, dtype=torch.float32)\n    y_t = torch.tensor(y, dtype=torch.float32)\n    ds = torch.utils.data.TensorDataset(X_t, y_t)\n    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle)\n\nmodel = MLPRegressor(in_dim=X_tr.shape[1], hidden=512, out_dim=3, p=0.1).to(cfg.DEVICE)\nloss_fn = WeightedMSEOn5(TARGET_WEIGHTS_5, gdm_equals_total=cfg.GDM_EQUALS_TOTAL).to(cfg.DEVICE)\nopt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-3)\n\ntr2_dl = make_tensor_loader(X_tr, y_tr, batch_size=256, shuffle=True)\nva2_dl = make_tensor_loader(X_va, y_va, batch_size=256, shuffle=False)\n\n\n@torch.no_grad()\ndef eval_on_embeddings(model, loader, device):\n    model.eval()\n    ys, ps = [], []\n    for xb, yb in loader:\n        xb = xb.to(device)\n        yb = yb.to(device)\n        pb = model(xb)\n        ys.append(yb.detach().cpu())\n        ps.append(pb.detach().cpu())\n    y3 = torch.cat(ys, 0)\n    p3 = torch.cat(ps, 0)\n    y5 = build_5_from_3(y3, cfg.GDM_EQUALS_TOTAL)\n    p5 = build_5_from_3(p3, cfg.GDM_EQUALS_TOTAL)\n    r2 = weighted_r2_score(y5, p5, TARGET_WEIGHTS_5)\n    mse = torch.mean((p3 - y3) ** 2).item()\n    return float(r2.item()), mse\n\ndef train_epochs(model, tr_loader, va_loader, epochs=30):\n    best_r2 = -1e9\n    best_state = None\n    for ep in range(1, epochs+1):\n        model.train()\n        for xb, yb in tr_loader:\n            xb = xb.to(cfg.DEVICE)\n            yb = yb.to(cfg.DEVICE)\n            pred = model(xb)\n            loss = loss_fn(pred, yb)\n            opt.zero_grad(set_to_none=True)\n            loss.backward()\n            opt.step()\n\n        r2, mse = eval_on_embeddings(model, va_loader, cfg.DEVICE)\n        print(f\"Epoch {ep:02d} | val weighted R2: {r2:.5f} | val mse(3): {mse:.5f}\")\n\n        if r2 > best_r2:\n            best_r2 = r2\n            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}\n\n    if best_state is not None:\n        model.load_state_dict(best_state)\n    print(\"Best val weighted R2:\", best_r2)\n\ntrain_epochs(model, tr2_dl, va2_dl, epochs=30)\n\n\n# ----------------------------\n# 9) Fit on full clean training set\n# ----------------------------\nfull_ds = BiomassDataset(train_df_wide, train_img_dir, IMG_COL, targets=cfg.TARGETS_3, pil_tfms=valid_pil_tfms)\nfull_dl = DataLoader(full_ds, batch_size=cfg.BATCH_SIZE, shuffle=False,  pin_memory=(cfg.DEVICE == \"cuda\"),\n    collate_fn=collate_pil,\n)\nX_full, y_full, _ = extract_embeddings_hf(dinov2, full_dl, cfg.DEVICE)\n\nmodel = MLPRegressor(in_dim=X_full.shape[1], hidden=512, out_dim=3, p=0.1).to(cfg.DEVICE)\nopt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-3)\nfull_loader = make_tensor_loader(X_full, y_full, batch_size=256, shuffle=True)\n\nfor ep in range(1, 21):\n    model.train()\n    for xb, yb in full_loader:\n        xb = xb.to(cfg.DEVICE)\n        yb = yb.to(cfg.DEVICE)\n        pred = model(xb)\n        loss = loss_fn(pred, yb)\n        opt.zero_grad(set_to_none=True)\n        loss.backward()\n        opt.step()\n    print(f\"Full-fit epoch {ep:02d} | loss: {loss.item():.5f}\")\n\n\n# ----------------------------\n# 10) Predict test and build submission\n# ----------------------------\n# Use unique image_path to compute embeddings once per image\ntest_unique = test_df.drop_duplicates(\"image_path\").reset_index(drop=True)\n\ntest_ds = BiomassDataset(test_unique, test_img_dir, IMG_COL, targets=None, pil_tfms=valid_pil_tfms)\ntest_dl = DataLoader(test_ds, batch_size=cfg.BATCH_SIZE, shuffle=False,\n                     num_workers=cfg.NUM_WORKERS,  pin_memory=(cfg.DEVICE == \"cuda\"),\n    collate_fn=collate_pil,\n)\n\nX_test, _, test_imgs = extract_embeddings_hf(dinov2, test_dl, cfg.DEVICE)\n\nmodel.eval()\nwith torch.no_grad():\n    p3 = model(torch.tensor(X_test, dtype=torch.float32, device=cfg.DEVICE)).cpu()\n\n# biomass can't be negative\np3 = torch.clamp(p3, min=0.0)\n\np5 = build_5_from_3(p3, gdm_equals_total=cfg.GDM_EQUALS_TOTAL).numpy()\npred_cols_5 = list(cfg.TARGET_ORDER_5)\n\npred_df_wide = pd.DataFrame(p5, columns=pred_cols_5)\npred_df_wide[\"image_path\"] = test_imgs\n\npred_long = pred_df_wide.melt(\n    id_vars=\"image_path\",\n    value_vars=pred_cols_5,\n    var_name=\"target_name\",\n    value_name=\"target\"\n)\n\nsubmission = test_df[[\"sample_id\", \"image_path\", \"target_name\"]].merge(\n    pred_long,\n    on=[\"image_path\", \"target_name\"],\n    how=\"left\"\n)[[\"sample_id\", \"target\"]]\n\n\nsubmission[\"target\"] = pd.to_numeric(submission[\"target\"], errors=\"coerce\").fillna(0.0) \n\nout_path = \"/kaggle/working/submission.csv\"\nsubmission.to_csv(out_path, index=False)\nprint(\"Saved:\", out_path, \"shape:\", submission.shape)\nprint(submission.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}