{"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":[{"sourceType":"competition","sourceId":126119,"databundleVersionId":14953781}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, glob, random, torch, timm, cv2\nimport numpy as np\nimport pandas as pd\nimport tifffile as tiff\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import f1_score\nfrom torch.optim.swa_utils import AveragedModel\nfrom tqdm import tqdm\n\n# ====================================================\n# 1. 配置参数 (CFG)\n# ====================================================\nclass CFG:\n    seed = 42\n    n_folds = 5\n    epochs = 40\n    warmup_epochs = 8\n    batch_size = 48 \n    lr = 1e-4 \n    weight_decay = 0.05 \n    mixup_alpha = 0.25\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # 路径确保绝对准确\n    INPUT_DIR = \"/kaggle/input/competitions/beyond-visible-spectrum-ai-for-agriculture-2026/Kaggle_Prepared\"\n    \n    hs_in_channels = 101\n    label_map = {\"Health\": 0, \"Rust\": 1, \"Other\": 2}\n    inv_label_map = {v: k for k, v in label_map.items()}\n\ndef seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything(CFG.seed)\n\n# ====================================================\n# 2. 脏数据检测\n# ====================================================\ndef find_all_bad_ids(root):\n    bad_ids = set()\n    try:\n        train_path = glob.glob(os.path.join(root, \"**/train\"), recursive=True)[0]\n        rgb_files = glob.glob(os.path.join(train_path, \"RGB/*.png\"))\n        for p in rgb_files:\n            bid = os.path.basename(p).split('.')[0]\n            img = cv2.imread(p)\n            if img is None or img.max() < 10 or img.std() < 0.1:\n                bad_ids.add(bid); continue\n            # 检测 TIF\n            for folder in [\"MS\", \"HS\"]:\n                f_path = p.replace(\"/RGB/\", f\"/{folder}/\").replace(\".png\", \".tif\")\n                if os.path.exists(f_path):\n                    tif = tiff.imread(f_path).astype(np.float32)\n                    if tif.max() < 1e-4: bad_ids.add(bid); break\n    except: pass\n    print(f\"Total filtered bad samples: {len(bad_ids)}\")\n    return bad_ids\n\n# ====================================================\n# 3. 精准 7通道统计量 (MS + NDVI + GNDVI)\n# ====================================================\ndef calculate_precise_stats(df):\n    print(\"Calculating Precise Stats (7-ch MS, 101-ch HS)...\")\n    ms_sum, ms_sq_sum = np.zeros(7), np.zeros(7)\n    hs_sum, hs_sq_sum = np.zeros(101), np.zeros(101)\n    p_ms = p_hs = 1e-8\n    for _, row in tqdm(df.iterrows(), total=len(df)):\n        m_raw = tiff.imread(row['ms_path']).astype(np.float32)\n        r, g, nir = m_raw[:,:,2], m_raw[:,:,1], m_raw[:,:,4]\n        m7 = np.concatenate([m_raw, ((nir-r)/(nir+r+1e-8))[...,None], ((nir-g)/(nir+g+1e-8))[...,None]], -1)\n        ms_sum += m7.sum(axis=(0,1)); ms_sq_sum += (m7**2).sum(axis=(0,1)); p_ms += m7.shape[0]*m7.shape[1]\n        h = tiff.imread(row['hs_path']).astype(np.float32)[:,:,10:111]\n        hs_sum += h.sum(axis=(0,1)); hs_sq_sum += (h**2).sum(axis=(0,1)); p_hs += h.shape[0]*h.shape[1]\n    return (torch.tensor(ms_sum/p_ms).float(), torch.tensor(np.sqrt(np.maximum(ms_sq_sum/p_ms-(ms_sum/p_ms)**2, 1e-6))).float(),\n            torch.tensor(hs_sum/p_hs).float(), torch.tensor(np.sqrt(np.maximum(hs_sq_sum/p_hs-(hs_sum/p_hs)**2, 1e-6))).float())\n\n# ====================================================\n# 4. 最终模型：HS主导型架构\n# ====================================================\nclass AgModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # 1. RGB: 降低权重, 保持 64x64 输入\n        self.rgb_net = timm.create_model('efficientnet_b0', pretrained=True, num_classes=0)\n        self.rgb_proj = nn.Sequential(nn.Linear(1280, 128), nn.LayerNorm(128), nn.GELU())\n        \n        # 2. MS: 64 维\n        self.ms_net = nn.Sequential(\n            nn.Conv2d(7, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),\n            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.GELU(),\n            nn.AdaptiveAvgPool2d(1), nn.Flatten()\n        )\n        \n        # 3. HS: 强化深度与宽度 (256维)\n        self.hs_conv = nn.Sequential(\n            nn.Conv2d(101, 128, 1), nn.BatchNorm2d(128), nn.GELU(),\n            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.GELU(),\n            nn.AdaptiveAvgPool2d(1), nn.Flatten()\n        )\n        \n        # 4. 分类器\n        total_dim = 128 + 64 + 256\n        self.classifier = nn.Sequential(\n            nn.Linear(total_dim, 256),\n            nn.LayerNorm(256),\n            nn.GELU(),\n            nn.Dropout(0.5),\n            nn.Linear(256, 3)\n        )\n\n    def forward(self, rgb, ms, hs):\n        f_rgb = self.rgb_proj(self.rgb_net(rgb))\n        f_ms = self.ms_net(ms)\n        f_hs = self.hs_conv(hs)\n        combined = torch.cat([f_rgb, f_ms, f_hs], dim=1)\n        return self.classifier(combined)\n\n# ====================================================\n# 5. Dataset (保持 64x64 + 4x 增强)\n# ====================================================\nclass AgDataset(Dataset):\n    def __init__(self, df, stats, is_train=True):\n        self.df, self.ms_m, self.ms_s, self.hs_m, self.hs_s, self.is_train = df, stats[0], stats[1], stats[2], stats[3], is_train\n    def __len__(self): return len(self.df) * 4 if self.is_train else len(self.df)\n    def __getitem__(self, idx):\n        real_idx, aug_type = (idx // 4, idx % 4) if self.is_train else (idx, 0)\n        row = self.df.iloc[real_idx]\n        \n        # 原生 64x64 处理\n        rgb = cv2.cvtColor(cv2.imread(row['rgb_path']), cv2.COLOR_BGR2RGB).astype(np.float32)/255.0\n        rgb = (rgb - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]\n        \n        m_raw = tiff.imread(row['ms_path']).astype(np.float32)\n        r, g, nir = m_raw[:,:,2], m_raw[:,:,1], m_raw[:,:,4]\n        m7 = np.concatenate([m_raw, ((nir-r)/(nir+r+1e-8))[...,None], ((nir-g)/(nir+g+1e-8))[...,None]], -1)\n        ms = (m7 - self.ms_m.numpy()) / (self.ms_s.numpy() + 1e-8)\n        \n        hs = (tiff.imread(row['hs_path']).astype(np.float32)[:,:,10:111] - self.hs_m.numpy()) / (self.hs_s.numpy()+1e-8)\n        \n        r_t, m_t, h_t = [torch.from_numpy(x).permute(2,0,1).float() for x in [rgb, ms, hs]]\n        if self.is_train:\n            if aug_type == 1: r_t, m_t, h_t = [torch.flip(x, [2]) for x in [r_t, m_t, h_t]]\n            elif aug_type == 2: r_t, m_t, h_t = [torch.flip(x, [1]) for x in [r_t, m_t, h_t]]\n            elif aug_type == 3: r_t, m_t, h_t = [torch.rot90(x, 1, [1, 2]) for x in [r_t, m_t, h_t]]\n            \n        if 'label' in self.df.columns:\n            return r_t, m_t, h_t, torch.tensor(row['label'], dtype=torch.long)\n        return r_t, m_t, h_t\n\n# ====================================================\n# 6. Mixup 函数 (返回混合后的图像和标签对)\n# ====================================================\ndef mixup_data(x, y, alpha=0.2):\n    lam = np.random.beta(alpha, alpha)\n    idx = torch.randperm(x[0].size(0)).to(CFG.device)\n    mixed_x = [lam * xi + (1 - lam) * xi[idx] for xi in x]\n    return mixed_x, y, y[idx], lam\n\n# ====================================================\n# 7. 训练、集成、预测主逻辑\n# ====================================================\ndef run_all():\n    bad_ids = find_all_bad_ids(CFG.INPUT_DIR)\n    train_path = glob.glob(os.path.join(CFG.INPUT_DIR, \"**/train\"), recursive=True)[0]\n    df = pd.DataFrame([{'id': os.path.basename(p).split('.')[0], 'rgb_path': p,\n                        'ms_path': p.replace(\"/RGB/\", \"/MS/\").replace(\".png\", \".tif\"),\n                        'hs_path': p.replace(\"/RGB/\", \"/HS/\").replace(\".png\", \".tif\"),\n                        'label': CFG.label_map[os.path.basename(p).split('_')[0]]} \n                       for p in glob.glob(os.path.join(train_path, \"RGB/*.png\")) if os.path.basename(p).split('.')[0] not in bad_ids])\n    \n    stats = calculate_precise_stats(df)\n    skf = StratifiedKFold(n_splits=CFG.n_folds, shuffle=True, random_state=CFG.seed)\n    \n    for fold, (t_idx, v_idx) in enumerate(skf.split(df, df['label'])):\n        print(f\"\\n### FOLD {fold} ###\")\n        train_loader = DataLoader(AgDataset(df.iloc[t_idx], stats, True), batch_size=CFG.batch_size, shuffle=True, num_workers=4)\n        val_loader = DataLoader(AgDataset(df.iloc[v_idx], stats, False), batch_size=CFG.batch_size, shuffle=False)\n        \n        model = AgModel().to(CFG.device)\n        ema_model = AveragedModel(model) # 必须加入 EMA\n        optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\n        # 改回最稳的 CosineAnnealingLR\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.epochs)\n        criterion = nn.CrossEntropyLoss(label_smoothing=0.1)\n        \n        best_f1 = 0\n        for epoch in range(CFG.epochs):\n            if epoch < CFG.warmup_epochs:\n                for p in model.rgb_net.parameters(): p.requires_grad = False\n            else:\n                for p in model.rgb_net.parameters(): p.requires_grad = True\n                \n            model.train()\n            curr_loss = 0\n            for r, m, h, y in train_loader:\n                r, m, h, y = [x.to(CFG.device) for x in [r, m, h, y]]\n                optimizer.zero_grad()\n                \n                # 原样本训练\n                out_orig = model(r, m, h)\n                loss_orig = criterion(out_orig, y)\n                \n                # 额外加入 Mixup \n                if random.random() < 0.5:\n                    mx, y_a, y_b, lam = mixup_data([r, m, h], y, CFG.mixup_alpha)\n                    out_mix = model(*mx)\n                    loss_mix = lam * criterion(out_mix, y_a) + (1-lam) * criterion(out_mix, y_b)\n                    loss = (loss_orig + loss_mix) / 2\n                else:\n                    loss = loss_orig\n                \n                loss.backward(); optimizer.step(); ema_model.update_parameters(model)\n                curr_loss += loss.item()\n            scheduler.step()\n            \n            # 验证 (使用 EMA 模型)\n            ema_model.eval(); all_p, all_y = [], []\n            with torch.no_grad():\n                for batch in val_loader:\n                    r, m, h, y = [x.to(CFG.device) for x in batch]\n                    out = ema_model(r, m, h)\n                    all_p.append(out.argmax(1).cpu().numpy()); all_y.append(y.cpu().numpy())\n            \n            f1 = f1_score(np.concatenate(all_y), np.concatenate(all_p), average='macro')\n            if f1 > best_f1:\n                best_f1 = f1\n                torch.save(ema_model.module.state_dict(), f\"best_fold{fold}.pt\")\n            print(f\"Epoch {epoch+1:02d} | Loss: {curr_loss/len(train_loader):.4f} | Val F1: {f1:.4f} | Best: {best_f1:.4f}\")\n\n    # 集成推理逻辑\n    print(\"\\n开始终极集成推理...\")\n    val_folder = glob.glob(os.path.join(CFG.INPUT_DIR, \"**/val\"), recursive=True)[0]\n    test_paths = glob.glob(os.path.join(val_folder, \"RGB/*.png\"))\n    test_df = pd.DataFrame([{'rgb_path': p, 'ms_path': p.replace(\"/RGB/\", \"/MS/\").replace(\".png\", \".tif\"),\n                             'hs_path': p.replace(\"/RGB/\", \"/HS/\").replace(\".png\", \".tif\")} for p in test_paths])\n    test_loader = DataLoader(AgDataset(test_df, stats, False), batch_size=CFG.batch_size, shuffle=False)\n    \n    models = []\n    for f in range(CFG.n_folds):\n        m = AgModel().to(CFG.device)\n        m.load_state_dict(torch.load(f\"best_fold{f}.pt\", map_location=CFG.device))\n        m.eval(); models.append(m)\n        \n    final_preds = []\n    with torch.no_grad():\n        for batch in tqdm(test_loader):\n            r, m, h = [x.to(CFG.device) for x in batch]\n            probs = torch.zeros((r.size(0), 3)).to(CFG.device)\n            for model in models:\n                # TTA \n                probs += F.softmax(model(r, m, h), dim=1)\n                probs += F.softmax(model(torch.flip(r,[3]), torch.flip(m,[3]), torch.flip(h,[3])), dim=1)\n                probs += F.softmax(model(torch.flip(r,[2]), torch.flip(m,[2]), torch.flip(h,[2])), dim=1)\n            final_preds.extend(probs.argmax(1).cpu().numpy())\n            \n    test_df['Category'] = [CFG.inv_label_map[p] for p in final_preds]\n    test_df['Id'] = test_df['hs_path'].apply(lambda x: os.path.basename(x))\n    test_df[['Id', 'Category']].to_csv(\"submission.csv\", index=False)\n    print(\"Done!\")\n\nif __name__ == \"__main__\":\n    run_all()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-26T04:53:39.217794Z","iopub.execute_input":"2026-02-26T04:53:39.218389Z","iopub.status.idle":"2026-02-26T05:14:41.083715Z","shell.execute_reply.started":"2026-02-26T04:53:39.218356Z","shell.execute_reply":"2026-02-26T05:14:41.082853Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}