{"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":663314,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":471723,"modelId":487624}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================================================================\n# CSIRO Image2Biomass - v4 CrossPVT T2T Mamba Inference (Self-contained)\n# -----------------------------------------------------------------------------\n# - 复刻 train.py 中的 CrossPVT_T2T_MambaDINO 结构\n# - 直接从 checkpoint 中读取 cfg 覆盖本地 CFG，保证结构一致\n# - 5-fold ensemble + TTA (原图/水平翻转/垂直翻转)\n# - 自动处理 DataParallel 的 module. 前缀\n# - 使用 parse_known_args() 规避 Kaggle/Colab 多余命令行参数\n# - 输出 submission.csv\n# =============================================================================\n\nimport os\nimport gc\nimport math\nimport argparse\nimport logging\nfrom dataclasses import dataclass, field\nfrom typing import List, Tuple, Dict, Optional\n\nimport cv2\nimport timm\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# =============================================================================\n# Logger\n# =============================================================================\nLOGGER = logging.getLogger(\"csiro_infer_v4\")\nif not LOGGER.handlers:\n    LOGGER.addHandler(logging.StreamHandler())\nLOGGER.setLevel(logging.INFO)\n\n\n# =============================================================================\n# 训练时使用的配置（会被 checkpoint 中的 cfg 覆盖）\n# =============================================================================\n@dataclass\nclass TrainCFG:\n    dropout: float = 0.1\n    hidden_ratio: float = 0.35\n\n    # DINO backbone 相关\n    dino_candidates: Tuple[str, ...] = (\n        \"vit_base_patch14_dinov2\",\n        \"vit_base_patch14_reg4_dinov2\",\n        \"vit_small_patch14_dinov2\",\n    )\n    small_grid: Tuple[int, int] = (4, 4)\n    big_grid: Tuple[int, int] = (2, 2)\n    t2t_depth: int = 2\n    cross_layers: int = 2\n    cross_heads: int = 6\n\n    # Pyramid + Mamba\n    pyramid_dims: Tuple[int, int, int] = (384, 512, 640)\n    mobilevit_heads: int = 4\n    mobilevit_depth: int = 2\n    sra_heads: int = 8\n    sra_ratio: int = 2\n    mamba_depth: int = 3\n    mamba_kernel: int = 5\n    aux_head: bool = True\n    aux_loss_weight: float = 0.4\n\n    # 目标列（用于辅助头 / 5 目标打包）\n    ALL_TARGET_COLS: Tuple[str, ...] = (\n        \"Dry_Green_g\",\n        \"Dry_Dead_g\",\n        \"Dry_Clover_g\",\n        \"GDM_g\",\n        \"Dry_Total_g\",\n    )\n\n\nCFG = TrainCFG()\n\n\ndef update_cfg_from_checkpoint(cfg_dict: dict):\n    \"\"\"\n    用 checkpoint 中保存的 cfg 覆盖当前 CFG 的同名字段，\n    确保推理结构与训练完全一致。\n    \"\"\"\n    global CFG\n    if not cfg_dict:\n        return\n    for k, v in cfg_dict.items():\n        if hasattr(CFG, k):\n            setattr(CFG, k, v)\n\n\n# =============================================================================\n# Model blocks（与 train.py 保持一致的简化实现）\n# =============================================================================\nclass FeedForward(nn.Module):\n    def __init__(self, dim, mlp_ratio=4.0, dropout=0.0):\n        super().__init__()\n        hid = int(dim * mlp_ratio)\n        self.net = nn.Sequential(\n            nn.Linear(dim, hid),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hid, dim),\n            nn.Dropout(dropout),\n        )\n\n    def forward(self, x):\n        return self.net(x)\n\n\nclass AttentionBlock(nn.Module):\n    def __init__(self, dim, heads=8, dropout=0.0, mlp_ratio=4.0):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)\n        self.norm2 = nn.LayerNorm(dim)\n        self.ff = FeedForward(dim, mlp_ratio=mlp_ratio, dropout=dropout)\n\n    def forward(self, x):\n        h = self.norm1(x)\n        attn_out, _ = self.attn(h, h, h, need_weights=False)\n        x = x + attn_out\n        x = x + self.ff(self.norm2(x))\n        return x\n\n\nclass MobileViTBlock(nn.Module):\n    \"\"\"\n    轻量 MobileViT：局部 CNN + 小型 Transformer（token 化和 fold back）\n    \"\"\"\n\n    def __init__(self, dim, heads=4, depth=2, patch=(2, 2), dropout=0.0):\n        super().__init__()\n        self.local = nn.Sequential(\n            nn.Conv2d(dim, dim, 3, padding=1, groups=dim),\n            nn.Conv2d(dim, dim, 1),\n            nn.GELU(),\n        )\n        self.patch = patch\n        self.transformer = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(depth)]\n        )\n        self.fuse = nn.Conv2d(dim * 2, dim, kernel_size=1)\n\n    def forward(self, x: torch.Tensor):\n        local_feat = self.local(x)\n        B, C, H, W = local_feat.shape\n        ph, pw = self.patch\n        new_h = math.ceil(H / ph) * ph\n        new_w = math.ceil(W / pw) * pw\n        if new_h != H or new_w != W:\n            local_feat = F.interpolate(local_feat, size=(new_h, new_w), mode=\"bilinear\", align_corners=False)\n            H, W = new_h, new_w\n\n        tokens = local_feat.unfold(2, ph, ph).unfold(3, pw, pw)  # B,C,nh,nw,ph,pw\n        tokens = tokens.contiguous().view(B, C, -1, ph, pw)\n        tokens = tokens.permute(0, 2, 3, 4, 1).reshape(B, -1, C)\n\n        for blk in self.transformer:\n            tokens = blk(tokens)\n\n        feat = tokens.view(B, -1, ph * pw, C).permute(0, 3, 1, 2)\n        nh = H // ph\n        nw = W // pw\n        feat = feat.view(B, C, nh, nw, ph, pw).permute(0, 1, 2, 4, 3, 5)\n        feat = feat.reshape(B, C, H, W)\n\n        if feat.shape[-2:] != x.shape[-2:]:\n            feat = F.interpolate(feat, size=x.shape[-2:], mode=\"bilinear\", align_corners=False)\n\n        out = self.fuse(torch.cat([x, feat], dim=1))\n        return out\n\n\nclass SpatialReductionAttention(nn.Module):\n    def __init__(self, dim, heads=8, sr_ratio=2, dropout=0.0):\n        super().__init__()\n        self.heads = heads\n        self.scale = (dim // heads) ** -0.5\n        self.q = nn.Linear(dim, dim)\n        self.kv = nn.Linear(dim, dim * 2)\n        self.sr_ratio = sr_ratio\n        if sr_ratio > 1:\n            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)\n            self.norm = nn.LayerNorm(dim)\n        else:\n            self.sr = None\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x, hw: Tuple[int, int]):\n        B, N, C = x.shape\n        q = self.q(x).reshape(B, N, self.heads, C // self.heads).permute(0, 2, 1, 3)\n\n        if self.sr is not None:\n            H, W = hw\n            feat = x.transpose(1, 2).reshape(B, C, H, W)\n            feat = self.sr(feat)\n            feat = feat.reshape(B, C, -1).transpose(1, 2)\n            feat = self.norm(feat)\n        else:\n            feat = x\n\n        kv = self.kv(feat)\n        k, v = kv.chunk(2, dim=-1)\n        k = k.reshape(B, -1, self.heads, C // self.heads).permute(0, 2, 3, 1)\n        v = v.reshape(B, -1, self.heads, C // self.heads).permute(0, 2, 1, 3)\n\n        attn = torch.matmul(q, k) * self.scale\n        attn = attn.softmax(dim=-1)\n        attn = self.drop(attn)\n        out = torch.matmul(attn, v).permute(0, 2, 1, 3).reshape(B, N, C)\n        out = self.proj(out)\n        return out\n\n\nclass PVTBlock(nn.Module):\n    def __init__(self, dim, heads=8, sr_ratio=2, dropout=0.0, mlp_ratio=4.0):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.sra = SpatialReductionAttention(dim, heads=heads, sr_ratio=sr_ratio, dropout=dropout)\n        self.norm2 = nn.LayerNorm(dim)\n        self.ff = FeedForward(dim, mlp_ratio=mlp_ratio, dropout=dropout)\n\n    def forward(self, x, hw: Tuple[int, int]):\n        x = x + self.sra(self.norm1(x), hw)\n        x = x + self.ff(self.norm2(x))\n        return x\n\n\nclass LocalMambaBlock(nn.Module):\n    \"\"\"\n    简化版 local Mamba：DW-Conv + gating + 线性映射\n    \"\"\"\n\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)\n        self.gate = nn.Linear(dim, dim)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x):\n        shortcut = x\n        x = self.norm(x)\n        g = torch.sigmoid(self.gate(x))\n        x = (x * g).transpose(1, 2)  # B, C, N\n        x = self.dwconv(x).transpose(1, 2)\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\nclass T2TRetokenizer(nn.Module):\n    \"\"\"\n    将 4x4 tile token 做局部 attention + 下采样到 2x2\n    \"\"\"\n\n    def __init__(self, dim, depth=2, heads=4, dropout=0.0):\n        super().__init__()\n        self.blocks = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(depth)]\n        )\n\n    def forward(self, tokens: torch.Tensor, grid_hw: Tuple[int, int]):\n        B, T, C = tokens.shape\n        H, W = grid_hw\n        feat_map = tokens.transpose(1, 2).reshape(B, C, H, W)\n        seq = feat_map.flatten(2).transpose(1, 2)\n        for blk in self.blocks:\n            seq = blk(seq)\n        seq_map = seq.transpose(1, 2).reshape(B, C, H, W)\n        pooled = F.adaptive_avg_pool2d(seq_map, (2, 2))\n        retokens = pooled.flatten(2).transpose(1, 2)\n        return retokens, seq_map\n\n\nclass CrossScaleFusion(nn.Module):\n    def __init__(self, dim, heads=6, dropout=0.0, layers=2):\n        super().__init__()\n        self.layers_s = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(layers)]\n        )\n        self.layers_b = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(layers)]\n        )\n        self.cross_s = nn.ModuleList(\n            [\n                nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True, kdim=dim, vdim=dim)\n                for _ in range(layers)\n            ]\n        )\n        self.cross_b = nn.ModuleList(\n            [\n                nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True, kdim=dim, vdim=dim)\n                for _ in range(layers)\n            ]\n        )\n        self.norm_s = nn.LayerNorm(dim)\n        self.norm_b = nn.LayerNorm(dim)\n\n    def forward(self, tok_s: torch.Tensor, tok_b: torch.Tensor):\n        B, Ts, C = tok_s.shape\n        Tb = tok_b.shape[1]\n        cls_s = tok_s.new_zeros(B, 1, C)\n        cls_b = tok_b.new_zeros(B, 1, C)\n        tok_s = torch.cat([cls_s, tok_s], dim=1)\n        tok_b = torch.cat([cls_b, tok_b], dim=1)\n\n        for ls, lb, cs, cb in zip(self.layers_s, self.layers_b, self.cross_s, self.cross_b):\n            tok_s = ls(tok_s)\n            tok_b = lb(tok_b)\n            q_s = self.norm_s(tok_s[:, :1])\n            q_b = self.norm_b(tok_b[:, :1])\n            cls_s_upd, _ = cs(\n                q_s,\n                torch.cat([tok_b, q_b], dim=1),\n                torch.cat([tok_b, q_b], dim=1),\n                need_weights=False,\n            )\n            cls_b_upd, _ = cb(\n                q_b,\n                torch.cat([tok_s, q_s], dim=1),\n                torch.cat([tok_s, q_s], dim=1),\n                need_weights=False,\n            )\n            tok_s = torch.cat([tok_s[:, :1] + cls_s_upd, tok_s[:, 1:]], dim=1)\n            tok_b = torch.cat([tok_b[:, :1] + cls_b_upd, tok_b[:, 1:]], dim=1)\n\n        tokens = torch.cat([tok_s[:, :1], tok_b[:, :1], tok_s[:, 1:], tok_b[:, 1:]], dim=1)\n        return tokens  # shape ~ (B, 2 + Ts + Tb, C)\n\n\nclass TileEncoder(nn.Module):\n    def __init__(self, backbone: nn.Module, input_res: int):\n        super().__init__()\n        self.backbone = backbone\n        self.input_res = input_res\n\n    def forward(self, x: torch.Tensor, grid: Tuple[int, int]):\n        B, C, H, W = x.shape\n        r, c = grid\n        hs = torch.linspace(0, H, steps=r + 1, device=x.device).round().long()\n        ws = torch.linspace(0, W, steps=c + 1, device=x.device).round().long()\n        tiles = []\n        for i in range(r):\n            for j in range(c):\n                rs, re = hs[i].item(), hs[i + 1].item()\n                cs, ce = ws[j].item(), ws[j + 1].item()\n                xt = x[:, :, rs:re, cs:ce]\n                if xt.shape[-2:] != (self.input_res, self.input_res):\n                    xt = F.interpolate(xt, size=(self.input_res, self.input_res), mode=\"bilinear\", align_corners=False)\n                tiles.append(xt)\n        tiles = torch.stack(tiles, dim=1)  # (B, T, C, H, W)\n        flat = tiles.view(-1, C, self.input_res, self.input_res)\n        feats = self.backbone(flat)\n        feats = feats.view(B, -1, feats.shape[-1])\n        return feats\n\n\nclass PyramidMixer(nn.Module):\n    def __init__(\n        self,\n        dim_in: int,\n        dims: Tuple[int, int, int],\n        mobilevit_heads: int = 4,\n        mobilevit_depth: int = 2,\n        sra_heads: int = 6,\n        sra_ratio: int = 2,\n        mamba_depth: int = 3,\n        mamba_kernel: int = 5,\n        dropout: float = 0.0,\n    ):\n        super().__init__()\n        c1, c2, c3 = dims\n        self.proj1 = nn.Linear(dim_in, c1)\n        self.mobilevit = MobileViTBlock(c1, heads=mobilevit_heads, depth=mobilevit_depth, dropout=dropout)\n        self.proj2 = nn.Linear(c1, c2)\n        self.pvt = PVTBlock(c2, heads=sra_heads, sr_ratio=sra_ratio, dropout=dropout, mlp_ratio=3.0)\n        self.mamba_local = LocalMambaBlock(c2, kernel_size=mamba_kernel, dropout=dropout)\n        self.proj3 = nn.Linear(c2, c3)\n        self.mamba_global = nn.ModuleList(\n            [LocalMambaBlock(c3, kernel_size=mamba_kernel, dropout=dropout) for _ in range(mamba_depth)]\n        )\n        self.final_attn = AttentionBlock(c3, heads=min(8, c3 // 64 + 1), dropout=dropout, mlp_ratio=2.0)\n\n    def _tokens_to_map(self, tokens: torch.Tensor, target_hw: Tuple[int, int]):\n        B, N, C = tokens.shape\n        H, W = target_hw\n        need = H * W\n        if N < need:\n            pad = tokens.new_zeros(B, need - N, C)\n            tokens = torch.cat([tokens, pad], dim=1)\n        tokens = tokens[:, :need, :]\n        feat_map = tokens.transpose(1, 2).reshape(B, C, H, W)\n        return feat_map\n\n    @staticmethod\n    def _fit_hw(n_tokens: int) -> Tuple[int, int]:\n        \"\"\"选择一个接近方形、满足 h*w>=n_tokens 的网格。\"\"\"\n        h = int(math.sqrt(n_tokens))\n        w = h\n        while h * w < n_tokens:\n            w += 1\n            if h * w < n_tokens:\n                h += 1\n        return h, w\n\n    def forward(self, tokens: torch.Tensor):\n        # 约 10 tokens -> 3x4 map\n        B, N, C = tokens.shape\n        map_hw = (3, 4)\n        feat_map = self._tokens_to_map(tokens, map_hw)\n\n        t1 = self.proj1(tokens)\n        m1 = self._tokens_to_map(t1, map_hw)\n        m1 = self.mobilevit(m1)\n        t1_out = m1.flatten(2).transpose(1, 2)[:, :N]\n\n        # Stage2: 下采样 token 数量（平均池化）\n        t2 = self.proj2(t1_out)\n        new_len = max(4, N // 2)\n        t2 = t2[:, :new_len] + F.adaptive_avg_pool1d(t2.transpose(1, 2), new_len).transpose(1, 2)\n        hw2 = self._fit_hw(t2.size(1))\n        if t2.size(1) < hw2[0] * hw2[1]:\n            pad = t2.new_zeros(B, hw2[0] * hw2[1] - t2.size(1), t2.size(2))\n            t2 = torch.cat([t2, pad], dim=1)\n        t2 = self.pvt(t2, hw2)\n        t2 = self.mamba_local(t2)\n\n        # Stage3: 全局\n        t3 = self.proj3(t2)\n        pooled = torch.stack([t3.mean(dim=1), t3.max(dim=1).values], dim=1)  # (B,2,C)\n        t3 = pooled\n        for blk in self.mamba_global:\n            t3 = blk(t3)\n        t3 = self.final_attn(t3)\n        global_feat = t3.mean(dim=1)\n        return global_feat, {\"stage1_map\": m1.detach(), \"stage2_tokens\": t2.detach(), \"stage3_tokens\": t3.detach()}\n\n\nclass CrossPVT_T2T_MambaDINO(nn.Module):\n    def __init__(self, dropout: float = 0.1, hidden_ratio: float = 0.35):\n        super().__init__()\n        self.backbone, self.feat_dim, self.backbone_name, self.input_res = self._build_dino_backbone()\n        self.tile_encoder = TileEncoder(self.backbone, self.input_res)\n        self.t2t = T2TRetokenizer(self.feat_dim, depth=CFG.t2t_depth, heads=CFG.cross_heads, dropout=dropout)\n        self.cross = CrossScaleFusion(\n            self.feat_dim, heads=CFG.cross_heads, dropout=dropout, layers=CFG.cross_layers\n        )\n        self.pyramid = PyramidMixer(\n            dim_in=self.feat_dim,\n            dims=CFG.pyramid_dims,\n            mobilevit_heads=CFG.mobilevit_heads,\n            mobilevit_depth=CFG.mobilevit_depth,\n            sra_heads=CFG.sra_heads,\n            sra_ratio=CFG.sra_ratio,\n            mamba_depth=CFG.mamba_depth,\n            mamba_kernel=CFG.mamba_kernel,\n            dropout=dropout,\n        )\n\n        combined = CFG.pyramid_dims[-1] * 2\n        self.combined_dim = combined\n        hidden = max(32, int(combined * hidden_ratio))\n\n        def head():\n            return nn.Sequential(\n                nn.Linear(combined, hidden),\n                nn.GELU(),\n                nn.Dropout(dropout),\n                nn.Linear(hidden, 1),\n            )\n\n        self.head_green = head()\n        self.head_clover = head()\n        self.head_dead = head()\n        self.score_head = nn.Sequential(nn.LayerNorm(combined), nn.Linear(combined, 1))\n        self.aux_head = (\n            nn.Sequential(nn.LayerNorm(CFG.pyramid_dims[1]), nn.Linear(CFG.pyramid_dims[1], 5))\n            if CFG.aux_head\n            else None\n        )\n        self.softplus = nn.Softplus(beta=1.0)\n\n        self.cross_gate_left = nn.Linear(CFG.pyramid_dims[-1], CFG.pyramid_dims[-1])\n        self.cross_gate_right = nn.Linear(CFG.pyramid_dims[-1], CFG.pyramid_dims[-1])\n\n    def _build_dino_backbone(self):\n        \"\"\"\n        只创建骨干结构，不加载预训练权重（pretrained=False），\n        避免在 Kaggle 无网环境下下载；权重由 checkpoint 提供。\n        \"\"\"\n        last_err = None\n        for name in CFG.dino_candidates:\n            for gp in [\"token\", \"avg\", \"__default__\"]:\n                try:\n                    if gp == \"__default__\":\n                        m = timm.create_model(name, pretrained=False, num_classes=0)\n                        gp_str = \"default\"\n                    else:\n                        m = timm.create_model(name, pretrained=False, num_classes=0, global_pool=gp)\n                        gp_str = gp\n                    feat = m.num_features\n                    input_res = self._infer_input_res(m)\n                    LOGGER.info(\n                        f\"✅ 使用 DINO 主干: {name} | global_pool={gp_str} | \"\n                        f\"feat_dim={feat} | input_res={input_res}\"\n                    )\n                    if hasattr(m, \"set_grad_checkpointing\"):\n                        m.set_grad_checkpointing(True)\n                    return m, feat, name, int(input_res)\n                except Exception as e:\n                    last_err = e\n                    continue\n        raise RuntimeError(f\"无法创建任何 DINO 主干。最后错误: {last_err}\")\n\n    @staticmethod\n    def _infer_input_res(m) -> int:\n        if hasattr(m, \"patch_embed\") and hasattr(m.patch_embed, \"img_size\"):\n            isz = m.patch_embed.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        if hasattr(m, \"img_size\"):\n            isz = m.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        dc = getattr(m, \"default_cfg\", {}) or {}\n        ins = dc.get(\"input_size\", None)\n        if ins:\n            if isinstance(ins, (tuple, list)) and len(ins) >= 2:\n                return int(ins[1])\n            return int(ins if isinstance(ins, (int, float)) else 224)\n        return 518\n\n    def _half_forward(self, x_half: torch.Tensor):\n        tiles_small = self.tile_encoder(x_half, CFG.small_grid)\n        tiles_big = self.tile_encoder(x_half, CFG.big_grid)\n        t2, stage1_map = self.t2t(tiles_small, CFG.small_grid)\n        fused = self.cross(t2, tiles_big)\n        feat, feat_maps = self.pyramid(fused)\n        feat_maps[\"stage1_map\"] = stage1_map\n        return feat, feat_maps\n\n    def _merge_heads(self, f_l: torch.Tensor, f_r: torch.Tensor):\n        g_l = torch.sigmoid(self.cross_gate_left(f_r))\n        g_r = torch.sigmoid(self.cross_gate_right(f_l))\n        f_l = f_l * g_l\n        f_r = f_r * g_r\n        f = torch.cat([f_l, f_r], dim=1)\n        green_pos = self.softplus(self.head_green(f))\n        clover_pos = self.softplus(self.head_clover(f))\n        dead_pos = self.softplus(self.head_dead(f))\n        gdm = green_pos + clover_pos\n        total = gdm + dead_pos\n        return total, gdm, green_pos, f\n\n    def _param_device_dtype(self):\n        try:\n            ref = next(self.parameters())\n            return ref.device, ref.dtype\n        except StopIteration:\n            return torch.device(\"cpu\"), torch.float32\n\n    def _empty_forward_output(self, device=None, dtype=None, return_features: bool = False):\n        if device is None or dtype is None:\n            device_p, dtype_p = self._param_device_dtype()\n            if device is None:\n                device = device_p\n            if dtype is None:\n                dtype = dtype_p\n\n        zero = torch.zeros(0, 1, device=device, dtype=dtype)\n        out = {\n            \"total\": zero,\n            \"gdm\": zero,\n            \"green\": zero,\n            \"score_feat\": torch.zeros(0, self.combined_dim, device=device, dtype=dtype),\n        }\n        if self.aux_head is not None:\n            out[\"aux\"] = torch.zeros(0, len(CFG.ALL_TARGET_COLS), device=device, dtype=dtype)\n        if return_features:\n            out[\"feature_maps\"] = {}\n        return out\n\n    def forward(self, *inputs, x_left=None, x_right=None, return_features: bool = False):\n        # 兼容 DataParallel 传参方式（单 tensor / 元组）\n        if inputs:\n            if len(inputs) == 1:\n                first = inputs[0]\n                if isinstance(first, (tuple, list)):\n                    if len(first) >= 1:\n                        x_left = first[0]\n                    if len(first) >= 2:\n                        x_right = first[1]\n                else:\n                    x_left = first\n            else:\n                x_left = inputs[0]\n                x_right = inputs[1]\n\n        if x_left is None:\n            return self._empty_forward_output(return_features=return_features)\n        if isinstance(x_left, torch.Tensor) and x_left.shape[0] == 0:\n            return self._empty_forward_output(return_features=return_features)\n\n        if x_right is None:\n            if isinstance(x_left, torch.Tensor):\n                if x_left.shape[1] % 2 != 0:\n                    raise ValueError(\"无法从单个张量推断左右分支，请显式提供 x_right。\")\n                x_left, x_right = torch.chunk(x_left, 2, dim=1)\n            else:\n                raise ValueError(\"缺少 x_right 输入。\")\n\n        feat_l, feats_l = self._half_forward(x_left)\n        feat_r, feats_r = self._half_forward(x_right)\n        total, gdm, green, f_concat = self._merge_heads(feat_l, feat_r)\n\n        out = {\n            \"total\": total,\n            \"gdm\": gdm,\n            \"green\": green,\n            \"score_feat\": f_concat,\n        }\n        if self.aux_head is not None:\n            aux_tokens = torch.cat([feats_l[\"stage2_tokens\"], feats_r[\"stage2_tokens\"]], dim=1)\n            aux_pred = self.softplus(self.aux_head(aux_tokens.mean(dim=1)))\n            out[\"aux\"] = aux_pred  # 顺序与 CFG.ALL_TARGET_COLS 对齐\n        if return_features:\n            out[\"feature_maps\"] = {\n                \"stage1_left\": feats_l.get(\"stage1_map\"),\n                \"stage1_right\": feats_r.get(\"stage1_map\"),\n                \"stage3_left\": feats_l.get(\"stage3_tokens\"),\n                \"stage3_right\": feats_r.get(\"stage3_tokens\"),\n            }\n        return out\n\n\n# =============================================================================\n# 推理配置\n# =============================================================================\nclass INF_CFG:\n    # 数据路径（Kaggle 默认，本地测试可修改）\n    BASE_PATH = \"/kaggle/input/csiro-biomass\"\n    TEST_CSV = os.path.join(BASE_PATH, \"test.csv\")\n    TEST_IMAGE_DIR = os.path.join(BASE_PATH, \"test\")\n\n    # 实验目录（你的权重所在目录）\n    EXPERIMENT_DIR = \"/kaggle/input/csiro/pytorch/default/12\"\n\n    # Checkpoint 路径（5-fold）\n    CKPT_PATTERN_FOLD_X = os.path.join(\n        EXPERIMENT_DIR, \"fold_{fold}\", \"checkpoints\", \"best_wr2.pt\"\n    )\n    CKPT_PATTERN_FOLDX = os.path.join(\n        EXPERIMENT_DIR, \"fold{fold}\", \"checkpoints\", \"best_wr2.pt\"\n    )\n    N_FOLDS = 5\n\n    # 输出\n    SUBMISSION_FILE = \"submission.csv\"\n\n    # 推理设置\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    BATCH_SIZE = 1\n    NUM_WORKERS = 0\n    MIXED_PRECISION = True\n\n    # TTA 设置\n    USE_TTA = True\n    TTA_TRANSFORMS = [\"original\", \"hflip\", \"vflip\"]\n\n    # 目标列顺序（与训练一致）\n    ALL_TARGET_COLS = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n\n\nprint(f\"Device: {INF_CFG.DEVICE}\")\nprint(f\"Experiment Dir: {INF_CFG.EXPERIMENT_DIR}\")\n\n\n# =============================================================================\n# 数据集\n# =============================================================================\nclass TestBiomassDataset(Dataset):\n    \"\"\"测试数据集：左右两路输入\"\"\"\n\n    def __init__(self, df: pd.DataFrame, transform, image_dir: str):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.image_dir = image_dir\n        self.paths = self.df[\"image_path\"].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        filename = os.path.basename(self.paths[idx])\n        full_path = os.path.join(self.image_dir, filename)\n\n        img = cv2.imread(full_path)\n        if img is None:\n            # 容错：若读图失败，用黑图占位\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # 左右切半（与训练一致）\n        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n\n        left_t = self.transform(image=left)[\"image\"]\n        right_t = self.transform(image=right)[\"image\"]\n\n        return left_t, right_t\n\n\n# =============================================================================\n# TTA 变换\n# =============================================================================\ndef get_tta_transforms(img_size: int) -> List[A.Compose]:\n    \"\"\"生成 TTA 变换列表\"\"\"\n    base = [\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2(),\n    ]\n\n    transforms = []\n\n    # 原图\n    transforms.append(\n        A.Compose([\n            A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA),\n            *base,\n        ])\n    )\n\n    # 水平翻转\n    transforms.append(\n        A.Compose([\n            A.HorizontalFlip(p=1.0),\n            A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA),\n            *base,\n        ])\n    )\n\n    # 垂直翻转\n    transforms.append(\n        A.Compose([\n            A.VerticalFlip(p=1.0),\n            A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA),\n            *base,\n        ])\n    )\n\n    return transforms\n\n\n# =============================================================================\n# 权重加载\n# =============================================================================\ndef strip_module_prefix(state_dict: dict) -> dict:\n    \"\"\"移除 DataParallel 的 module. 前缀\"\"\"\n    if not state_dict:\n        return state_dict\n\n    keys = list(state_dict.keys())\n    if all(k.startswith(\"module.\") for k in keys):\n        return {k[len(\"module.\"):]: v for k, v in state_dict.items()}\n    return state_dict\n\n\ndef load_checkpoint(path: str) -> dict:\n    \"\"\"加载 checkpoint\"\"\"\n    if not os.path.exists(path):\n        raise FileNotFoundError(f\"Checkpoint not found: {path}\")\n\n    try:\n        state = torch.load(path, map_location=\"cpu\", weights_only=False)\n    except TypeError:\n        state = torch.load(path, map_location=\"cpu\")\n\n    return state\n\n\ndef load_model_from_checkpoint(ckpt_path: str) -> nn.Module:\n    \"\"\"从 checkpoint 加载模型\"\"\"\n    print(f\"\\n加载 checkpoint: {os.path.basename(ckpt_path)}\")\n\n    state = load_checkpoint(ckpt_path)\n\n    # 用 checkpoint 中的 cfg 覆盖当前 CFG（保证结构对应）\n    cfg_dict = state.get(\"cfg\", {})\n    update_cfg_from_checkpoint(cfg_dict)\n\n    dropout = cfg_dict.get(\"dropout\", CFG.dropout)\n    hidden_ratio = cfg_dict.get(\"hidden_ratio\", CFG.hidden_ratio)\n\n    # 创建模型（结构与训练一致）\n    model = CrossPVT_T2T_MambaDINO(dropout=dropout, hidden_ratio=hidden_ratio)\n\n    # 提取模型状态\n    model_state = state.get(\"model_state\")\n    if model_state is None:\n        # 如果 checkpoint 直接是 state_dict\n        model_state = state\n\n    # 移除 module. 前缀\n    model_state = strip_module_prefix(model_state)\n\n    # 加载权重\n    missing_keys, unexpected_keys = model.load_state_dict(model_state, strict=False)\n\n    if missing_keys:\n        print(f\"  ⚠️  缺失的键: {len(missing_keys)} 个\")\n        if len(missing_keys) <= 10:\n            for k in missing_keys[:10]:\n                print(f\"    - {k}\")\n\n    if unexpected_keys:\n        print(f\"  ⚠️  意外的键: {len(unexpected_keys)} 个\")\n        if len(unexpected_keys) <= 10:\n            for k in unexpected_keys[:10]:\n                print(f\"    - {k}\")\n\n    model.to(INF_CFG.DEVICE)\n    model.eval()\n\n    # 获取输入分辨率\n    input_res = getattr(model, \"input_res\", 518)\n    backbone_name = getattr(model, \"backbone_name\", \"unknown\")\n\n    print(f\"  ✓ 模型加载成功 | backbone={backbone_name} | input_res={input_res}\")\n\n    return model\n\n\n# =============================================================================\n# 推理\n# =============================================================================\ndef pack5_targets(total: torch.Tensor, gdm: torch.Tensor, green: torch.Tensor) -> torch.Tensor:\n    \"\"\"将 total, gdm, green 打包为 5 个目标\"\"\"\n    clover = gdm - green\n    dead = total - gdm\n    return torch.cat([green, dead, clover, gdm, total], dim=1)\n\n\n@torch.no_grad()\ndef predict_one_view(models: List[nn.Module], loader: DataLoader) -> np.ndarray:\n    \"\"\"对单个 TTA 视角进行预测\"\"\"\n    preds_list = []\n    amp_dtype = \"cuda\" if INF_CFG.DEVICE.type == \"cuda\" else \"cpu\"\n\n    for xl, xr in tqdm(loader, desc=\"  Predicting\", leave=False):\n        xl = xl.to(INF_CFG.DEVICE, non_blocking=True)\n        xr = xr.to(INF_CFG.DEVICE, non_blocking=True)\n\n        # 拼接为单 tensor（与训练时的 DataParallel 调用方式一致）\n        x_cat = torch.cat([xl, xr], dim=1)\n\n        per_model_preds = []\n\n        with torch.amp.autocast(amp_dtype, enabled=INF_CFG.MIXED_PRECISION):\n            for model in models:\n                out = model(x_cat, return_features=False)\n\n                total = out[\"total\"]\n                gdm = out[\"gdm\"]\n                green = out[\"green\"]\n\n                # 打包为 5 个目标\n                five = pack5_targets(total, gdm, green)\n\n                # 非负约束\n                five = torch.clamp(five, min=0.0)\n\n                per_model_preds.append(five.float().cpu())\n\n        # 5-fold ensemble 平均\n        stacked = torch.mean(torch.stack(per_model_preds, dim=0), dim=0)\n        preds_list.append(stacked.numpy())\n\n    return np.concatenate(preds_list, axis=0)\n\n\ndef run_inference(test_df: pd.DataFrame, image_dir: str) -> np.ndarray:\n    \"\"\"运行完整推理流程（5-fold ensemble + TTA）\"\"\"\n    print(\"\\n\" + \"=\" * 80)\n    print(\"开始推理\")\n    print(\"=\" * 80)\n\n    # 加载所有 fold 的模型\n    print(\"\\n加载模型 (5-fold)...\")\n    models = []\n    input_res = None\n\n    for fold in range(INF_CFG.N_FOLDS):\n        ckpt_path = INF_CFG.CKPT_PATTERN_FOLD_X.format(fold=fold)\n        if not os.path.exists(ckpt_path):\n            ckpt_path = INF_CFG.CKPT_PATTERN_FOLDX.format(fold=fold)\n\n        if not os.path.exists(ckpt_path):\n            print(f\"  ⚠️  Fold {fold} checkpoint 不存在，尝试路径:\")\n            print(f\"    - {INF_CFG.CKPT_PATTERN_FOLD_X.format(fold=fold)}\")\n            print(f\"    - {INF_CFG.CKPT_PATTERN_FOLDX.format(fold=fold)}\")\n            continue\n\n        model = load_model_from_checkpoint(ckpt_path)\n        models.append(model)\n\n        if input_res is None:\n            input_res = getattr(model, \"input_res\", 518)\n            print(f\"  输入分辨率: {input_res}\")\n\n    if len(models) == 0:\n        print(\"\\n❌ 错误：没有找到任何可用的 checkpoint！\")\n        print(f\"   请检查实验目录: {INF_CFG.EXPERIMENT_DIR}\")\n        for fold in range(INF_CFG.N_FOLDS):\n            print(f\"     - {INF_CFG.CKPT_PATTERN_FOLD_X.format(fold=fold)}\")\n            print(f\"     - {INF_CFG.CKPT_PATTERN_FOLDX.format(fold=fold)}\")\n        raise RuntimeError(\"没有找到任何可用的 checkpoint！\")\n\n    print(f\"\\n✓ 成功加载 {len(models)} 个模型\")\n\n    # TTA 推理\n    if INF_CFG.USE_TTA:\n        tta_transforms = get_tta_transforms(input_res)\n        print(f\"\\n使用 TTA: {len(tta_transforms)} 个视角\")\n\n        per_view_preds = []\n\n        for i, transform in enumerate(tta_transforms):\n            view_name = INF_CFG.TTA_TRANSFORMS[i] if i < len(INF_CFG.TTA_TRANSFORMS) else f\"view_{i+1}\"\n            print(f\"\\n--- TTA 视角 {i+1}/{len(tta_transforms)}: {view_name} ---\")\n\n            ds = TestBiomassDataset(test_df, transform, image_dir)\n            dl = DataLoader(\n                ds,\n                batch_size=INF_CFG.BATCH_SIZE,\n                shuffle=False,\n                num_workers=INF_CFG.NUM_WORKERS,\n                pin_memory=True,\n            )\n\n            view_pred = predict_one_view(models, dl)\n            per_view_preds.append(view_pred)\n\n        final_pred = np.mean(per_view_preds, axis=0)\n        print(f\"\\n✓ TTA 完成，最终预测形状: {final_pred.shape}\")\n    else:\n        transform = get_tta_transforms(input_res)[0]\n        ds = TestBiomassDataset(test_df, transform, image_dir)\n        dl = DataLoader(\n            ds,\n            batch_size=INF_CFG.BATCH_SIZE,\n            shuffle=False,\n            num_workers=INF_CFG.NUM_WORKERS,\n            pin_memory=True,\n        )\n        final_pred = predict_one_view(models, dl)\n\n    return final_pred\n\n\n# =============================================================================\n# 生成提交文件\n# =============================================================================\ndef create_submission(final_pred: np.ndarray, test_long: pd.DataFrame, test_unique: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"生成提交文件\"\"\"\n    print(\"\\n\" + \"=\" * 80)\n    print(\"生成提交文件\")\n    print(\"=\" * 80)\n\n    # 提取各目标\n    green = final_pred[:, 0]\n    dead = final_pred[:, 1]\n    clover = final_pred[:, 2]\n    gdm = final_pred[:, 3]\n    total = final_pred[:, 4]\n\n    # 非负 & NaN/Inf 处理\n    def clean(x):\n        x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n        return np.maximum(0, x)\n\n    green, dead, clover, gdm, total = map(clean, [green, dead, clover, gdm, total])\n\n    wide = pd.DataFrame(\n        {\n            \"image_path\": test_unique[\"image_path\"],\n            \"Dry_Green_g\": green,\n            \"Dry_Dead_g\": dead,\n            \"Dry_Clover_g\": clover,\n            \"GDM_g\": gdm,\n            \"Dry_Total_g\": total,\n        }\n    )\n\n    long_preds = wide.melt(\n        id_vars=[\"image_path\"],\n        value_vars=INF_CFG.ALL_TARGET_COLS,\n        var_name=\"target_name\",\n        value_name=\"target\",\n    )\n\n    sub = pd.merge(\n        test_long[[\"sample_id\", \"image_path\", \"target_name\"]],\n        long_preds,\n        on=[\"image_path\", \"target_name\"],\n        how=\"left\",\n    )[[\"sample_id\", \"target\"]]\n\n    sub[\"target\"] = np.nan_to_num(sub[\"target\"], nan=0.0, posinf=0.0, neginf=0.0)\n    sub.to_csv(INF_CFG.SUBMISSION_FILE, index=False)\n\n    print(f\"\\n✓ 提交文件已保存: {INF_CFG.SUBMISSION_FILE}\")\n    print(f\"  样本数: {len(sub)}\")\n    print(f\"  预测统计:\")\n    print(f\"    Dry_Green_g:   mean={green.mean():.2f}, std={green.std():.2f}, min={green.min():.2f}, max={green.max():.2f}\")\n    print(f\"    Dry_Dead_g:    mean={dead.mean():.2f}, std={dead.std():.2f}, min={dead.min():.2f}, max={dead.max():.2f}\")\n    print(f\"    Dry_Clover_g:  mean={clover.mean():.2f}, std={clover.std():.2f}, min={clover.min():.2f}, max={clover.max():.2f}\")\n    print(f\"    GDM_g:         mean={gdm.mean():.2f}, std={gdm.std():.2f}, min={gdm.min():.2f}, max={gdm.max():.2f}\")\n    print(f\"    Dry_Total_g:   mean={total.mean():.2f}, std={total.std():.2f}, min={total.min():.2f}, max={total.max():.2f}\")\n    print(f\"\\n前 10 行预览:\")\n    print(sub.head(10).to_string())\n    return sub\n\n\n# =============================================================================\n# 主函数（使用 parse_known_args 避免 SystemExit:2）\n# =============================================================================\ndef parse_args():\n    parser = argparse.ArgumentParser(description=\"CSIRO v4 CrossPVT T2T Mamba Inference\")\n\n    parser.add_argument(\n        \"--test-csv\",\n        type=str,\n        default=None,\n        help=\"测试集 CSV 路径（默认: INF_CFG.TEST_CSV）\",\n    )\n    parser.add_argument(\n        \"--test-image-dir\",\n        type=str,\n        default=None,\n        help=\"测试图像目录（默认: INF_CFG.TEST_IMAGE_DIR）\",\n    )\n    parser.add_argument(\n        \"--experiment-dir\",\n        type=str,\n        default=None,\n        help=\"实验目录（checkpoint 所在位置，默认: INF_CFG.EXPERIMENT_DIR）\",\n    )\n    parser.add_argument(\n        \"--output\",\n        type=str,\n        default=None,\n        help=\"输出文件路径（默认: INF_CFG.SUBMISSION_FILE）\",\n    )\n    parser.add_argument(\n        \"--batch-size\",\n        type=int,\n        default=None,\n        help=\"批次大小（默认: INF_CFG.BATCH_SIZE）\",\n    )\n    parser.add_argument(\n        \"--no-tta\",\n        action=\"store_true\",\n        help=\"禁用 TTA\",\n    )\n\n    # 关键修改：用 parse_known_args 吞掉 Kaggle/Colab 多余参数（如 -f /tmp/xxx.json）\n    args, _ = parser.parse_known_args()\n    return args\n\n\ndef main():\n    args = parse_args()\n\n    # 更新配置\n    if args.test_csv:\n        INF_CFG.TEST_CSV = args.test_csv\n    if args.test_image_dir:\n        INF_CFG.TEST_IMAGE_DIR = args.test_image_dir\n    if args.experiment_dir:\n        INF_CFG.EXPERIMENT_DIR = args.experiment_dir\n        INF_CFG.CKPT_PATTERN_FOLD_X = os.path.join(\n            INF_CFG.EXPERIMENT_DIR, \"fold_{fold}\", \"checkpoints\", \"best_wr2.pt\"\n        )\n        INF_CFG.CKPT_PATTERN_FOLDX = os.path.join(\n            INF_CFG.EXPERIMENT_DIR, \"fold{fold}\", \"checkpoints\", \"best_wr2.pt\"\n        )\n    if args.output:\n        INF_CFG.SUBMISSION_FILE = args.output\n    if args.batch_size:\n        INF_CFG.BATCH_SIZE = args.batch_size\n    if args.no_tta:\n        INF_CFG.USE_TTA = False\n\n    print(\"=\" * 80)\n    print(\"CSIRO Image2Biomass - v4 CrossPVT T2T Mamba Inference\")\n    print(\"=\" * 80)\n    print(f\"测试 CSV: {INF_CFG.TEST_CSV}\")\n    print(f\"测试图像目录: {INF_CFG.TEST_IMAGE_DIR}\")\n    print(f\"实验目录: {INF_CFG.EXPERIMENT_DIR}\")\n    print(f\"输出文件: {INF_CFG.SUBMISSION_FILE}\")\n    print(f\"批次大小: {INF_CFG.BATCH_SIZE}\")\n    print(f\"使用 TTA: {INF_CFG.USE_TTA}\")\n\n    # 加载测试数据\n    print(\"\\n加载测试数据...\")\n    if not os.path.exists(INF_CFG.TEST_CSV):\n        raise FileNotFoundError(f\"测试 CSV 不存在: {INF_CFG.TEST_CSV}\")\n\n    test_long = pd.read_csv(INF_CFG.TEST_CSV)\n    test_unique = test_long.drop_duplicates(subset=[\"image_path\"]).reset_index(drop=True)\n    print(f\"✓ 找到 {len(test_unique)} 张独立测试图像\")\n    print(f\"  总测试样本数: {len(test_long)}\")\n\n    # 运行推理\n    final_pred = run_inference(test_unique, INF_CFG.TEST_IMAGE_DIR)\n\n    # 生成提交文件\n    submission = create_submission(final_pred, test_long, test_unique)\n\n    print(\"\\n\" + \"=\" * 80)\n    print(\"推理完成！\")\n    print(\"=\" * 80)\n\n    # 清理\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}