{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":724178,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":551090,"modelId":563700}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ========= SwinV2 (only) — load folds + infer + make submission =========\nimport os, glob, math, cv2\nimport numpy as np\nimport pandas as pd\nimport torch, torch.nn as nn\nimport torch.nn.functional as F\nimport timm\n\n# ----------------- CONFIG -----------------\nMODEL_NAME  = \"swinv2_large_window12to24_192to384.ms_in22k_ft_in1k\"\nIMG_SIZE    = 384\nBATCH_SIZE  = 4\nNUM_WORKERS = 2\n\nTEST_CSV    = \"/kaggle/input/csiro-biomass/test.csv\"\nWEIGHTS_DIR = \"/kaggle/input/swinv2/pytorch/default/1/swinv2\"  # has best_model_fold0.pth ...\nOUT_CSV     = \"submission.csv\"\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = 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], np.float32)\nSTD  = np.array([0.229, 0.224, 0.225], np.float32)\n\n# ----------------- image utils -----------------\ndef clean_image(img_rgb):\n    h, w = img_rgb.shape[:2]\n    img = img_rgb[: int(h * 0.90), :]\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    mask = cv2.inRange(hsv, np.array([5,150,150]), np.array([25,255,255]))\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, size):\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)  # CHW\n\ndef resolve_path(rel_path, base_dirs):\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\n# ----------------- dataset -----------------\nfrom torch.utils.data import Dataset, DataLoader\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 blocks (from your nb, feature-path) -----------------\ntry:\n    from timm.layers import DropPath\nexcept Exception:\n    from timm.models.layers import DropPath\n\nclass Local2DTokenMixerConvNextBlock(nn.Module):\n    def __init__(self, dim, kernel_size=7, mlp_ratio=4.0, dropout=0.0, drop_path=0.0, layer_scale_init=1e-6):\n        super().__init__()\n        self.dwconv = nn.Conv2d(dim, dim, kernel_size, padding=kernel_size//2, groups=dim, bias=True)\n        self.norm = nn.LayerNorm(dim, eps=1e-6)\n        hidden = int(dim * mlp_ratio)\n        self.pw1 = nn.Linear(dim, hidden)\n        self.act = nn.GELU()\n        self.pw2 = nn.Linear(hidden, dim)\n        self.drop = nn.Dropout(dropout)\n        self.drop_path = DropPath(drop_path) if drop_path and drop_path > 0 else nn.Identity()\n        self.gamma = nn.Parameter(layer_scale_init * torch.ones(dim)) if layer_scale_init and layer_scale_init > 0 else None\n\n    def forward(self, x_grid):  # (B,H,W,C)\n        shortcut = x_grid\n        x = x_grid.permute(0,3,1,2).contiguous()\n        x = self.dwconv(x)\n        x = x.permute(0,2,3,1).contiguous()\n        x = self.norm(x)\n        x = self.pw1(x); x = self.act(x); x = self.drop(x); x = self.pw2(x)\n        if self.gamma is not None:\n            x = x * self.gamma\n        x = self.drop(x)\n        x = self.drop_path(x)\n        return shortcut + x\n\nclass GeM2d(nn.Module):\n    def __init__(self, p=3.0, eps=1e-6):\n        super().__init__()\n        self.p = nn.Parameter(torch.ones(1) * float(p))\n        self.eps = float(eps)\n    def forward(self, x):  # (B,C,H,W)->(B,C)\n        x = x.clamp(min=self.eps).pow(self.p)\n        x = x.mean(dim=(2,3)).pow(1.0 / self.p)\n        return x\n\nclass SeparableConvGNAct(nn.Module):\n    def __init__(self, c, kernel_size=3, gn_groups=32):\n        super().__init__()\n        g = min(int(gn_groups), int(c))\n        while g > 1 and (c % g) != 0: g -= 1\n        self.dw = nn.Conv2d(c, c, kernel_size, padding=kernel_size//2, groups=c, bias=False)\n        self.pw = nn.Conv2d(c, c, 1, bias=False)\n        self.gn = nn.GroupNorm(g, c)\n        self.act = nn.GELU()\n    def forward(self, x):\n        x = self.dw(x); x = self.pw(x); x = self.gn(x); x = self.act(x)\n        return x\n\ndef _resize_like(x, ref):\n    if x.shape[-2:] == ref.shape[-2:]:\n        return x\n    if x.shape[-2] > ref.shape[-2] or x.shape[-1] > ref.shape[-1]:\n        return F.interpolate(x, size=ref.shape[-2:], mode=\"area\")\n    return F.interpolate(x, size=ref.shape[-2:], mode=\"nearest\")\n\nclass StereoCrossGate2D(nn.Module):\n    def __init__(self, c, kernel_size=3, gn_groups=32, layer_scale_init=1e-3):\n        super().__init__()\n        g = min(int(gn_groups), int(c))\n        while g > 1 and (c % g) != 0: g -= 1\n        self.normL = nn.GroupNorm(g, c)\n        self.normR = nn.GroupNorm(g, c)\n        pad = kernel_size // 2\n        self.gate_from_L = nn.Sequential(\n            nn.Conv2d(c, c, kernel_size, padding=pad, groups=c, bias=False),\n            nn.Conv2d(c, c, 1, bias=True),\n        )\n        self.gate_from_R = nn.Sequential(\n            nn.Conv2d(c, c, kernel_size, padding=pad, groups=c, bias=False),\n            nn.Conv2d(c, c, 1, bias=True),\n        )\n        self.ls = nn.Parameter(torch.ones(c) * float(layer_scale_init))\n\n    def forward(self, xL, xR):\n        l = self.normL(xL); r = self.normR(xR)\n        gL = torch.sigmoid(self.gate_from_L(l))\n        gR = torch.sigmoid(self.gate_from_R(r))\n        ls = self.ls.view(1,-1,1,1).to(xL.dtype)\n        xL = xL + ls * (xL * gR)\n        xR = xR + ls * (xR * gL)\n        return xL, xR\n\nclass FuseToP16(nn.Module):\n    def __init__(self, c, gn_groups=32):\n        super().__init__()\n        self.w8  = nn.Conv2d(c, 1, 1, bias=True)\n        self.w16 = nn.Conv2d(c, 1, 1, bias=True)\n        self.w32 = nn.Conv2d(c, 1, 1, bias=True)\n        self.out = SeparableConvGNAct(c, kernel_size=3, gn_groups=gn_groups)\n    def forward(self, p8, p16, p32):\n        p8r  = _resize_like(p8,  p16)\n        p32r = _resize_like(p32, p16)\n        w = torch.cat([self.w8(p8r), self.w16(p16), self.w32(p32r)], dim=1)  # (B,3,H,W)\n        w = torch.softmax(w, dim=1)\n        fused = w[:,0:1]*p8r + w[:,1:2]*p16 + w[:,2:3]*p32r\n        return self.out(fused)\n\nclass SeamBlend2D(nn.Module):\n    def __init__(self, c):\n        super().__init__()\n        self.dw = nn.Conv2d(c, c, kernel_size=(3,9), padding=(1,4), groups=c, bias=False)\n        self.pw = nn.Conv2d(c, c, 1, bias=True)\n        self.act = nn.GELU()\n    def forward(self, x):\n        return x + self.pw(self.act(self.dw(x)))\n\nclass BiomassSwinV2(nn.Module):\n    def __init__(self, model_name, pretrained=False, head_dropout=0.1, fuse_dim=None):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, features_only=True, out_indices=(0,1,2,3))\n        fi = self.backbone.feature_info\n        chans = list(fi.channels())\n        reds  = list(fi.reduction())\n\n        def pick(target):\n            return int(min(range(len(reds)), key=lambda i: abs(int(reds[i]) - int(target))))\n        self.idx8, self.idx16, self.idx32 = pick(8), pick(16), pick(32)\n        ch8, ch16, ch32 = int(chans[self.idx8]), int(chans[self.idx16]), int(chans[self.idx32])\n\n        self.nf = int(fuse_dim or ch16)\n\n        def proj(in_ch):\n            g = min(32, self.nf)\n            while g > 1 and (self.nf % g) != 0: g -= 1\n            return nn.Sequential(\n                nn.Conv2d(in_ch, self.nf, 1, bias=False),\n                nn.GroupNorm(g, self.nf),\n                nn.GELU(),\n            )\n        self.proj8, self.proj16, self.proj32 = proj(ch8), proj(ch16), proj(ch32)\n\n        self.stereo8  = StereoCrossGate2D(self.nf, kernel_size=3, gn_groups=32, layer_scale_init=1e-3)\n        self.stereo16 = StereoCrossGate2D(self.nf, kernel_size=3, gn_groups=32, layer_scale_init=1e-3)\n        self.stereo32 = StereoCrossGate2D(self.nf, kernel_size=3, gn_groups=32, layer_scale_init=1e-3)\n\n        self.fuse_to_p16 = FuseToP16(self.nf, gn_groups=32)\n        self.seam16 = SeamBlend2D(self.nf)\n\n        self.fusion2d = nn.Sequential(\n            Local2DTokenMixerConvNextBlock(self.nf, kernel_size=7, mlp_ratio=4.0, dropout=0.0, drop_path=0.10, layer_scale_init=1e-6),\n            Local2DTokenMixerConvNextBlock(self.nf, kernel_size=7, mlp_ratio=4.0, dropout=0.0, drop_path=0.10, layer_scale_init=1e-6),\n        )\n\n        self.gem2d = GeM2d(p=3.0)\n        self.feat_ln = nn.LayerNorm(self.nf)\n\n        def pos_head():\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf//2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf//2, 1),\n                nn.Softplus(),\n            )\n        self.head_1 = pos_head()  # green\n        self.head_2 = pos_head()  # clover\n        self.head_3 = pos_head()  # dead\n \n    @staticmethod\n    def _ensure_nchw(x: torch.Tensor, in_ch: int) -> torch.Tensor:\n        # If x is NHWC (B,H,W,C) convert to NCHW (B,C,H,W)\n        if x.ndim == 4 and x.shape[1] != in_ch and x.shape[-1] == in_ch:\n            return x.permute(0, 3, 1, 2).contiguous()\n        return x\n\n    def forward(self, left, right):\n        fl = self.backbone(left)\n        fr = self.backbone(right)\n\n        f8_l  = self._ensure_nchw(fl[self.idx8],  self.proj8[0].in_channels)\n        f16_l = self._ensure_nchw(fl[self.idx16], self.proj16[0].in_channels)\n        f32_l = self._ensure_nchw(fl[self.idx32], self.proj32[0].in_channels)\n    \n        f8_r  = self._ensure_nchw(fr[self.idx8],  self.proj8[0].in_channels)\n        f16_r = self._ensure_nchw(fr[self.idx16], self.proj16[0].in_channels)\n        f32_r = self._ensure_nchw(fr[self.idx32], self.proj32[0].in_channels)\n    \n        p8_l  = self.proj8(f8_l);   p8_r  = self.proj8(f8_r)\n        p16_l = self.proj16(f16_l); p16_r = self.proj16(f16_r)\n        p32_l = self.proj32(f32_l); p32_r = self.proj32(f32_r)\n\n\n        p8_l,  p8_r  = self.stereo8(p8_l,  p8_r)\n        p16_l, p16_r = self.stereo16(p16_l, p16_r)\n        p32_l, p32_r = self.stereo32(p32_l, p32_r)\n\n        p8  = torch.cat([p8_l,  p8_r],  dim=3)\n        p16 = torch.cat([p16_l, p16_r], dim=3)\n        p32 = torch.cat([p32_l, p32_r], dim=3)\n\n        p16_fused = self.fuse_to_p16(p8, p16, p32)\n        p16_fused = self.seam16(p16_fused)\n\n        grid = p16_fused.permute(0,2,3,1).contiguous()\n        grid = self.fusion2d(grid)\n        p16m = grid.permute(0,3,1,2).contiguous()\n\n        feat = self.gem2d(p16m)\n        feat = self.feat_ln(feat)\n\n        green  = self.head_1(feat)\n        clover = self.head_2(feat)\n        dead   = self.head_3(feat)\n        gdm    = green + clover\n        total  = gdm + dead\n        return total, gdm, green, clover, dead\n\n# ----------------- inference helpers -----------------\ndef load_sd(path):\n    s = torch.load(path, map_location=\"cpu\")\n    if isinstance(s, dict) and (\"model_state_dict\" in s or \"state_dict\" in s):\n        s = s.get(\"model_state_dict\", s.get(\"state_dict\"))\n    return s\n\ndef postprocess_3_to_5(pred3):\n    total = pred3[:,0]; gdm = pred3[:,1]; green = pred3[:,2]\n    clover = np.maximum(0.0, gdm - green)\n    dead   = np.maximum(0.0, total - gdm)\n    return np.stack([green, dead, clover, gdm, total], axis=1).astype(np.float32)\n\n@torch.inference_mode()\ndef predict_ckpt(model, loader, ckpt):\n    model.load_state_dict(load_sd(ckpt), strict=False)\n    model.to(DEVICE).eval()\n    out = 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                outs = model(l, r)\n        else:\n            outs = model(l, r)\n        total, gdm, green = outs[0].view(-1), outs[1].view(-1), outs[2].view(-1)\n        pred3 = torch.stack([total, gdm, green], dim=1).float().cpu().numpy()\n        b = l.size(0)\n        out[off:off+b] = pred3\n        off += b\n    return out\n\n# ----------------- run -----------------\ntest_df = pd.read_csv(TEST_CSV)\nuniq_imgs = test_df[\"image_path\"].drop_duplicates().values\n\nbase_dirs = [\n    \"/kaggle/input/csiro-biomass\",\n    \"/kaggle/input/csiro-biomass/test\",\n    \"/kaggle/input/csiro-biomass/test_images\",\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 checkpoints like best_model_fold*.pth in: {WEIGHTS_DIR}\")\n\nmodel = BiomassSwinV2(MODEL_NAME, pretrained=False, head_dropout=0.1)\npred3_sum = np.zeros((len(ds), 3), np.float32)\nfor p in ckpts:\n    pred3_sum += predict_ckpt(model, dl, p)\npred3 = pred3_sum / float(len(ckpts))\npred5 = postprocess_3_to_5(pred3)\n\npreds_wide = pd.DataFrame(pred5, 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}]}