{"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"},{"sourceId":724049,"sourceType":"modelInstanceVersion","modelInstanceId":550988,"modelId":563608}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ========= EVA02 (only) — load folds + infer + make submission =========\nimport os, math, glob, cv2\nimport numpy as np\nimport pandas as pd\nimport torch, torch.nn as nn\nimport timm\nfrom torch.utils.data import Dataset, DataLoader\n\n# ----------------- CONFIG -----------------\nMODEL_NAME  = \"eva02_large_patch14_448.mim_in22k_ft_in22k_in1k\"\nIMG_SIZE    = 448\nBATCH_SIZE  = 4\nNUM_WORKERS = 2\n\nTEST_CSV  = \"/kaggle/input/csiro-biomass/test.csv\"\n# folder that contains your trained weights: best_model_fold0.pth, best_model_fold1.pth, ...\nWEIGHTS_DIR = \"/kaggle/input/csiro-eva02/pytorch/default/1/eva02\"\nOUT_CSV  = \"submission.csv\"\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntorch.backends.cudnn.benchmark = True\ntorch.set_float32_matmul_precision(\"high\")\n\nALL_TARGET_COLS = [\"Dry_Green_g\",\"Dry_Dead_g\",\"Dry_Clover_g\",\"GDM_g\",\"Dry_Total_g\"]\nMEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)\nSTD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)\n\n# ----------------- image utils -----------------\ndef clean_image(img_rgb: np.ndarray) -> np.ndarray:\n    h, w = img_rgb.shape[:2]\n    img = img_rgb[: int(h * 0.90), :]  # drop bottom 10%\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n    if mask.sum() > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n    return img\n\ndef preprocess_half(img_rgb: np.ndarray, size: int) -> torch.Tensor:\n    # resize -> normalize -> CHW float32\n    im = cv2.resize(img_rgb, (size, size), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255.0\n    im = (im - MEAN) / STD\n    return torch.from_numpy(im).permute(2, 0, 1)  # (3,H,W)\n\n# ----------------- dataset -----------------\ndef resolve_path(rel_path: str, base_dirs):\n    # try exact rel path under each base; fallback to basename under each base\n    for bd in base_dirs:\n        p = os.path.join(bd, rel_path)\n        if os.path.exists(p): return p\n    bn = os.path.basename(rel_path)\n    for bd in base_dirs:\n        p = os.path.join(bd, bn)\n        if os.path.exists(p): return p\n    return None\n\nclass TestDS(Dataset):\n    def __init__(self, image_paths, base_dirs):\n        self.image_paths = list(image_paths)\n        self.base_dirs = list(base_dirs)\n\n    def __len__(self): return len(self.image_paths)\n\n    def __getitem__(self, i):\n        rel = self.image_paths[i]\n        fp = resolve_path(rel, self.base_dirs)\n        img = cv2.imread(fp) if fp else None\n        if img is None:\n            img = np.zeros((1000, 2000, 3), np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        mid = img.shape[1] // 2\n        left  = preprocess_half(img[:, :mid], IMG_SIZE)\n        right = preprocess_half(img[:, mid:], IMG_SIZE)\n        return left, right, rel\n\n# ----------------- model (EVA02 token mode only) -----------------\nclass Local2DTokenMixerBlock(nn.Module):\n    def __init__(self, dim, kernel_size=5, dropout=0.1):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        self.gate = nn.Linear(dim, dim)\n        self.dwconv = nn.Conv2d(dim, dim, kernel_size, padding=kernel_size // 2, groups=dim, bias=True)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x_grid):  # (B,H,W,D)\n        shortcut = x_grid\n        x = self.norm(x_grid)\n        x = x * torch.sigmoid(self.gate(x))\n        x = x.permute(0, 3, 1, 2)   # (B,D,H,W)\n        x = self.dwconv(x)\n        x = x.permute(0, 2, 3, 1)   # (B,H,W,D)\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name: str):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=False, num_classes=0)\n        self.nf = int(self.backbone.num_features)\n        self.num_prefix = int(getattr(self.backbone, \"num_prefix_tokens\", 1))\n\n        gh = gw = None\n        if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"grid_size\"):\n            gs = self.backbone.patch_embed.grid_size\n            if isinstance(gs, (tuple, list)) and len(gs) == 2:\n                gh, gw = int(gs[0]), int(gs[1])\n        self.gh, self.gw = gh, gw\n\n        self.fusion2d = nn.Sequential(\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n        )\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        self.feat_ln = nn.LayerNorm(self.nf)\n\n        def make_pos_head():\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(0.2),\n                nn.Linear(self.nf // 2, 1),\n                nn.Softplus(),\n            )\n        self.head_green  = make_pos_head()\n        self.head_clover = make_pos_head()\n        self.head_dead   = make_pos_head()\n\n    @staticmethod\n    def _infer_grid_hw(num_patches: int):\n        gh = int(math.isqrt(num_patches))\n        while gh > 1 and (num_patches % gh) != 0:\n            gh -= 1\n        gw = num_patches // gh\n        return gh, gw\n\n    def _grid_from_input(self, x: torch.Tensor, num_patches: int):\n        if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"patch_size\"):\n            ps = self.backbone.patch_embed.patch_size\n            ph, pw = (ps if isinstance(ps, (tuple, list)) else (ps, ps))\n            gh = int(x.shape[-2] // ph)\n            gw = int(x.shape[-1] // pw)\n            if gh * gw == num_patches:\n                return gh, gw\n        return self._infer_grid_hw(num_patches)\n\n    @staticmethod\n    def _concat_halves(p_l, p_r, gh, gw):\n        B, Np, D = p_l.shape\n        pl2 = p_l.view(B, gh, gw, D)\n        pr2 = p_r.view(B, gh, gw, D)\n        full2 = torch.cat([pl2, pr2], dim=2)      # (B,gh,2gw,D)\n        return full2.reshape(B, gh * (2 * gw), D)\n\n    def forward(self, left, right):\n        x_l = self.backbone.forward_features(left)   # (B,N,D)\n        x_r = self.backbone.forward_features(right)  # (B,N,D)\n\n        p_l = x_l[:, self.num_prefix:, :]\n        p_r = x_r[:, self.num_prefix:, :]\n        Np = int(p_l.size(1))\n\n        gh, gw = self.gh, self.gw\n        if (gh is None) or (gw is None) or (gh * gw != Np):\n            gh, gw = self._grid_from_input(left, Np)\n\n        p_seq = self._concat_halves(p_l, p_r, gh, gw)  # (B, gh*(2gw), D)\n        B, N, D = p_seq.shape\n        H, W = gh, 2 * gw\n\n        with torch.amp.autocast(\"cuda\", enabled=False):\n            p_grid = p_seq.float().view(B, H, W, D).contiguous()\n            p_grid = self.fusion2d(p_grid)\n            p_fused = p_grid.view(B, H * W, D)\n\n            feat = self.pool(p_fused.transpose(1, 2)).squeeze(-1)\n            feat = self.feat_ln(feat)\n\n        green  = self.head_green(feat)\n        clover = self.head_clover(feat)\n        dead   = self.head_dead(feat)\n        gdm    = green + clover\n        total  = gdm + dead\n        return total, gdm, green  # keep it lean (derive clover/dead later)\n\n# ----------------- load folds + infer -----------------\ndef load_sd(path):\n    state = torch.load(path, map_location=\"cpu\")\n    if isinstance(state, dict) and (\"model_state_dict\" in state or \"state_dict\" in state):\n        state = state.get(\"model_state_dict\", state.get(\"state_dict\"))\n    return state\n\n@torch.inference_mode()\ndef predict_one_checkpoint(model, loader, ckpt_path):\n    model.load_state_dict(load_sd(ckpt_path), strict=False)\n    model.to(DEVICE).eval()\n\n    preds = np.zeros((len(loader.dataset), 3), np.float32)  # [total, gdm, green]\n    off = 0\n    for l, r, _ in loader:\n        l = l.to(DEVICE, non_blocking=True)\n        r = r.to(DEVICE, non_blocking=True)\n        if DEVICE.type == \"cuda\":\n            with torch.autocast(\"cuda\", dtype=torch.bfloat16):\n                total, gdm, green = model(l, r)\n        else:\n            total, gdm, green = model(l, r)\n\n        b = l.size(0)\n        pred3 = torch.stack([total.view(-1), gdm.view(-1), green.view(-1)], dim=1).float().cpu().numpy()\n        preds[off:off+b] = pred3\n        off += b\n    return preds\n\ndef postprocess_3_to_5(pred3):\n    total = pred3[:, 0]\n    gdm   = pred3[:, 1]\n    green = pred3[:, 2]\n    clover = np.maximum(0.0, gdm - green)\n    dead   = np.maximum(0.0, total - gdm)\n    # order required:\n    return np.stack([green, dead, clover, gdm, total], axis=1).astype(np.float32)\n\n# ----------------- run -----------------\ntest_df = pd.read_csv(TEST_CSV)\nuniq_imgs = test_df[\"image_path\"].drop_duplicates().values\n\n# base dirs to try (add more if you keep images elsewhere)\nbase_dirs = [\n    \"/kaggle/input/csiro-biomass\",              # sometimes image_path is relative under this\n    \"/kaggle/input/csiro-biomass/test_images\",  # common layout\n    \"/kaggle/input/csiro-biomass/test\",         # alt layout\n]\n\nds = TestDS(uniq_imgs, base_dirs)\ndl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n\nckpts = sorted(glob.glob(os.path.join(WEIGHTS_DIR, \"best_model_fold*.pth\")))\nif not ckpts:\n    raise FileNotFoundError(f\"No fold checkpoints found in: {WEIGHTS_DIR}\")\n\nmodel = BiomassModel(MODEL_NAME)\npreds_sum = np.zeros((len(ds), 3), np.float32)\n\nfor p in ckpts:\n    preds_sum += predict_one_checkpoint(model, dl, p)\npreds_3 = preds_sum / float(len(ckpts))\npreds_5 = postprocess_3_to_5(preds_3)\n\npreds_wide = pd.DataFrame(preds_5, columns=ALL_TARGET_COLS)\npreds_wide.insert(0, \"image_path\", uniq_imgs)\n\npreds_long = preds_wide.melt(\n    id_vars=[\"image_path\"],\n    value_vars=ALL_TARGET_COLS,\n    var_name=\"target_name\",\n    value_name=\"target\",\n)\n\nsub = (\n    test_df[[\"sample_id\", \"image_path\", \"target_name\"]]\n    .merge(preds_long, on=[\"image_path\", \"target_name\"], how=\"left\")[[\"sample_id\", \"target\"]]\n    .fillna(0.0)\n    .sort_values(\"sample_id\")\n    .reset_index(drop=True)\n)\n\nsub.to_csv(OUT_CSV, index=False)\nprint(\"saved:\", OUT_CSV)\nsub.head()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}