{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":695833,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":527730,"modelId":541787},{"sourceId":698248,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":529648,"modelId":543645}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# CSIRO Biomass - 0.56-style training recipe\n# Dual backbone: convnext_base + tf_efficientnetv2_s (local safetensors)\n# 3 outputs (Green/Clover/Dead) + derived GDM/Total\n# 20 epochs, KFold, TTA, ensemble\n# ============================================================\n\nTRAIN = True\nDEBUG = False  # if True: 2 folds, 5 epochs\n\nimport os\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.optim import AdamW\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nimport timm\n\ntqdm.pandas()\n\n# ----------------------------\n# Paths\n# ----------------------------\nLOCAL = os.environ.get(\"KAGGLE_KERNEL_RUN_TYPE\", \"\") == \"\"\nDATA_ROOT = \"../input/\" if LOCAL else \"/kaggle/input/csiro-biomass/\"\nOUTPUT_DIR = \"./trained_models\"\nos.makedirs(OUTPUT_DIR, exist_ok=True)\n\nCONVNEXT_CKPT = \"/kaggle/input/convnext-base-imagenet22k-imagenet1k-weights/pytorch/default/1/convnext_base_pretrained.safetensors\"\nEFFNETV2_CKPT = \"/kaggle/input/effnetv2/pytorch/default/1/tf_efficientnetv2_s_pretrained.safetensors\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Device:\", device)\n\n# ----------------------------\n# Data: train wide (one row per image/sample)\n# ----------------------------\ntrain_df = pd.read_csv(f\"{DATA_ROOT}/train.csv\")\ntrain_df[[\"sample_id_prefix\", \"sample_id_suffix\"]] = train_df.sample_id.str.split(\"__\", expand=True)\n\nMETA_COLS = [\"sample_id_prefix\", \"image_path\", \"Sampling_Date\", \"State\", \"Species\", \"Pre_GSHH_NDVI\", \"Height_Ave_cm\"]\nwide = (\n    train_df.pivot_table(\n        index=META_COLS,\n        columns=\"target_name\",\n        values=\"target\",\n        aggfunc=\"first\",\n    )\n    .reset_index()\n)\n\nREQ = [\"Dry_Green_g\", \"Dry_Clover_g\", \"Dry_Dead_g\"]\nfor c in REQ:\n    if c not in wide.columns:\n        wide[c] = 0.0\n\n# KFold\nfrom sklearn.model_selection import KFold\nNFOLD = 5\nkf = KFold(n_splits=NFOLD, shuffle=True, random_state=42)\nwide[\"fold\"] = -1\nfor f, (_, va_idx) in enumerate(kf.split(wide)):\n    wide.loc[va_idx, \"fold\"] = f\n\nprint(\"Train rows:\", len(wide))\n\n# ----------------------------\n# Metric: weighted R2 on 5 targets (3 + derived)\n# ----------------------------\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = np.array([0.1, 0.1, 0.1, 0.2, 0.5], dtype=np.float32)\n    r2s = []\n    for i in range(5):\n        yt = y_true[:, i]\n        yp = y_pred[:, i]\n        ss_res = np.sum((yt - yp) ** 2)\n        ss_tot = np.sum((yt - np.mean(yt)) ** 2)\n        r2 = 1.0 - (ss_res / ss_tot) if ss_tot > 0 else 0.0\n        r2s.append(r2)\n    r2s = np.array(r2s, dtype=np.float32)\n    return float(np.sum(r2s * weights) / np.sum(weights)), r2s\n\ndef calc_metric(outputs3: np.ndarray, targets3: np.ndarray):\n    y_true = np.column_stack([\n        targets3,\n        targets3[:, 0] + targets3[:, 1],                 # GDM\n        targets3[:, 0] + targets3[:, 1] + targets3[:, 2]  # Total\n    ])\n    y_pred = np.column_stack([\n        outputs3,\n        outputs3[:, 0] + outputs3[:, 1],\n        outputs3[:, 0] + outputs3[:, 1] + outputs3[:, 2]\n    ])\n    return weighted_r2_score(y_true, y_pred)\n\n# ----------------------------\n# Loss: same idea as 0.56 notebook\n# ----------------------------\nclass BiomassLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.base = nn.SmoothL1Loss()\n\n    def forward(self, out, y):\n        lg = self.base(out[:, 0], y[:, 0])\n        lc = self.base(out[:, 1], y[:, 1])\n        ld = self.base(out[:, 2], y[:, 2])\n\n        pred_gdm = out[:, 0] + out[:, 1]\n        true_gdm = y[:, 0] + y[:, 1]\n        l_gdm = self.base(pred_gdm, true_gdm)\n\n        pred_total = out.sum(dim=1)\n        true_total = y.sum(dim=1)\n        l_total = self.base(pred_total, true_total)\n\n        return (1.0*lg + 1.0*lc + 1.0*ld + 0.5*l_gdm + 1.0*l_total)\n\n# ----------------------------\n# Dataset / Dataloader: RESIZE(512,512) like your 0.56 code\n# ----------------------------\ndef make_transform(img_size=512, aug=True):\n    if aug:\n        return transforms.Compose([\n            transforms.Resize((img_size, img_size)),\n            transforms.RandomHorizontalFlip(p=0.5),\n            transforms.RandomVerticalFlip(p=0.5),\n            transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n        ])\n    else:\n        return transforms.Compose([\n            transforms.Resize((img_size, img_size)),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n        ])\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df, img_size=512, aug=True, train=True):\n        self.df = df.reset_index(drop=True)\n        self.train = train\n        self.tfm = make_transform(img_size=img_size, aug=aug)\n        self.img_size = img_size\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img = Image.open(os.path.join(DATA_ROOT, row[\"image_path\"])).convert(\"RGB\")\n        img = self.tfm(img)\n\n        if self.train:\n            y = torch.tensor([row[\"Dry_Green_g\"], row[\"Dry_Clover_g\"], row[\"Dry_Dead_g\"]], dtype=torch.float32)\n            return img, y\n        else:\n            return img\n\ndef make_loader(df, img_size=512, batch_size=8, shuffle=True, aug=True, train=True):\n    ds = BiomassDataset(df, img_size=img_size, aug=aug, train=train)\n    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, num_workers=2, pin_memory=True)\n\n# ----------------------------\n# Safetensors load into timm model\n# ----------------------------\ndef load_safetensors(path: str):\n    import safetensors.torch as st\n    sd = st.load_file(path)\n    # strip common wrappers if present\n    if all(k.startswith(\"model.\") for k in sd.keys()):\n        sd = {k.replace(\"model.\", \"\", 1): v for k,v in sd.items()}\n    if all(k.startswith(\"module.\") for k in sd.keys()):\n        sd = {k.replace(\"module.\", \"\", 1): v for k,v in sd.items()}\n    return sd\n\ndef strip_head(sd: dict):\n    drop = (\"head.\", \"fc.\", \"classifier.\")\n    return {k:v for k,v in sd.items() if not k.startswith(drop)}\n\ndef init_pretrained(model, ckpt_path):\n    if not ckpt_path or not Path(ckpt_path).exists():\n        print(\"[init] ckpt missing:\", ckpt_path)\n        return\n    sd = strip_head(load_safetensors(ckpt_path))\n    missing, unexpected = model.load_state_dict(sd, strict=False)\n    print(f\"[init] loaded: {Path(ckpt_path).name} | missing={len(missing)} unexpected={len(unexpected)}\")\n\ndef get_model(model_name, ckpt_path):\n    m = timm.create_model(model_name, pretrained=False, num_classes=3)\n    init_pretrained(m, ckpt_path)\n    return m\n\n# ----------------------------\n# Train / Validate (torch.amp, grad accum)\n# ----------------------------\ndef train_epoch(model, loader, criterion, optimizer, scaler, accum=4):\n    model.train()\n    total = 0.0\n    optimizer.zero_grad(set_to_none=True)\n\n    for step, (x, y) in enumerate(loader):\n        x = x.to(device, non_blocking=True)\n        y = y.to(device, non_blocking=True)\n\n        with torch.amp.autocast(device_type=\"cuda\", enabled=(device.type==\"cuda\")):\n            out = model(x)\n            loss = criterion(out, y) / accum\n\n        scaler.scale(loss).backward()\n\n        if (step + 1) % accum == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad(set_to_none=True)\n\n        total += float(loss.item()) * accum\n\n    if len(loader) % accum != 0:\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad(set_to_none=True)\n\n    return total / len(loader)\n\n@torch.no_grad()\ndef validate(model, loader, criterion):\n    model.eval()\n    total = 0.0\n    outs, trgs = [], []\n\n    for x, y in loader:\n        x = x.to(device, non_blocking=True)\n        y = y.to(device, non_blocking=True)\n\n        with torch.amp.autocast(device_type=\"cuda\", enabled=(device.type==\"cuda\")):\n            out = model(x)\n            loss = criterion(out, y)\n\n        total += float(loss.item())\n        outs.append(out.detach().cpu())\n        trgs.append(y.detach().cpu())\n\n    outs = torch.cat(outs).numpy()\n    trgs = torch.cat(trgs).numpy()\n    w_r2, r2s = calc_metric(outs, trgs)\n    return total/len(loader), w_r2, r2s\n\ndef train_fold(df, fold, model_name, ckpt_path, img_size=512, real_bs=8, accum=4, lr=1e-3,\n               epochs=20, patience=5, tag=\"model\"):\n    tr_df = df[df.fold != fold].copy()\n    va_df = df[df.fold == fold].copy()\n\n    tr_loader = make_loader(tr_df, img_size=img_size, batch_size=real_bs, shuffle=True, aug=True, train=True)\n    va_loader = make_loader(va_df, img_size=img_size, batch_size=real_bs, shuffle=False, aug=False, train=True)\n\n    model = get_model(model_name, ckpt_path).to(device)\n    criterion = BiomassLoss()\n    optimizer = AdamW(model.parameters(), lr=lr)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=\"max\", factor=0.5, patience=max(1,patience//2))\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=(device.type==\"cuda\"))\n\n    best = -1e9\n    bad = 0\n\n    for ep in range(epochs):\n        tr_loss = train_epoch(model, tr_loader, criterion, optimizer, scaler, accum=accum)\n        va_loss, wR2, r2s = validate(model, va_loader, criterion)\n        scheduler.step(wR2)\n\n        # print EVERY epoch (no confusion)\n        print(f\"{tag} | fold {fold} | ep {ep:03d}/{epochs} | tr={tr_loss:.4f} va={va_loss:.4f} \"\n              f\"wR2={wR2:.4f} lr={optimizer.param_groups[0]['lr']:.2e} \"\n              f\"r2=[{', '.join([f'{x:.3f}' for x in r2s])}]\")\n\n        if wR2 > best:\n            best = wR2\n            bad = 0\n            torch.save(model.state_dict(), f\"{OUTPUT_DIR}/{tag}_best_fold{fold}.pth\")\n        else:\n            bad += 1\n\n        if bad >= patience:\n            print(f\"Early stop {tag} fold {fold} | best wR2={best:.4f}\")\n            break\n\n    print(f\"Done {tag} fold {fold} | best wR2={best:.4f}\")\n    return best\n\n# ----------------------------\n# Train both backbones\n# ----------------------------\nSPECS = [\n    {\"tag\":\"convnext_base\", \"name\":\"convnext_base\", \"ckpt\":CONVNEXT_CKPT, \"img\":512, \"bs\":6, \"accum\":4, \"lr\":3e-4},\n    {\"tag\":\"effnetv2_s\",    \"name\":\"tf_efficientnetv2_s\", \"ckpt\":EFFNETV2_CKPT, \"img\":512, \"bs\":8, \"accum\":4, \"lr\":1e-3},\n]\n\nEPOCHS = 5 if DEBUG else 20\nRUN_FOLDS = 2 if DEBUG else NFOLD\n\nif TRAIN:\n    for spec in SPECS:\n        print(\"\\n\" + \"=\"*80)\n        print(f\"Training {spec['tag']} | img={spec['img']} | lr={spec['lr']}\")\n        scores = []\n        for f in range(RUN_FOLDS):\n            scores.append(train_fold(\n                wide, f,\n                model_name=spec[\"name\"],\n                ckpt_path=spec[\"ckpt\"],\n                img_size=spec[\"img\"],\n                real_bs=spec[\"bs\"],\n                accum=spec[\"accum\"],\n                lr=spec[\"lr\"],\n                epochs=EPOCHS,\n                patience=5,\n                tag=spec[\"tag\"]\n            ))\n        print(spec[\"tag\"], \"CV avg:\", float(np.mean(scores)))\n\n# ----------------------------\n# Test wide\n# ----------------------------\ntest_df = pd.read_csv(f\"{DATA_ROOT}/test.csv\")\ntest_df[[\"sample_id_prefix\", \"sample_id_suffix\"]] = test_df.sample_id.str.split(\"__\", expand=True)\n\n# One row per image\ntest_wide = test_df[[\"sample_id_prefix\", \"image_path\"]].drop_duplicates().reset_index(drop=True)\n\n# ----------------------------\n# TTA predict + fold ensemble + backbone ensemble\n# ----------------------------\n@torch.no_grad()\ndef predict_tta(model, loader):\n    model.eval()\n    preds = []\n    for x in loader:\n        x = x.to(device, non_blocking=True)\n        with torch.amp.autocast(device_type=\"cuda\", enabled=(device.type==\"cuda\")):\n            p1 = model(x)\n            p2 = model(torch.flip(x, dims=[3]))\n            p3 = model(torch.flip(x, dims=[2]))\n            p4 = model(torch.flip(x, dims=[2,3]))\n            p = (p1+p2+p3+p4)/4.0\n        preds.append(p.detach().cpu())\n    p = torch.cat(preds).numpy()\n    return np.clip(p, 0, None)\n\ndef predict_backbone(spec):\n    # loader\n    ds = BiomassDataset(test_wide, img_size=spec[\"img\"], aug=False, train=False)\n    loader = DataLoader(ds, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)\n\n    # fold models\n    fold_files = sorted(Path(OUTPUT_DIR).glob(f\"{spec['tag']}_best_fold*.pth\"))\n    if len(fold_files) == 0:\n        raise FileNotFoundError(f\"No fold models found for {spec['tag']} in {OUTPUT_DIR}\")\n\n    all_preds = []\n    for ff in fold_files:\n        model = get_model(spec[\"name\"], spec[\"ckpt\"]).to(device)\n        model.load_state_dict(torch.load(ff, map_location=\"cpu\"), strict=True)\n        all_preds.append(predict_tta(model, loader))\n\n    return np.mean(all_preds, axis=0)\n\npreds_list = []\nfor spec in SPECS:\n    preds_list.append(predict_backbone(spec))\n\npreds3 = np.mean(preds_list, axis=0)\n\ntest_wide[\"Dry_Green_g\"]  = preds3[:, 0]\ntest_wide[\"Dry_Clover_g\"] = preds3[:, 1]\ntest_wide[\"Dry_Dead_g\"]   = preds3[:, 2]\n\n# PHYSICS CONSISTENCY (CRITICAL)\ntest_wide[\"GDM_g\"]       = test_wide[\"Dry_Green_g\"] + test_wide[\"Dry_Clover_g\"]\ntest_wide[\"Dry_Total_g\"] = test_wide[\"GDM_g\"] + test_wide[\"Dry_Dead_g\"]\n\n# submission\ncols = [\"Dry_Clover_g\",\"Dry_Dead_g\",\"Dry_Green_g\",\"Dry_Total_g\",\"GDM_g\"]\nsub = test_wide.set_index(\"sample_id_prefix\")[cols].stack().reset_index()\nsub.columns = [\"sample_id_prefix\",\"target_name\",\"target\"]\nsub[\"sample_id\"] = sub[\"sample_id_prefix\"] + \"__\" + sub[\"target_name\"]\nsub = sub[[\"sample_id\",\"target\"]]\nsub.to_csv(\"submission.csv\", index=False)\n\nprint(sub.head())\nprint(\"Saved submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-28T05:57:48.713801Z","iopub.execute_input":"2025-12-28T05:57:48.714403Z","iopub.status.idle":"2025-12-28T06:17:22.584212Z","shell.execute_reply.started":"2025-12-28T05:57:48.714364Z","shell.execute_reply":"2025-12-28T06:17:22.582614Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}