{"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":293150034,"sourceType":"kernelVersion"}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Image2Biomass — INFER + SUBMIT notebook (Kaggle, PYTORCH)\n#\n# This notebook:\n# 1) Loads trained fold checkpoints from the TRAIN notebook (/kaggle/working/i2b/models/*.pth)\n# 2) Runs inference on unique test images with the SAME preprocessing (cleaning + tiling/TTA)\n# 3) Ensembles across all found checkpoints\n# 4) Writes submission.csv in the LONG format expected by the competition\n#\n# IMPORTANT (Kaggle workflow):\n# - If this is a separate notebook, you must \"Add Data\" -> your saved dataset that contains i2b/models/*.pth.\n# - By default we look in /kaggle/input/i2b-artifacts, then fall back to /kaggle/working/i2b.\n\n# %% [0] Setup\nimport os, gc, math\n\n# ---- Offline safety (separate notebook, NO internet) ----\n# Kaggle internet should already be OFF, but these env vars prevent some libs from even trying.\nos.environ.setdefault(\"HF_HUB_DISABLE_TELEMETRY\", \"1\")\nos.environ.setdefault(\"HF_HUB_OFFLINE\", \"1\")\nos.environ.setdefault(\"TRANSFORMERS_OFFLINE\", \"1\")\nos.environ.setdefault(\"HF_DATASETS_OFFLINE\", \"1\")\nos.environ.setdefault(\"WANDB_MODE\", \"disabled\")\n\nfrom typing import List, Dict\n\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\n\nimport timm\nfrom timm.data import resolve_data_config\n\nfrom torchvision import transforms as tvt\n\n# Optional: OpenCV for image cleaning (timestamp/artifact removal)\ntry:\n    import cv2\n    _HAS_CV2 = True\nexcept Exception as _e:\n    cv2 = None\n    _HAS_CV2 = False\n    print(\"[warn] cv2 not available; image cleaning disabled:\", repr(_e))\n\n# ------------------\n# Paths\n# ------------------\nDATA_DIR = \"/kaggle/input/csiro-biomass\"\nif not os.path.exists(DATA_DIR):\n    base = \"/kaggle/input\"\n    cands = [os.path.join(base, d) for d in os.listdir(base) if (\"csiro\" in d.lower() and \"biomass\" in d.lower())]\n    if cands:\n        DATA_DIR = cands[0]\nprint(\"Using DATA_DIR:\", DATA_DIR)\n\nTEST_CSV = os.path.join(DATA_DIR, \"test.csv\")\n\n# Where trained artifacts live (this NOTEBOOK must be separate and OFFLINE)\n# Preferred workflow:\n#   1) In your TRAIN notebook: Save Version -> \"Save output\" ON\n#   2) In this INFER notebook: Add Data -> that saved dataset\n# This code will auto-detect a dataset that contains i2b/models/*.pth.\n\ndef _autodetect_artifacts_root() -> str:\n    base = \"/kaggle/input\"\n    if not os.path.exists(base):\n        return \"/kaggle/working/i2b\"\n\n    # Common layouts:\n    #   /kaggle/input/<dataset>/i2b/models/*.pth\n    #   /kaggle/input/<dataset>/models/*.pth\n    for d in sorted(os.listdir(base)):\n        root = os.path.join(base, d)\n        if not os.path.isdir(root):\n            continue\n        if os.path.exists(os.path.join(root, \"i2b\", \"models\")):\n            return os.path.join(root, \"i2b\")\n        if os.path.exists(os.path.join(root, \"models\")):\n            # Heuristic: any fold checkpoint inside?\n            try:\n                fns = os.listdir(os.path.join(root, \"models\"))\n            except Exception:\n                fns = []\n            if any((fn.endswith(\".pth\") and \"fold\" in fn) for fn in fns):\n                return root\n\n    return \"/kaggle/working/i2b\"\n\nARTIFACT_DIR = _autodetect_artifacts_root()\nprint(\"Using ARTIFACT_DIR:\", ARTIFACT_DIR)\n\nMODELS_DIR = os.path.join(ARTIFACT_DIR, \"models\")\nif not os.path.exists(MODELS_DIR) and os.path.exists(os.path.join(ARTIFACT_DIR, \"i2b\", \"models\")):\n    MODELS_DIR = os.path.join(ARTIFACT_DIR, \"i2b\", \"models\")\nprint(\"Using MODELS_DIR:\", MODELS_DIR)\nOUT_SUB_PATH = \"/kaggle/working/submission.csv\"\n\n# ------------------\n# Competition schema\n# ------------------\nSAMPLE_ID_COL = \"sample_id\"\nIMG_COL = \"image_path\"\nTNAME_COL = \"target_name\"\n\nTARGET_COLS = [\n    \"Dry_Green_g\",\n    \"Dry_Dead_g\",\n    \"Dry_Clover_g\",\n    \"GDM_g\",\n    \"Dry_Total_g\",\n]\nTNAME_TO_IDX = {t: i for i, t in enumerate(TARGET_COLS)}\n\n# ------------------\n# Inference / TTA settings (match TRAIN defaults)\n# ------------------\nBATCH_SIZE = 16\nNUM_WORKERS = 2\nAMP = True\n\nVAL_TTA_MCROPS = 5\nVAL_TTA_MODE = \"five\"    # \"five\" or \"grid\"\nVAL_RESIZE = 448         # must be >= IMG_SIZE; preserves texture\n\n# ------------------\n# Optional image cleaning (MUST match TRAIN notebook if enabled there)\n# ------------------\nUSE_CLEANING = True\nCLEAN_CROP_BOTTOM_PCT = 0.10\nCLEAN_INPAINT_ORANGE = True\nCLEAN_ORANGE_LOWER = (5, 150, 150)\nCLEAN_ORANGE_UPPER = (25, 255, 255)\nCLEAN_DILATE_ITER = 2\nCLEAN_CLAHE = False\nCLEAN_WHITEBALANCE = False\n\n# ------------------\n# Device\n# ------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"torch:\", torch.__version__)\nprint(\"cuda available:\", torch.cuda.is_available())\nif torch.cuda.is_available():\n    print(\"GPU:\", torch.cuda.get_device_name(0))\n\n# %% [1] Load test (LONG)\n\ntest_long = pd.read_csv(TEST_CSV)\nneed = {IMG_COL, TNAME_COL}\nmissing = sorted(list(need - set(test_long.columns)))\nif missing:\n    raise ValueError(f\"test.csv missing columns: {missing}\")\n\nif SAMPLE_ID_COL not in test_long.columns:\n    # fallback (rare)\n    SAMPLE_ID_COL = test_long.columns[0]\n\n# Unique images to run inference on\n# (test is LONG: one row per (image_path, target_name))\ntest_img_df = test_long.drop_duplicates(subset=[IMG_COL]).reset_index(drop=True)\nprint(\"test_long:\", test_long.shape, \"| unique images:\", test_img_df.shape)\n\n# %% [2] Find checkpoints\n\ndef find_checkpoints(models_dir: str) -> List[str]:\n    if not os.path.exists(models_dir):\n        return []\n    ckpts = []\n    for fn in os.listdir(models_dir):\n        if fn.endswith(\".pth\") and (\"fold\" in fn):\n            ckpts.append(os.path.join(models_dir, fn))\n    ckpts = sorted(ckpts)\n    return ckpts\n\nckpt_paths = find_checkpoints(MODELS_DIR)\nif not ckpt_paths:\n    raise FileNotFoundError(\n        f\"No .pth checkpoints found in {MODELS_DIR}. \"\n        \"Make sure you added the dataset with i2b/models/*.pth, or run inference in the same notebook after training.\"\n    )\n\nprint(\"Found checkpoints:\")\nfor p in ckpt_paths:\n    print(\" -\", p)\n\n# %% [3] Helpers (loading, cleaning, crops)\n\ndef safe_torch_load(path: str):\n    # Torch 2.6+ supports weights_only; older versions may not.\n    try:\n        return torch.load(path, map_location=\"cpu\", weights_only=True)\n    except TypeError:\n        return torch.load(path, map_location=\"cpu\")\n\n\ndef _grey_world_white_balance(img_rgb: np.ndarray) -> np.ndarray:\n    img = img_rgb.astype(np.float32)\n    mean = img.reshape(-1, 3).mean(axis=0)\n    scale = float(mean.mean()) / (mean + 1e-6)\n    img = img * scale.reshape(1, 1, 3)\n    return np.clip(img, 0.0, 255.0).astype(np.uint8)\n\n\ndef _apply_clahe_rgb(img_rgb: np.ndarray, clip_limit: float = 2.0, grid: int = 8) -> np.ndarray:\n    lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)\n    l, a, b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=float(clip_limit), tileGridSize=(int(grid), int(grid)))\n    l2 = clahe.apply(l)\n    lab2 = cv2.merge((l2, a, b))\n    return cv2.cvtColor(lab2, cv2.COLOR_LAB2RGB)\n\n\ndef clean_image_pil(img: Image.Image) -> Image.Image:\n    if (not USE_CLEANING) or (not _HAS_CV2):\n        return img\n\n    arr = np.asarray(img)\n    if arr.ndim != 3 or arr.shape[2] != 3:\n        return img\n\n    h, w = arr.shape[:2]\n\n    # 1) Crop bottom artifacts\n    if CLEAN_CROP_BOTTOM_PCT and CLEAN_CROP_BOTTOM_PCT > 0:\n        cut = int(round(float(h) * (1.0 - float(CLEAN_CROP_BOTTOM_PCT))))\n        cut = max(1, min(int(h), int(cut)))\n        arr = arr[:cut, :, :]\n\n    # 2) Lighting normalization (optional)\n    if CLEAN_WHITEBALANCE:\n        arr = _grey_world_white_balance(arr)\n    if CLEAN_CLAHE:\n        arr = _apply_clahe_rgb(arr)\n\n    # 3) Inpaint orange timestamp\n    if CLEAN_INPAINT_ORANGE:\n        hsv = cv2.cvtColor(arr, cv2.COLOR_RGB2HSV)\n        lower = np.array(CLEAN_ORANGE_LOWER, dtype=np.uint8)\n        upper = np.array(CLEAN_ORANGE_UPPER, dtype=np.uint8)\n        mask = cv2.inRange(hsv, lower, upper)\n\n        if CLEAN_DILATE_ITER and CLEAN_DILATE_ITER > 0:\n            kernel = np.ones((3, 3), np.uint8)\n            mask = cv2.dilate(mask, kernel, iterations=int(CLEAN_DILATE_ITER))\n\n        if int(mask.sum()) > 0:\n            arr = cv2.inpaint(arr, mask, 3, cv2.INPAINT_TELEA)\n\n    return Image.fromarray(arr)\n\n\ndef ensure_min_side(img: Image.Image, target_min: int) -> Image.Image:\n    w, h = img.size\n    ms = min(w, h)\n    if ms >= target_min:\n        return img\n    scale = float(target_min) / float(ms)\n    new_w = int(round(w * scale))\n    new_h = int(round(h * scale))\n    return img.resize((new_w, new_h), resample=Image.BICUBIC)\n\n\ndef resize_for_tiling(img: Image.Image, resize_min: int, crop_size: int) -> Image.Image:\n    resize_min = int(resize_min)\n    crop_size = int(crop_size)\n    target_min = max(resize_min, crop_size)\n    img = ensure_min_side(img, target_min)\n\n    if resize_min and resize_min > crop_size:\n        w, h = img.size\n        scale = float(resize_min) / float(min(w, h))\n        if abs(scale - 1.0) > 1e-6:\n            new_w = int(round(w * scale))\n            new_h = int(round(h * scale))\n            img = img.resize((new_w, new_h), resample=Image.BICUBIC)\n\n    return img\n\n\ndef crop_positions(w: int, h: int, crop: int, n: int, mode: str):\n    max_x = max(0, w - crop)\n    max_y = max(0, h - crop)\n\n    if n <= 1:\n        return [(max_x // 2, max_y // 2)]\n\n    if mode == \"five\":\n        pos = [\n            (0, 0),\n            (max_x, 0),\n            (0, max_y),\n            (max_x, max_y),\n            (max_x // 2, max_y // 2),\n        ]\n        return pos[:n]\n\n    # grid\n    g = int(math.ceil(math.sqrt(n)))\n    xs = np.linspace(0, max_x, g)\n    ys = np.linspace(0, max_y, g)\n    pos = []\n    for yy in ys:\n        for xx in xs:\n            pos.append((int(round(xx)), int(round(yy))))\n    return pos[:n]\n\n# %% [4] Model (same as TRAIN)\n\nINDEP_OUT_DIM = 3  # DG, DD, DC\n\nclass ViTRegressor(nn.Module):\n    def __init__(self, model_name: str, out_dim: int = INDEP_OUT_DIM, pretrained: bool = False):\n        super().__init__()\n        self.out_dim = int(out_dim)\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool=\"avg\")\n        n = self.backbone.num_features\n        self.head = nn.Sequential(\n            nn.LayerNorm(n),\n            nn.Linear(n, 512),\n            nn.GELU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, self.out_dim),\n        )\n\n    def forward(self, x):\n        feat = self.backbone(x)\n        return self.head(feat)\n\n\ndef modelspace3_to_grams5(yhat3_model: torch.Tensor) -> torch.Tensor:\n    # yhat3_model: log1p grams for (DG, DD, DC)\n    g3 = torch.expm1(yhat3_model).clamp_min(0.0)\n    dg = g3[:, 0]\n    dd = g3[:, 1]\n    dc = g3[:, 2]\n    gdm = dg + dc\n    total = dd + gdm\n    g5 = torch.stack([dg, dd, dc, gdm, total], dim=1)\n    return g5\n\n# %% [5] Build normalization from the first checkpoint's model_name\n\nfirst_ckpt = safe_torch_load(ckpt_paths[0])\nmodel_name = str(first_ckpt.get(\"model_name\", \"vit_large_patch14_clip_224.openai\"))\nimg_size = int(first_ckpt.get(\"img_size\", 224))\nprint(\"Checkpoint model_name:\", model_name)\nprint(\"Checkpoint img_size:\", img_size)\n\n# Use timm config for mean/std (avoid mismatch)\n_cfg_model = timm.create_model(model_name, pretrained=False)\ncfg = resolve_data_config({}, model=_cfg_model)\ndel _cfg_model\n\nIMG_SIZE = int(img_size)\nMEAN = cfg.get(\"mean\")\nSTD = cfg.get(\"std\")\n\nval_post = tvt.Compose([\n    tvt.ToTensor(),\n    tvt.Normalize(mean=MEAN, std=STD),\n])\n\n# %% [6] Dataset / Loader (unique images, multi-crop)\n\nclass TestTiledDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, n_crops: int, crop_mode: str):\n        self.df = df.reset_index(drop=True)\n        self.n_crops = int(max(1, n_crops))\n        self.crop_mode = str(crop_mode)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, i: int):\n        rel = str(self.df.loc[i, IMG_COL])\n        path = os.path.join(DATA_DIR, rel)\n\n        img = Image.open(path).convert(\"RGB\")\n        img = clean_image_pil(img)\n        img = resize_for_tiling(img, VAL_RESIZE, IMG_SIZE)\n\n        w, h = img.size\n        positions = crop_positions(w, h, int(IMG_SIZE), self.n_crops, self.crop_mode)\n        crops = [img.crop((x0, y0, x0 + int(IMG_SIZE), y0 + int(IMG_SIZE))) for (x0, y0) in positions]\n\n        xs = torch.stack([val_post(c) for c in crops], dim=0)  # (K,3,H,W)\n        return xs, rel\n\n\ndef make_loader(df: pd.DataFrame) -> DataLoader:\n    ds = TestTiledDataset(df, n_crops=VAL_TTA_MCROPS, crop_mode=VAL_TTA_MODE)\n    return DataLoader(\n        ds,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n        drop_last=False,\n        persistent_workers=(NUM_WORKERS > 0),\n        prefetch_factor=2,\n    )\n\nloader = make_loader(test_img_df)\n\n# %% [7] Predict per checkpoint, ensemble\n\n@torch.no_grad()\ndef predict_one(model: nn.Module, loader: DataLoader) -> Dict[str, np.ndarray]:\n    model.eval()\n    out: Dict[str, np.ndarray] = {}\n\n    for xb, rels in loader:\n        # xb: (B,K,3,H,W)\n        b, k, c, h, w = xb.shape\n        xb = xb.view(b * k, c, h, w).to(device, non_blocking=True)\n\n        with torch.amp.autocast('cuda', enabled=(AMP and device.type == 'cuda')):\n            yhat3 = model(xb)                           # (b*k,3) model-space\n            g3 = torch.expm1(yhat3).clamp_min(0.0)      # (b*k,3) grams\n\n        g3 = g3.view(b, k, 3).mean(1)                   # (b,3) grams\n        dg = g3[:, 0]\n        dd = g3[:, 1]\n        dc = g3[:, 2]\n        gdm = dg + dc\n        total = dd + gdm\n        g5 = torch.stack([dg, dd, dc, gdm, total], dim=1)  # (b,5) grams\n\n        g5 = g5.detach().cpu().numpy().astype(np.float32)\n\n        for j in range(b):\n            out[str(rels[j])] = g5[j]\n\n    return out\n\n\ndef load_model_from_ckpt(path: str) -> nn.Module:\n    ckpt = safe_torch_load(path)\n    mn = str(ckpt.get(\"model_name\", model_name))\n    od = int(ckpt.get(\"out_dim\", INDEP_OUT_DIM))\n\n    m = ViTRegressor(mn, out_dim=od, pretrained=False)\n\n    # state_dict may be saved under different keys depending on TRAIN code\n    sd = ckpt.get(\"state_dict\")\n    if sd is None:\n        sd = ckpt.get(\"model_state\")\n    if sd is None:\n        sd = ckpt  # last resort (rare): raw state_dict\n\n    if not isinstance(sd, dict):\n        raise KeyError(f\"Checkpoint missing a usable state_dict: {path}\")\n\n    # Handle DataParallel 'module.' prefix\n    if any(k.startswith(\"module.\") for k in sd.keys()):\n        sd = {k[len(\"module.\"):]: v for k, v in sd.items()}\n\n    # Try strict load first; if it fails, fall back to non-strict and print what happened.\n    try:\n        m.load_state_dict(sd, strict=True)\n    except RuntimeError as e:\n        missing, unexpected = m.load_state_dict(sd, strict=False)\n        print(f\"[warn] non-strict load for {os.path.basename(path)}\")\n        print(\"  missing:\", list(missing)[:10], (\"...\" if len(missing) > 10 else \"\"))\n        print(\"  unexpected:\", list(unexpected)[:10], (\"...\" if len(unexpected) > 10 else \"\"))\n        print(\"  strict error:\", repr(e))\n\n    m.to(device)\n    if torch.cuda.device_count() > 1:\n        m = nn.DataParallel(m)\n\n    return m\n\n# Ensemble: sum predictions then divide\nsum_pred: Dict[str, np.ndarray] = {str(ip): np.zeros((5,), np.float32) for ip in test_img_df[IMG_COL].astype(str).values}\n\nfor ci, ckpt_path in enumerate(ckpt_paths):\n    print(f\"[{ci+1}/{len(ckpt_paths)}] loading:\", ckpt_path)\n    model = load_model_from_ckpt(ckpt_path).to(device)\n\n    pred_map = predict_one(model, loader)\n    for k, v in pred_map.items():\n        sum_pred[k] += v\n\n    del model\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\nn_models = float(len(ckpt_paths))\nfor k in sum_pred:\n    sum_pred[k] /= n_models\n\n# %% [8] Build submission (LONG)\n\nout_target = np.zeros((len(test_long),), np.float32)\n\n# Fast vectorized map: create a DataFrame of preds in wide form then merge\npred_wide = pd.DataFrame({\n    IMG_COL: list(sum_pred.keys()),\n    **{t: [sum_pred[ip][j] for ip in sum_pred.keys()] for j, t in enumerate(TARGET_COLS)}\n})\n\nmerged = test_long.merge(pred_wide, on=IMG_COL, how=\"left\")\nif merged[TARGET_COLS[0]].isna().any():\n    n_missing = int(merged[TARGET_COLS[0]].isna().sum())\n    raise RuntimeError(f\"Missing predictions for {n_missing} test rows after merge (unexpected).\")\n\n# Pick the correct column per row target_name\nfor t in TARGET_COLS:\n    merged.loc[merged[TNAME_COL] == t, \"target\"] = merged.loc[merged[TNAME_COL] == t, t]\n\nsub = merged[[SAMPLE_ID_COL, \"target\"]].copy()\nsub.to_csv(OUT_SUB_PATH, index=False)\n\nprint(\"Saved:\", OUT_SUB_PATH)\nprint(sub.head())\n\n# %% [9] Notes\n# - Public LB will tell us if CV is honest or if we're overfitting quirks in the grouped split.\n# - If LB << CV, the first things to try are: stronger image cleaning, different split scheme, and/or lighter TTA.\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-17T00:02:48.063684Z","iopub.execute_input":"2026-01-17T00:02:48.063937Z","iopub.status.idle":"2026-01-17T00:03:26.579275Z","shell.execute_reply.started":"2026-01-17T00:02:48.063902Z","shell.execute_reply":"2026-01-17T00:03:26.578579Z"}},"outputs":[],"execution_count":null}]}