{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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":14321240,"sourceType":"datasetVersion","datasetId":9106446,"isSourceIdPinned":true},{"sourceId":14326842,"sourceType":"datasetVersion","datasetId":9146377},{"sourceId":14346204,"sourceType":"datasetVersion","datasetId":9160141},{"sourceId":14354940,"sourceType":"datasetVersion","datasetId":9166225},{"sourceId":14356746,"sourceType":"datasetVersion","datasetId":9167438},{"sourceId":14356747,"sourceType":"datasetVersion","datasetId":9167439},{"sourceId":14356755,"sourceType":"datasetVersion","datasetId":9167446},{"sourceId":14356763,"sourceType":"datasetVersion","datasetId":9167451},{"sourceId":14365012,"sourceType":"datasetVersion","datasetId":9173023},{"sourceId":14365036,"sourceType":"datasetVersion","datasetId":9173043},{"sourceId":14365052,"sourceType":"datasetVersion","datasetId":9173057},{"sourceId":14365056,"sourceType":"datasetVersion","datasetId":9173060},{"sourceId":14365059,"sourceType":"datasetVersion","datasetId":9173062},{"sourceId":14374896,"sourceType":"datasetVersion","datasetId":9180149},{"sourceId":14374898,"sourceType":"datasetVersion","datasetId":9180150},{"sourceId":14374900,"sourceType":"datasetVersion","datasetId":9180152},{"sourceId":14374901,"sourceType":"datasetVersion","datasetId":9180153},{"sourceId":14374932,"sourceType":"datasetVersion","datasetId":9180174},{"sourceId":14379268,"sourceType":"datasetVersion","datasetId":9182940},{"sourceId":14379269,"sourceType":"datasetVersion","datasetId":9182941},{"sourceId":14379271,"sourceType":"datasetVersion","datasetId":9182943},{"sourceId":14383062,"sourceType":"datasetVersion","datasetId":9185575},{"sourceId":14390079,"sourceType":"datasetVersion","datasetId":9190203},{"sourceId":14396581,"sourceType":"datasetVersion","datasetId":9194316},{"sourceId":14429191,"sourceType":"datasetVersion","datasetId":9216228},{"sourceId":14429192,"sourceType":"datasetVersion","datasetId":9216229},{"sourceId":14434099,"sourceType":"datasetVersion","datasetId":9219608},{"sourceId":14434100,"sourceType":"datasetVersion","datasetId":9219609,"isSourceIdPinned":true},{"sourceId":14441153,"sourceType":"datasetVersion","datasetId":9224206},{"sourceId":14441492,"sourceType":"datasetVersion","datasetId":9224484},{"sourceId":14443371,"sourceType":"datasetVersion","datasetId":9225843},{"sourceId":14445476,"sourceType":"datasetVersion","datasetId":9227274},{"sourceId":14445478,"sourceType":"datasetVersion","datasetId":9227276},{"sourceId":14459938,"sourceType":"datasetVersion","datasetId":9235783},{"sourceId":14462976,"sourceType":"datasetVersion","datasetId":9237841},{"sourceId":14470491,"sourceType":"datasetVersion","datasetId":9242419},{"sourceId":14471556,"sourceType":"datasetVersion","datasetId":9243163},{"sourceId":14474552,"sourceType":"datasetVersion","datasetId":9245168},{"sourceId":14474553,"sourceType":"datasetVersion","datasetId":9245169},{"sourceId":14487973,"sourceType":"datasetVersion","datasetId":9253660},{"sourceId":14487975,"sourceType":"datasetVersion","datasetId":9253662},{"sourceId":14487977,"sourceType":"datasetVersion","datasetId":9253664},{"sourceId":14488044,"sourceType":"datasetVersion","datasetId":9253714},{"sourceId":14488047,"sourceType":"datasetVersion","datasetId":9253716},{"sourceId":14509843,"sourceType":"datasetVersion","datasetId":9267495},{"sourceId":14509844,"sourceType":"datasetVersion","datasetId":9267496},{"sourceId":14509845,"sourceType":"datasetVersion","datasetId":9267497},{"sourceId":14516897,"sourceType":"datasetVersion","datasetId":9271746},{"sourceId":14516898,"sourceType":"datasetVersion","datasetId":9271747},{"sourceId":14516899,"sourceType":"datasetVersion","datasetId":9271748},{"sourceId":14529886,"sourceType":"datasetVersion","datasetId":9280225},{"sourceId":14529887,"sourceType":"datasetVersion","datasetId":9280226},{"sourceId":14529888,"sourceType":"datasetVersion","datasetId":9280227},{"sourceId":14529889,"sourceType":"datasetVersion","datasetId":9280228},{"sourceId":14539490,"sourceType":"datasetVersion","datasetId":9286452},{"sourceId":14539492,"sourceType":"datasetVersion","datasetId":9286454},{"sourceId":14539494,"sourceType":"datasetVersion","datasetId":9286456},{"sourceId":14539495,"sourceType":"datasetVersion","datasetId":9286457},{"sourceId":14546041,"sourceType":"datasetVersion","datasetId":9290490},{"sourceId":14547864,"sourceType":"datasetVersion","datasetId":9289607},{"sourceId":14555202,"sourceType":"datasetVersion","datasetId":9296633},{"sourceId":14555203,"sourceType":"datasetVersion","datasetId":9296634},{"sourceId":14555206,"sourceType":"datasetVersion","datasetId":9296637},{"sourceId":14572582,"sourceType":"datasetVersion","datasetId":9308640},{"sourceId":14605549,"sourceType":"datasetVersion","datasetId":9329376},{"sourceId":14605551,"sourceType":"datasetVersion","datasetId":9329377},{"sourceId":14615541,"sourceType":"datasetVersion","datasetId":9335745},{"sourceId":14615548,"sourceType":"datasetVersion","datasetId":9335749},{"sourceId":14615614,"sourceType":"datasetVersion","datasetId":9335790},{"sourceId":14615615,"sourceType":"datasetVersion","datasetId":9335791},{"sourceId":14619915,"sourceType":"datasetVersion","datasetId":9336278,"isSourceIdPinned":true},{"sourceId":14628977,"sourceType":"datasetVersion","datasetId":9344736},{"sourceId":14628978,"sourceType":"datasetVersion","datasetId":9344737},{"sourceId":14629115,"sourceType":"datasetVersion","datasetId":9344833},{"sourceId":14629117,"sourceType":"datasetVersion","datasetId":9344835},{"sourceId":258611181,"sourceType":"kernelVersion"},{"sourceId":258617852,"sourceType":"kernelVersion"},{"sourceId":294159977,"sourceType":"kernelVersion"},{"sourceId":4537,"sourceType":"modelInstanceVersion","modelInstanceId":3329,"modelId":986}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %%\n\"\"\"\ninference-ULTIMATE-2mix-noMeta-Fusion-FDAM-Mamba-DINOv2.py\n==========================================================\n推理代码 - 适配 train-Autodl-ULTIMATE-2mix-noMeta-Fusion-FDAM-Mamba-DINOv2.py 训练的模型\n\n关键特点:\n1. 融合三路特征: 整图 + 左半边 + 右半边\n2. 共享 backbone (DINOv2-base, Patch14)\n3. 使用 FDAMBlock (频域动态注意力调制)\n4. 使用 MobileMambaBlock (SSM 上下文建模)\n5. 使用 MambaFDAMFusion (双流互补融合)\n\"\"\"\n\n# %%\nimport os\nimport gc\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport timm\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nfrom tqdm.auto import tqdm\n\n# 尝试导入 mamba_ssm，如果不行则使用 fallback\ntry:\n    from mamba_ssm import Mamba\n    HAS_MAMBA = True\n    print(\"[Info] mamba_ssm imported successfully.\")\nexcept ImportError:\n    HAS_MAMBA = False\n    print(\"[Warning] mamba_ssm not found. Using Conv1D fallback for MobileMambaBlock.\")\n\n# %%\n# ==========================================\n# 0. 全局配置\n# ==========================================\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Using Device:\", device)\n\n# 必须与训练时保持一致\nIMG_SIZE = 448\nIMG_WIDTH = 448\nIMG_HEIGHT = 896\nTILE_SIZE = 224\nPATCH_SIZE = 14  # DINOv2 patch size\n\nDATA_DIR = \"/kaggle/input/csiro-biomass\"\nMODEL_DIR = \"/kaggle/input/fdam-mamba-dinov2\"\n\n# %%\n# ==========================================\n# 1. 模型定义\n# ==========================================\n\nclass FDAMBlock(nn.Module):\n    \"\"\"\n    Frequency-Dynamic Attention Modulation (FDAM) Block\n    核心功能: 增强高频纹理信息 (Texture/Edge Enhancement)\n    \"\"\"\n    def __init__(self, dim):\n        super().__init__()\n        self.dim = dim\n        self.scale = 0.02\n        \n        # 1. 频域门控 (Frequency Gating/Scaling)\n        self.freq_weight = nn.Sequential(\n            nn.Linear(dim, dim // 4),\n            nn.GELU(),\n            nn.Linear(dim // 4, dim),\n            nn.Sigmoid()\n        )\n        \n        # 2. 空间高通相关性 (Spatial High-Pass)\n        self.high_pass_extractor = nn.Sequential(\n            nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim, bias=False),\n            nn.BatchNorm2d(dim),\n            nn.GELU()\n        )\n        \n        # 3. 特征融合\n        self.fusion = nn.Linear(dim * 2, dim)\n        self.norm = nn.LayerNorm(dim)\n\n    def forward(self, x):\n        \"\"\"\n        x: (B, N, D)  Input tokens\n        \"\"\"\n        B, N, D = x.shape\n        H_feat = int(np.sqrt(N)) # 假设是正方形 (对于 DINOv2 Tiled 这是 true)\n        W_feat = H_feat\n        \n        is_square = (H_feat * W_feat == N)\n        \n        if is_square:\n            # Reshape 2D\n            x_2d = x.transpose(1, 2).view(B, D, H_feat, W_feat)\n            \n            # === Branch 1: Frequency Domain Scaling ===\n            # rfft2 对于实数输入\n            x_fft = torch.fft.rfft2(x_2d, norm='ortho')\n            \n            # 获取幅值\n            x_amp = torch.abs(x_fft)\n            x_pha = torch.angle(x_fft)\n            \n            # 使用通道注意力调整幅值\n            x_amp_pool = torch.mean(x_amp, dim=(2, 3)) # (B, D)\n            freq_gate = self.freq_weight(x_amp_pool).unsqueeze(-1).unsqueeze(-1)\n            \n            # 增强高频\n            x_amp_scaled = x_amp * (1 + freq_gate)\n            \n            # 重建\n            x_fft_scaled = x_amp_scaled * torch.exp(1j * x_pha)\n            x_freq_enhanced = torch.fft.irfft2(x_fft_scaled, s=(H_feat, W_feat), norm='ortho')\n            \n            # === Branch 2: Spatial High-Pass ===\n            x_spatial_high = self.high_pass_extractor(x_2d)\n            \n            # Merge branches\n            x_freq = x_freq_enhanced.flatten(2).transpose(1, 2)\n            x_spat = x_spatial_high.flatten(2).transpose(1, 2)\n            \n            combined = torch.cat([x_freq, x_spat], dim=-1)\n            out = self.fusion(combined)\n            \n            return self.norm(x + out) # Residual\n            \n        else:\n            return x\n\nclass MobileMambaBlock(nn.Module):\n    \"\"\"\n    Mobile Mamba Block\n    核心功能: 利用 SSM 捕捉长距离依赖\n    \"\"\"\n    def __init__(self, dim, d_state=16, d_conv=4, expand=2):\n        super().__init__()\n        self.dim = dim\n        self.fallback = not HAS_MAMBA\n        \n        if not self.fallback:\n            self.mamba = Mamba(\n                d_model=dim,\n                d_state=d_state,\n                d_conv=d_conv,\n                expand=expand\n            )\n        else:\n            # Fallback: Gated Depthwise Conv1D\n            hidden_dim = int(dim * expand)\n            self.in_proj = nn.Linear(dim, hidden_dim * 2)\n            self.conv1d = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=d_conv, padding=d_conv//2, groups=hidden_dim)\n            self.act = nn.SiLU()\n            self.out_proj = nn.Linear(hidden_dim, dim)\n            \n        self.norm = nn.LayerNorm(dim)\n\n    def forward(self, x):\n        \"\"\"\n        x: (B, N, D)\n        \"\"\"\n        res = x\n        x = self.norm(x)\n        \n        if not self.fallback:\n            x = self.mamba(x)\n        else:\n            # Gated Conv Fallback\n            x_proj = self.in_proj(x).transpose(1, 2) # (B, 2*H, N)\n            x1, x2 = x_proj.chunk(2, dim=1)\n            x1 = self.act(self.conv1d(x1)[..., :x.shape[1]]) \n            x = (x1 * F.silu(x2)).transpose(1, 2)\n            x = self.out_proj(x)\n            \n        return res + x\n\nclass MambaFDAMFusion(nn.Module):\n    \"\"\"\n    Dual-Stream Complementary Fusion: Mamba + FDAM\n    \"\"\"\n    def __init__(self, dim, heads=4, dropout=0.2):\n        super().__init__()\n        \n        # Stream 1: Global Context (SSM)\n        self.mamba_stream = MobileMambaBlock(dim)\n        \n        # Stream 2: Local Texture (FDAM)\n        self.fdam_stream = FDAMBlock(dim)\n        \n        # Stream Fusion Gate\n        self.fusion_gate = nn.Sequential(\n            nn.Linear(dim * 2, dim),\n            nn.Sigmoid()\n        )\n        \n        # Cross Attention\n        self.cross_attn = nn.MultiheadAttention(dim, num_heads=heads, dropout=dropout, batch_first=True)\n        self.norm_cross = nn.LayerNorm(dim)\n        self.norm1 = nn.LayerNorm(dim)\n        \n        # FFN\n        self.ffn = nn.Sequential(\n            nn.Linear(dim, dim * 4),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim * 4, dim),\n            nn.Dropout(dropout)\n        )\n        self.norm2 = nn.LayerNorm(dim)\n\n    def forward(self, x_query, x_key):\n        \"\"\"\n        x_query: (B, N_q, D) - fine-grained tokens (small tiles)\n        x_key:   (B, N_k, D) - coarse tokens (big tiles)\n        \"\"\"\n        # 1. 对 Query (Small Tiles) 进行双流增强\n        \n        # Mamba Path (Context)\n        q_mamba = self.mamba_stream(x_query)\n        \n        # FDAM Path (Texture)\n        q_fdam = self.fdam_stream(x_query)\n        \n        # Gated Fusion\n        concat_feat = torch.cat([q_mamba, q_fdam], dim=-1)\n        gate = self.fusion_gate(concat_feat)\n        q_enhanced = gate * q_mamba + (1 - gate) * q_fdam\n        \n        # 2. Cross Attention with Coarse Tokens\n        k_norm = self.norm1(x_key)\n        q_norm = self.norm1(q_enhanced)\n        \n        attn_out, _ = self.cross_attn(q_norm, k_norm, k_norm)\n        x_out = q_enhanced + attn_out\n        \n        # 3. FFN\n        x_out = x_out + self.ffn(self.norm2(x_out))\n        \n        return x_out\n\nclass AttentionPooling(nn.Module):\n    def __init__(self, in_dim):\n        super().__init__()\n        self.attention = nn.Sequential(\n            nn.Linear(in_dim, 128),\n            nn.Tanh(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        w = self.attention(x)\n        w = F.softmax(w, dim=1)\n        return torch.sum(x * w, dim=1)\n\nclass TiledDINOv2MambaFDAM(nn.Module):\n    \"\"\"\n    Tiled DINOv2 + Mamba-FDAM Fusion Main Model\n    \"\"\"\n    def __init__(self, model_name='vit_base_patch14_dinov2.lvd142m'):\n        super().__init__()\n        \n        self.patch_size = PATCH_SIZE\n        self.tile_size = TILE_SIZE\n        \n        # Shared Backbone\n        self.backbone = timm.create_model(\n            model_name, pretrained=False, num_classes=0, img_size=self.tile_size, global_pool=''\n        )\n        self.in_features = self.backbone.num_features\n        self.attn_pool = AttentionPooling(self.in_features)\n        \n        # Grids\n        self.grid_big_split = (2, 2)\n        self.grid_small_split = (4, 4)\n        self.grid_big_whole = (2, 4)\n        self.grid_small_whole = (4, 8)\n        \n        # === Mamba + FDAM Fusion Modules ===\n        self.fusion_split = MambaFDAMFusion(self.in_features)\n        self.fusion_whole = MambaFDAMFusion(self.in_features)\n        \n        # Heads\n        # TiledDINOv2MambaFDAM 聚合策略: \n        # process_stream 返回的是 enhanced small tokens 的 mean pooling\n        # feature dim = D (in_features)\n        # Total streams = 3 (Whole, Left, Right)\n        # Final dim = D * 3\n        self.final_dim = self.in_features * 3\n        \n        self.head_green = self._make_head()\n        self.head_dead = self._make_head()\n        self.head_clover = self._make_head()\n        \n        self.aux_head = nn.Sequential(\n            nn.LayerNorm(self.final_dim),\n            nn.Linear(self.final_dim, 256), nn.GELU(),\n            nn.Linear(256, 5), nn.Softplus()\n        )\n\n    def _make_head(self):\n        return nn.Sequential(\n            nn.LayerNorm(self.final_dim),\n            nn.Linear(self.final_dim, 512),\n            nn.GELU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1),\n            nn.Softplus()\n        )\n\n    def _get_tiles(self, x, grid):\n        B, C, H, W = x.shape\n        rows, cols = grid\n        h_step, w_step = H // rows, W // cols\n        tiles = []\n        for r in range(rows):\n            for c in range(cols):\n                rs, re = r*h_step, (r+1)*h_step\n                cs, ce = c*w_step, (c+1)*w_step\n                tile = x[:, :, rs:re, cs:ce]\n                if tile.shape[-2:] != (self.tile_size, self.tile_size):\n                    tile = F.interpolate(tile, (self.tile_size, self.tile_size), mode='bilinear', align_corners=False)\n                tiles.append(tile)\n        return torch.cat(tiles, dim=0)\n\n    def extract_features(self, x, grid):\n        batch_tiles = self._get_tiles(x, grid)\n        feat = self.backbone.forward_features(batch_tiles)\n        if isinstance(feat, dict): feat = feat['x_norm_patchtokens']\n        \n        # DINOv2: skip CLS token if present (usually DINOv2 doesn't have it in feature map unless explicitly added, but AttnPool handles shape)\n        # However, timm dinov2 forward_features returns (B, N, D).\n        # We need pooling per tile.\n        \n        # Attention Pooling\n        pooled = self.attn_pool(feat) # (N_tiles * B, D)\n        \n        # Reshape to (B, N_tiles, D)\n        B = x.shape[0]\n        N = grid[0] * grid[1]\n        pooled = pooled.view(N, B, -1).permute(1, 0, 2) \n        return pooled\n\n    def process_stream(self, x, grid_small, grid_big, fusion_module):\n        # Extract features\n        feat_s = self.extract_features(x, grid_small) # (B, N_s, D)\n        feat_b = self.extract_features(x, grid_big)   # (B, N_b, D)\n        \n        # Apply Mamba-FDAM Fusion\n        enhanced_s = fusion_module(feat_s, feat_b)\n        \n        # Global Pooling for Final Prediction\n        return enhanced_s.mean(dim=1) # (B, D)\n\n    def forward(self, whole, left, right):\n        # 1. Whole Stream\n        f_whole = self.process_stream(\n            whole, self.grid_small_whole, self.grid_big_whole, self.fusion_whole\n        )\n        \n        # 2. Left Stream\n        f_left = self.process_stream(\n            left, self.grid_small_split, self.grid_big_split, self.fusion_split\n        )\n        \n        # 3. Right Stream\n        f_right = self.process_stream(\n            right, self.grid_small_split, self.grid_big_split, self.fusion_split\n        )\n        \n        # Combine\n        feat = torch.cat([f_whole, f_left, f_right], dim=1) # (B, D*3)\n        \n        # Heads\n        green = self.head_green(feat)\n        dead = self.head_dead(feat)\n        clover = self.head_clover(feat)\n        gdm = green + clover\n        total = gdm + dead\n        \n        main_pred = torch.cat([green, dead, clover, gdm, total], dim=1)\n        aux_pred = self.aux_head(feat)\n        \n        return main_pred, aux_pred\n\n# %%\n# ==========================================\n# 2. 数据准备\n# ==========================================\nprint(\"Preparing Data...\")\n\ntest_df = pd.read_csv(os.path.join(DATA_DIR, \"test.csv\"))\ntest_df[\"image_path\"] = test_df[\"image_path\"].apply(lambda x: os.path.join(DATA_DIR, x))\n\ntarget_name_to_idx = {\n    \"Dry_Green_g\": 0, \"Dry_Dead_g\": 1, \"Dry_Clover_g\": 2, \"GDM_g\": 3, \"Dry_Total_g\": 4\n}\ntest_df[\"target_idx\"] = test_df[\"target_name\"].map(target_name_to_idx)\n\ngroup_df = test_df.groupby(\"image_path\")\nprint(f\"Test Data Prepared: {len(group_df)} unique images.\")\n\n# %%\n# ==========================================\n# 3. Dataset & DataLoader\n# ==========================================\ntransform_split = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\ntransform_whole = transforms.Compose([\n    transforms.Resize((IMG_HEIGHT, IMG_WIDTH)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\nclass TestDataset(Dataset):\n    def __init__(self, grouped_df, transform_split, transform_whole):\n        self.items = list(grouped_df)\n        self.transform_split = transform_split\n        self.transform_whole = transform_whole\n\n    def __len__(self):\n        return len(self.items)\n\n    def __getitem__(self, idx):\n        img_path, df = self.items[idx]\n        try:\n            img = Image.open(img_path).convert(\"RGB\")\n            w, h = img.size\n            mid = w // 2\n            left_img = img.crop((0, 0, mid, h))\n            right_img = img.crop((mid, 0, w, h))\n            whole_img = img\n        except Exception as e:\n            left_img = Image.new('RGB', (IMG_SIZE, IMG_SIZE))\n            right_img = Image.new('RGB', (IMG_SIZE, IMG_SIZE))\n            whole_img = Image.new('RGB', (IMG_WIDTH, IMG_HEIGHT))\n        \n        left_t = self.transform_split(left_img)\n        right_t = self.transform_split(right_img)\n        whole_t = self.transform_whole(whole_img)\n        \n        return whole_t, left_t, right_t, df\n\ndef collate_fn(batch):\n    wholes = [b[0] for b in batch]\n    lefts = [b[1] for b in batch]\n    rights = [b[2] for b in batch]\n    dfs = [b[3] for b in batch]\n    \n    return torch.stack(wholes), torch.stack(lefts), torch.stack(rights), dfs\n\ntest_dataset = TestDataset(group_df, transform_split, transform_whole)\ntest_loader = DataLoader(\n    test_dataset, \n    batch_size=2, \n    shuffle=False,\n    collate_fn=collate_fn,\n    num_workers=2\n)\n\nprint(f\"DataLoader Ready: {len(test_loader)} batches.\")\n\n# %%\n# ==========================================\n# 4. 加载模型\n# ==========================================\nprint(\"Loading Models...\")\nmodels = []\n\nif os.path.exists(MODEL_DIR):\n    models_paths = sorted([\n        os.path.join(MODEL_DIR, p) for p in os.listdir(MODEL_DIR) if p.endswith(\".pth\")\n    ])\n    \n    print(f\"Found {len(models_paths)} models in {MODEL_DIR}\")\n    \n    for path in models_paths:\n        try:\n            model = TiledDINOv2MambaFDAM()\n            state_dict = torch.load(path, map_location=device)\n            model.load_state_dict(state_dict, strict=False)\n            model.to(device)\n            model.eval()\n            models.append(model)\n            print(f\"  Loaded: {os.path.basename(path)}\")\n        except Exception as e:\n            print(f\"  Skipping {path}: {e}\")\nelse:\n    print(f\"Warning: {MODEL_DIR} does not exist.\")\n\nprint(f\"Successfully loaded {len(models)} models.\")\n\n# %%\n# ==========================================\n# 5. 物理约束后处理\n# ==========================================\ndef postprocess_predictions(preds_direct):\n    pred_green = preds_direct[:, 0]\n    pred_gdm = preds_direct[:, 3]\n    pred_total = preds_direct[:, 4]\n    \n    pred_clover = np.maximum(0, pred_gdm - pred_green)\n    pred_dead = np.maximum(0, pred_total - pred_gdm)\n    \n    return np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1)\n\n# %%\n# ==========================================\n# 6. 推理\n# ==========================================\nprint(\"Starting Inference (No TTA)...\")\n\npreds_list = []\n\nwith torch.no_grad():\n    for wholes, lefts, rights, df_batch in tqdm(test_loader, desc=\"Inference\"):\n        wholes = wholes.to(device)\n        lefts = lefts.to(device)\n        rights = rights.to(device)\n        \n        batch_preds_sum = torch.zeros((wholes.size(0), 5), device=device)\n        \n        if len(models) > 0:\n            for model in models:\n                output, _ = model(wholes, lefts, rights)\n                batch_preds_sum += output\n            preds_avg = batch_preds_sum / len(models)\n        else:\n            preds_avg = torch.zeros((wholes.size(0), 5), device=device)\n        \n        preds_avg_np = preds_avg.cpu().numpy()\n        preds_processed = postprocess_predictions(preds_avg_np)\n        \n        for i, df in enumerate(df_batch):\n            batch_pred = preds_processed[i]\n            for _, row in df.iterrows():\n                pred_val = float(batch_pred[row[\"target_idx\"]])\n                if pred_val < 0: pred_val = 0.0\n                preds_list.append({\n                    \"sample_id\": row[\"sample_id\"],               \n                    \"target\": pred_val\n                })\n\n# %%\n# ==========================================\n# 7. 保存结果\n# ==========================================\nprint(\"Saving Submission...\")\n\nsubmission = pd.DataFrame(preds_list)\nsubmission.to_csv(\"submission.csv\", index=False)\n\nprint(f\"Submission saved: {len(submission)} rows\")\nprint(submission.head(10))\n\n# %%\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\n\nprint(\"Done!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T18:18:18.474065Z","iopub.execute_input":"2026-01-26T18:18:18.474428Z","iopub.status.idle":"2026-01-26T18:18:39.131707Z","shell.execute_reply.started":"2026-01-26T18:18:18.474398Z","shell.execute_reply":"2026-01-26T18:18:39.130646Z"}},"outputs":[],"execution_count":null}]}