{"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":725785,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":552452,"modelId":565013}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ===== ConvNeXtV2 ONLY — exact nb model/layers + aux heads + TTA + fold avg =====\nimport os, gc, cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\n\n# --- paths / settings (only what’s used) ---\nBASE_PATH      = \"/kaggle/input/csiro-biomass\"\nTEST_IMAGE_DIR = f\"{BASE_PATH}/test\"\nTEST_CSV       = f\"{BASE_PATH}/test.csv\"\n\nMODEL_DIR   = \"/kaggle/input/csiro-cnvnextv2/pytorch/default/1/convnextv2\"\nPREFIX      = \"convnextv2\"\nWEIGHTS_DIR = os.path.join(MODEL_DIR, PREFIX)\n\nMODEL_NAME   = \"convnextv2_large.fcmae_ft_in22k_in1k_384\"\nIMG_SIZE     = 384\nVAL_TTA_TIMES = 2\nN_FOLDS      = 5\n\nBATCH_SIZE   = 8\nNUM_WORKERS  = 2\n\n# reproduce nb inference: aux OFF (but aux heads exist in model class below)\nUSE_AUX_INFER = False          # <-- set True only if you want model() to also return aux dict\nNUM_STATES    = 0              # <-- only matters if USE_AUX_INFER=True and you want aux_state head\n\nALL_TARGET_COLS = [\"Dry_Green_g\",\"Dry_Dead_g\",\"Dry_Clover_g\",\"GDM_g\",\"Dry_Total_g\"]\nIMAGENET_MEAN = [0.485, 0.456, 0.406]\nIMAGENET_STD  = [0.229, 0.224, 0.225]\n\ncv2.setNumThreads(0)\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = True\ntorch.backends.cudnn.benchmark = True\ntorch.set_float32_matmul_precision(\"high\")\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\namp_dtype = torch.bfloat16 if (device.type == \"cuda\" and torch.cuda.is_bf16_supported()) else torch.float16\n\n# --- exact nb helpers ---\ndef clean_image(img):\n    h, w = img.shape[:2]\n    img = img[0:int(h * 0.90), :]\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 np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n    return img\n\ndef get_tta_transforms(img_size=512, mode=0, mean=IMAGENET_MEAN, std=IMAGENET_STD):\n    h = int(img_size) if not isinstance(img_size, (tuple, list)) else int(img_size[0])\n    w = int(img_size) if not isinstance(img_size, (tuple, list)) else int(img_size[1])\n    transforms_list = [A.Resize(h, w)]\n\n    if mode == 1:\n        transforms_list.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        transforms_list.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        transforms_list.append(\n            A.Rotate(limit=(90, 90), p=1.0,\n                     interpolation=cv2.INTER_LINEAR,\n                     border_mode=cv2.BORDER_REFLECT_101)\n        )\n\n    transforms_list += [\n        A.Normalize(mean=mean, std=std, max_pixel_value=255.0),\n        ToTensorV2(),\n    ]\n    return A.Compose(transforms_list, p=1.0)\n\n# --- layers used by ConvNeXtV2 feature-path (from nb) ---\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=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):\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\n    def forward(self, x):\n        x = x.clamp(min=self.eps).pow(self.p)\n        return x.mean(dim=(2, 3)).pow(1.0 / self.p)\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:\n            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\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:\n            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)\n        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        return xL + ls * (xL * gR), xR + ls * (xR * gL)\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\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)\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\n    def forward(self, x):\n        return x + self.pw(self.act(self.dw(x)))\n\n# --- BiomassModel feature-path + gcd + AUX heads exactly like nb ---\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=False, fuse_dim=None, use_aux=False, num_states=4, head_dropout=0.1):\n        super().__init__()\n        self.use_aux = bool(use_aux)\n        self.num_states = int(num_states)\n\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_idx(target):\n            return int(min(range(len(reds)), key=lambda i: abs(int(reds[i]) - int(target))))\n\n        self.idx8  = _pick_idx(8)\n        self.idx16 = _pick_idx(16)\n        self.idx32 = _pick_idx(32)\n\n        ch8  = int(chans[self.idx8])\n        ch16 = int(chans[self.idx16])\n        ch32 = 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:\n                g -= 1\n            return nn.Sequential(\n                nn.Conv2d(in_ch, self.nf, kernel_size=1, bias=False),\n                nn.GroupNorm(g, self.nf),\n                nn.GELU(),\n            )\n\n        self.proj8  = _proj(ch8)\n        self.proj16 = _proj(ch16)\n        self.proj32 = _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, eps=1e-6)\n\n        def make_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\n        # gcd: predict green/clover/dead -> derive gdm/total\n        self.head_1 = make_pos_head()\n        self.head_2 = make_pos_head()\n        self.head_3 = make_pos_head()\n\n        # --- AUX heads (exact nb) ---\n        def make_reg_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            )\n\n        def make_cls_head(n_classes):\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, n_classes),\n            )\n\n        self.aux_height = make_reg_head() if self.use_aux else None\n        self.aux_ndvi   = make_reg_head() if self.use_aux else None\n        self.aux_state  = make_cls_head(self.num_states) if (self.use_aux and self.num_states >= 2) else None\n\n    @staticmethod\n    def _ensure_nchw(x, in_ch):\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        feats_l = self.backbone(left)\n        feats_r = self.backbone(right)\n\n        f8_l  = self._ensure_nchw(feats_l[self.idx8],  self.proj8[0].in_channels)\n        f16_l = self._ensure_nchw(feats_l[self.idx16], self.proj16[0].in_channels)\n        f32_l = self._ensure_nchw(feats_l[self.idx32], self.proj32[0].in_channels)\n\n        f8_r  = self._ensure_nchw(feats_r[self.idx8],  self.proj8[0].in_channels)\n        f16_r = self._ensure_nchw(feats_r[self.idx16], self.proj16[0].in_channels)\n        f32_r = self._ensure_nchw(feats_r[self.idx32], self.proj32[0].in_channels)\n\n        p8_l,  p16_l,  p32_l  = self.proj8(f8_l),  self.proj16(f16_l),  self.proj32(f32_l)\n        p8_r,  p16_r,  p32_r  = self.proj8(f8_r),  self.proj16(f16_r),  self.proj32(f32_r)\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.seam16(self.fuse_to_p16(p8, p16, p32))\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.feat_ln(self.gem2d(p16m))\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\n        if not self.use_aux:\n            return total, gdm, green, clover, dead\n\n        aux = {\n            \"height\": self.aux_height(feat) if self.aux_height is not None else None,\n            \"Pre_GSHH_NDVI\": self.aux_ndvi(feat) if self.aux_ndvi is not None else None,\n            \"state_logits\": self.aux_state(feat) if self.aux_state is not None else None,\n        }\n        return total, gdm, green, clover, dead, aux\n\n# --- inference dataset/utils (nb) ---\nclass BiomassTestDataset(Dataset):\n    def __init__(self, img_dir):\n        self.paths = sorted([\n            os.path.join(img_dir, f) for f in os.listdir(img_dir)\n            if f.lower().endswith((\".jpg\", \".jpeg\", \".png\"))\n        ])\n        self.filenames = [os.path.basename(p) for p in self.paths]\n\n    def __len__(self): return len(self.paths)\n\n    def __getitem__(self, idx):\n        img = cv2.imread(self.paths[idx])\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        mid = img.shape[1] // 2\n        return img[:, :mid].copy(), img[:, mid:].copy(), self.filenames[idx]\n\ndef _collate_raw_np(batch):\n    lefts, rights, fns = zip(*batch)\n    return list(lefts), list(rights), list(fns)\n\ndef _postprocess_3_to_5(preds3):\n    pred_total = preds3[:, 0]\n    pred_gdm   = preds3[:, 1]\n    pred_green = preds3[:, 2]\n    pred_clover = np.maximum(0.0, pred_gdm - pred_green)\n    pred_dead   = np.maximum(0.0, pred_total - pred_gdm)\n    return np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1).astype(np.float32)\n\ndef _make_submission_from_preds(test_csv_path, all_target_cols, preds_5, filenames, out_csv_path):\n    test_df = pd.read_csv(test_csv_path)\n    ex = test_df[\"image_path\"].iloc[0]\n    if \"/\" in ex:\n        prefix = ex.rsplit(\"/\", 1)[0] + \"/\"\n        filenames = [prefix + fn for fn in filenames]\n\n    wide = pd.DataFrame(preds_5, columns=all_target_cols)\n    wide.insert(0, \"image_path\", filenames)\n\n    long = wide.melt(id_vars=[\"image_path\"], var_name=\"target_name\", value_name=\"target\")\n    sub = test_df[[\"sample_id\", \"image_path\", \"target_name\"]].merge(\n        long, on=[\"image_path\", \"target_name\"], how=\"left\"\n    )[[\"sample_id\", \"target\"]]\n\n    sub[\"target\"] = sub[\"target\"].fillna(0.0)\n    sub = sub.sort_values(\"sample_id\").reset_index(drop=True)\n    sub.to_csv(out_csv_path, index=False)\n    return sub\n\n@torch.inference_mode()\ndef predict_fold_tta(model, loader, img_size, tta_steps=1):\n    N = len(loader.dataset)\n    preds_sum = np.zeros((N, 3), dtype=np.float32)\n\n    tta_steps = int(max(1, min(int(tta_steps), 4)))\n    tfms = [get_tta_transforms(img_size, mode=m, mean=IMAGENET_MEAN, std=IMAGENET_STD) for m in range(tta_steps)]\n    nb = (device.type == \"cuda\")\n\n    for tfm in tfms:\n        offset = 0\n        for left_list, right_list, _ in loader:\n            bs = len(left_list)\n            lt = torch.stack([tfm(image=im)[\"image\"] for im in left_list], 0).to(device, non_blocking=nb)\n            rt = torch.stack([tfm(image=im)[\"image\"] for im in right_list], 0).to(device, non_blocking=nb)\n\n            if nb:\n                with torch.autocast(\"cuda\", dtype=amp_dtype):\n                    outs = model(lt, rt)\n            else:\n                outs = model(lt, rt)\n\n            # works for (..,dead) and (..,dead,aux): only need first 3\n            total, gdm, green = outs[0].view(-1), outs[1].view(-1), outs[2].view(-1)\n            preds_sum[offset:offset+bs] += torch.stack([total, gdm, green], 1).float().cpu().numpy()\n            offset += bs\n\n            del lt, rt, outs, total, gdm, green\n\n    preds_sum /= float(tta_steps)\n    return preds_sum.astype(np.float32)\n\ndef load_model_for_fold(weight_path):\n    model = BiomassModel(\n        MODEL_NAME,\n        pretrained=False,\n        fuse_dim=None,\n        use_aux=USE_AUX_INFER,\n        num_states=NUM_STATES,\n        head_dropout=0.1,\n    )\n    state = torch.load(weight_path, map_location=\"cpu\")\n    if isinstance(state, dict) and (\"model_state_dict\" in state or \"state_dict\" in state):\n        sd = state[\"model_state_dict\"] if \"model_state_dict\" in state else state[\"state_dict\"]\n    else:\n        sd = state\n    model.load_state_dict(sd, strict=False)  # nb\n    return model\n\n# --- run convnextv2 folds + TTA ---\nds = BiomassTestDataset(TEST_IMAGE_DIR)\ndl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS,\n                pin_memory=(device.type == \"cuda\"), persistent_workers=(NUM_WORKERS > 0),\n                prefetch_factor=4 if NUM_WORKERS > 0 else None, collate_fn=_collate_raw_np)\n\npreds3_sum = np.zeros((len(ds), 3), np.float32)\nused = 0\n\nfor fold in range(N_FOLDS):\n    wp = os.path.join(WEIGHTS_DIR, f\"best_model_fold{fold}.pth\")\n    if not os.path.exists(wp):\n        continue\n    m = load_model_for_fold(wp).to(device).eval()\n    preds3_sum += predict_fold_tta(m, dl, IMG_SIZE, tta_steps=VAL_TTA_TIMES)\n    used += 1\n    del m\n    if device.type == \"cuda\":\n        torch.cuda.empty_cache()\n    gc.collect()\n\nif used == 0:\n    raise FileNotFoundError(f\"No weights found in {WEIGHTS_DIR}\")\n\npreds3 = preds3_sum / float(used)\npreds5 = _postprocess_3_to_5(preds3)\n\nsub = _make_submission_from_preds(TEST_CSV, ALL_TARGET_COLS, preds5, ds.filenames, \"submission.csv\")\nprint(f\"✓ submission.csv | folds={used} | shape={sub.shape}\")\nsub.head()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}