{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":13900620,"sourceType":"datasetVersion","datasetId":8856212},{"sourceId":663314,"sourceType":"modelInstanceVersion","modelInstanceId":471723,"modelId":487624}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\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@dataclass\nclass Config:\n    dropout: float = 0.1\n    hidden_ratio: float = 0.35\n    \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_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    base_path: str = \"/kaggle/input/csiro-biomass\"\n    test_csv: str = \"/kaggle/input/csiro-biomass/test.csv\"\n    test_image_dir: str = \"/kaggle/input/csiro-biomass/test\"\n    \n    experiment_dir_a: str = \"/kaggle/input/csiro/pytorch/default/12\"\n    ckpt_pattern_fold_x_a: str = \"/kaggle/input/csiro/pytorch/default/12/fold_{fold}/checkpoints/best_wr2.pt\"\n    ckpt_pattern_foldx_a: str = \"/kaggle/input/csiro/pytorch/default/12/fold{fold}/checkpoints/best_wr2.pt\"\n    n_folds_a: int = 5\n    \n    model_dir_b: str = \"/kaggle/input/csiro-mvp-models\"\n    model_paths_b: List[str] = field(default_factory=lambda: [\n        f\"/kaggle/input/csiro-mvp-models/model{i}.pth\" for i in range(1, 11)\n    ])\n    weight_a: float = 0.95\n    weight_b: float = 0.075\n    \n    batch_size: int = 1\n    num_workers: int = 0\n    mixed_precision: bool = True\n    use_tta: bool = True\n    \n    submission_file: str = \"submission.csv\"\n    \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    @property\n    def device(self):\n        return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    \n    @property\n    def ckpts_a(self):\n        return self.model_paths_b[:5]\n    \n    @property\n    def ckpts_b(self):\n        return self.model_paths_b[5:]\n\n\nCFG = Config()\n\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    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)\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    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)\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    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\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)\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 = 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        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        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        t3 = self.proj3(t2)\n        pooled = torch.stack([t3.mean(dim=1), t3.max(dim=1).values], dim=1)\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        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                    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\"Cannot create DINO backbone. Last error: {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 forward(self, *inputs, x_left=None, x_right=None, return_features: bool = False):\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 {}\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(\"Cannot infer left/right branches from single tensor.\")\n                x_left, x_right = torch.chunk(x_left, 2, dim=1)\n            else:\n                raise ValueError(\"Missing x_right input.\")\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\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\nclass FiLM(nn.Module):\n    def __init__(self, feat_dim):\n        super().__init__()\n        hidden = max(32, feat_dim // 2)\n        self.mlp = nn.Sequential(\n            nn.Linear(feat_dim, hidden), nn.ReLU(inplace=True), nn.Linear(hidden, feat_dim * 2)\n        )\n\n    def forward(self, context):\n        gamma_beta = self.mlp(context)\n        return torch.chunk(gamma_beta, 2, dim=1)\n\n\nclass BaseDINO(nn.Module):\n    def __init__(self, backbone_name):\n        super().__init__()\n        self.dropout = 0.30\n        self.hidden_ratio = 0.25\n        self.grid = (2, 2)\n        self.backbone_name = backbone_name\n        \n        self.backbone = timm.create_model(backbone_name, pretrained=False, num_classes=0)\n        self.feat_dim = self.backbone.num_features\n        self.input_size = self._get_input_size(self.backbone)\n        self.combined_dim = self.feat_dim * 2\n        hidden_size = max(8, int(self.combined_dim * self.hidden_ratio))\n\n        def make_head():\n            return nn.Sequential(\n                nn.Linear(self.combined_dim, hidden_size),\n                nn.ReLU(inplace=True),\n                nn.Dropout(self.dropout),\n                nn.Linear(hidden_size, 1),\n            )\n\n        self.head_green = make_head()\n        self.head_clover = make_head()\n        self.head_dead = make_head()\n        self.softplus = nn.Softplus(beta=1.0)\n\n    def _get_input_size(self, model):\n        if hasattr(model, \"patch_embed\") and hasattr(model.patch_embed, \"img_size\"):\n            size = model.patch_embed.img_size\n            return int(size if isinstance(size, (int, float)) else size[0])\n        \n        if hasattr(model, \"img_size\"):\n            size = model.img_size\n            return int(size if isinstance(size, (int, float)) else size[0])\n        \n        cfg = getattr(model, \"default_cfg\", {}) or {}\n        input_size = cfg.get(\"input_size\", None)\n        \n        if input_size:\n            if isinstance(input_size, (tuple, list)) and len(input_size) >= 2:\n                return int(input_size[1])\n            return int(input_size if isinstance(input_size, (int, float)) else 224)\n        \n        arch = cfg.get(\"architecture\", \"\") or str(type(model))\n        return 518 if \"dinov2\" in arch.lower() or \"dinov3\" in arch.lower() else 224\n\n    def merge_features(self, left_feat, right_feat):\n        combined = torch.cat([left_feat, right_feat], dim=1)\n        green = self.softplus(self.head_green(combined))\n        clover = self.softplus(self.head_clover(combined))\n        dead = self.softplus(self.head_dead(combined))\n        gdm = green + clover\n        total = gdm + dead\n        return total, gdm, green\n\n\nclass TiledFiLMDINO(BaseDINO):\n    def __init__(self, backbone_name):\n        super().__init__(backbone_name)\n        self.film_left = FiLM(self.feat_dim)\n        self.film_right = FiLM(self.feat_dim)\n\n    def _split_dimension(self, length, parts):\n        step = length // parts\n        segments = []\n        start = 0\n        \n        for _ in range(parts - 1):\n            segments.append((start, start + step))\n            start += step\n        \n        segments.append((start, length))\n        return segments\n\n    def _extract_tile_features(self, x):\n        B, C, H, W = x.shape\n        rows, cols = self.grid\n        row_segments = self._split_dimension(H, rows)\n        col_segments = self._split_dimension(W, cols)\n        features = []\n        \n        for (rs, re) in row_segments:\n            for (cs, ce) in col_segments:\n                tile = x[:, :, rs:re, cs:ce]\n                if tile.shape[-2:] != (self.input_size, self.input_size):\n                    tile = F.interpolate(tile, size=(self.input_size, self.input_size), mode=\"bilinear\")\n                feat = self.backbone(tile)\n                features.append(feat)\n        \n        return torch.stack(features, dim=0).permute(1, 0, 2)\n\n    def _process_stream(self, x, film_layer):\n        tiles = self._extract_tile_features(x)\n        context = tiles.mean(dim=1)\n        gamma, beta = film_layer(context)\n        modulated = tiles * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)\n        return modulated.mean(dim=1)\n\n    def forward(self, left_img, right_img):\n        left_feat = self._process_stream(left_img, self.film_left)\n        right_feat = self._process_stream(right_img, self.film_right)\n        return self.merge_features(left_feat, right_feat)\n\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.img_dir = img_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.img_dir, filename)\n        \n        img = cv2.imread(full_path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        h, w, _ = img.shape\n        mid = w // 2\n        left_img = img[:, :mid]\n        right_img = img[:, mid:]\n\n        left_tensor = self.transform(image=left_img)[\"image\"]\n        right_tensor = self.transform(image=right_img)[\"image\"]\n        \n        return left_tensor, right_tensor\n\n\ndef get_tta_transforms_v4(img_size: int):\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    transforms.append(\n        A.Compose([\n            A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA),\n            *base,\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    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\ndef get_tta_transforms_mvp(img_size):\n    norm = [A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2()]\n    return [\n        A.Compose([A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.HorizontalFlip(p=1.0), A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.VerticalFlip(p=1.0), A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.RandomRotate90(p=1.0), A.Resize(img_size, img_size), *norm]),\n    ]\n\n\ndef strip_module_prefix(state_dict: dict) -> dict:\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 clean_state_dict_mvp(state_dict):\n    if not state_dict:\n        return state_dict\n    \n    cleaned_dict = {}\n    \n    for k, v in state_dict.items():\n        if k.startswith(\"module.\"):\n            k = k[7:]\n        \n        if k.startswith(\"student.\"):\n            k = k[8:]\n        \n        skip_prefixes = (\"txt_enc.\", \"img_proj.\", \"txt_film\", \"teacher.\", \"momentum_teacher.\")\n        if any(k.startswith(prefix) for prefix in skip_prefixes):\n            continue\n            \n        cleaned_dict[k] = v\n    \n    return cleaned_dict\n\n\ndef load_checkpoint_v4(path: str) -> nn.Module:\n    state = torch.load(path, map_location=\"cpu\", weights_only=False)\n    cfg_dict = state.get(\"cfg\", {})\n    dropout = cfg_dict.get(\"dropout\", CFG.dropout)\n    hidden_ratio = cfg_dict.get(\"hidden_ratio\", CFG.hidden_ratio)\n\n    model = CrossPVT_T2T_MambaDINO(dropout=dropout, hidden_ratio=hidden_ratio)\n\n    model_state = state.get(\"model_state\")\n    if model_state is None:\n        model_state = state\n\n    model_state = strip_module_prefix(model_state)\n    model.load_state_dict(model_state, strict=False)\n    model.to(CFG.device)\n    model.eval()\n\n    return model\n\n\ndef load_checkpoint_mvp(checkpoint_path):\n    if not os.path.exists(checkpoint_path):\n        return None\n    \n    try:\n        raw_state = torch.load(checkpoint_path, map_location=CFG.device, weights_only=False)\n    except Exception:\n        return None\n    \n    if isinstance(raw_state, dict):\n        if 'state_dict' in raw_state:\n            state_dict = raw_state['state_dict']\n        elif 'model' in raw_state:\n            state_dict = raw_state['model']\n        else:\n            state_dict = raw_state\n    else:\n        state_dict = raw_state\n    \n    state_dict = clean_state_dict_mvp(state_dict)\n    \n    if not state_dict:\n        return None\n    \n    # Используем те же backbone, что и в исходном рабочем коде\n    backbones = [\n        \"vit_base_patch14_reg4_dinov2\",\n        \"vit_base_patch14_reg4_dinov3\",\n        \"vit_base_patch14_dinov3\",\n    ]\n    \n    for backbone in backbones:\n        try:\n            model = TiledFiLMDINO(backbone)\n            result = model.load_state_dict(state_dict, strict=False)\n            \n            missing = [k for k in result.missing_keys if not k.startswith('backbone.pos_embed')]\n            \n            if len(missing) == 0:\n                model.to(CFG.device)\n                model.eval()\n                return model\n                \n        except Exception:\n            continue\n    \n    return None\n\n\n@torch.no_grad()\ndef predict_one_view_v4(models: List[nn.Module], loader: DataLoader) -> np.ndarray:\n    preds_list = []\n    amp_dtype = \"cuda\" if CFG.device.type == \"cuda\" else \"cpu\"\n    \n    for xl, xr in tqdm(loader, leave=False):\n        xl = xl.to(CFG.device, non_blocking=True)\n        xr = xr.to(CFG.device, non_blocking=True)\n        x_cat = torch.cat([xl, xr], dim=1)\n        \n        per_model_preds = []\n        \n        with torch.amp.autocast(amp_dtype, enabled=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                dead = total - gdm\n                clover = gdm - green\n                five = torch.cat([green, dead, clover, gdm, total], dim=1)\n                \n                per_model_preds.append(five.float().cpu())\n        \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\n@torch.no_grad()\ndef predict_one_view_mvp(models: List[nn.Module], loader: DataLoader) -> np.ndarray:\n    preds = []\n    use_amp = CFG.device.type == \"cuda\"\n    \n    for left_imgs, right_imgs in tqdm(loader, leave=False):\n        left_imgs = left_imgs.to(CFG.device, non_blocking=True)\n        right_imgs = right_imgs.to(CFG.device, non_blocking=True)\n        batch_preds = []\n        \n        with torch.amp.autocast(\"cuda\", enabled=use_amp):\n            for model in models:\n                total, gdm, green = model(left_imgs, right_imgs)\n                dead = torch.clamp(total - gdm, min=0.0)\n                clover = torch.clamp(gdm - green, min=0.0)\n                pred = torch.cat([green, dead, clover, gdm, total], dim=1)\n                batch_preds.append(pred.clamp(0.05, 400.0).cpu())\n        \n        preds.append(torch.stack(batch_preds).mean(dim=0).numpy())\n    \n    return np.concatenate(preds, axis=0)\n\n\ndef run_inference_v4(test_df: pd.DataFrame, image_dir: str) -> np.ndarray:\n    models = []\n    \n    for fold in range(CFG.n_folds_a):\n        ckpt_path = CFG.ckpt_pattern_fold_x_a.format(fold=fold)\n        if not os.path.exists(ckpt_path):\n            ckpt_path = CFG.ckpt_pattern_foldx_a.format(fold=fold)\n        \n        if not os.path.exists(ckpt_path):\n            continue\n        \n        model = load_checkpoint_v4(ckpt_path)\n        models.append(model)\n    \n    if not models:\n        raise RuntimeError(\"No V4 models loaded!\")\n    \n    input_size = getattr(models[0], \"input_res\", 518)\n    \n    if CFG.use_tta:\n        tta_transforms = get_tta_transforms_v4(input_size)\n        per_view_preds = []\n        \n        for transform in tta_transforms:\n            ds = BiomassDataset(test_df, transform, image_dir)\n            dl = DataLoader(\n                ds,\n                batch_size=CFG.batch_size,\n                shuffle=False,\n                num_workers=CFG.num_workers,\n                pin_memory=True,\n            )\n            view_pred = predict_one_view_v4(models, dl)\n            per_view_preds.append(view_pred)\n        \n        final_pred = np.mean(per_view_preds, axis=0)\n    else:\n        transform = get_tta_transforms_v4(input_size)[0]\n        ds = BiomassDataset(test_df, transform, image_dir)\n        dl = DataLoader(\n            ds,\n            batch_size=CFG.batch_size,\n            shuffle=False,\n            num_workers=CFG.num_workers,\n            pin_memory=True,\n        )\n        final_pred = predict_one_view_v4(models, dl)\n    \n    return final_pred\n\n\ndef run_inference_mvp(checkpoint_paths, df, img_dir) -> np.ndarray:\n    models = []\n    \n    for ckpt_path in checkpoint_paths:\n        model = load_checkpoint_mvp(ckpt_path)\n        if model is not None:\n            models.append(model)\n    \n    if not models:\n        raise ValueError(\"No MVP models loaded!\")\n    \n    input_size = models[0].input_size\n    \n    tta_preds = []\n    for transform in get_tta_transforms_mvp(input_size):\n        ds = BiomassDataset(df, transform, img_dir)\n        dl = DataLoader(ds, batch_size=CFG.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n        tta_preds.append(predict_one_view_mvp(models, dl))\n    \n    return np.mean(tta_preds, axis=0)\n\n\ndef create_submission(final_pred: np.ndarray, test_long: pd.DataFrame, test_unique: pd.DataFrame) -> pd.DataFrame:\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    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=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(CFG.submission_file, index=False)\n    \n    return sub\n\n\ntest_df = pd.read_csv(CFG.test_csv)\nunique_df = test_df.drop_duplicates(subset=[\"image_path\"]).reset_index(drop=True)\n\ntry:\n    pred_v4 = run_inference_v4(unique_df, CFG.test_image_dir)\nexcept Exception as e:\n    print(f\"V4 inference failed: {e}\")\n    pred_v4 = None\n\ntry:\n    pred_mvp_a = run_inference_mvp(CFG.ckpts_a, unique_df, CFG.test_image_dir)\n    pred_mvp_b = run_inference_mvp(CFG.ckpts_b, unique_df, CFG.test_image_dir)\n    pred_mvp = CFG.weight_a * pred_mvp_a + CFG.weight_b * pred_mvp_b\nexcept Exception as e:\n    print(f\"MVP inference failed: {e}\")\n    pred_mvp = None\n    \nfinal_pred = (pred_v4 + pred_mvp) / 2.0\n    \nsubmission = create_submission(final_pred, test_df, unique_df)\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:35:00.805447Z","iopub.execute_input":"2025-12-15T14:35:00.806122Z","iopub.status.idle":"2025-12-15T14:35:52.104568Z","shell.execute_reply.started":"2025-12-15T14:35:00.806096Z","shell.execute_reply":"2025-12-15T14:35:52.103858Z"}},"outputs":[],"execution_count":null}]}