{"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":"gpu","dataSources":[{"sourceId":126119,"databundleVersionId":14953781,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, re, random\nfrom dataclasses import dataclass\nfrom typing import Dict, Optional, List, Tuple\n\nimport numpy as np\nimport pandas as pd\n\nimport cv2\nimport tifffile as tiff\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport timm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:06.173973Z","iopub.execute_input":"2026-01-02T02:58:06.174197Z","iopub.status.idle":"2026-01-02T02:58:18.610166Z","shell.execute_reply.started":"2026-01-02T02:58:06.174176Z","shell.execute_reply":"2026-01-02T02:58:18.609495Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@dataclass\nclass CFG:\n    ROOT: str = \"/kaggle/input/beyond-visible-spectrum-ai-for-agriculture-2026/Kaggle_Prepared\"  \n    TRAIN_DIR: str = \"train\"\n    VAL_DIR: str = \"val\"\n\n    USE_RGB: bool = True\n    USE_MS: bool  = True\n    USE_HS: bool  = True\n\n    LOAD_SIZE:int = 64\n    IMG_SIZE: int = 64\n\n    BATCH_SIZE: int = 128  \n    EPOCHS: int = 100\n    LR: float = 2e-4\n    WD: float = 0.05\n\n    NUM_WORKERS: int = 4\n    SEED: int = 3557\n\n    RGB_BACKBONE: str = \"convnext_base\"  \n    AMP: bool = True\n\n    HS_DROP_FIRST: int = 10\n    HS_DROP_LAST: int = 14\n\n    OUT_DIR: str = \"/kaggle/working/\"\n    BEST_CKPT: str = \"best.pt\"\n\n\nLABELS = [\"Health\", \"Rust\", \"Other\"]\nLBL2ID = {k: i for i, k in enumerate(LABELS)}\nID2LBL = {i: k for k, i in LBL2ID.items()}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.611359Z","iopub.execute_input":"2026-01-02T02:58:18.611604Z","iopub.status.idle":"2026-01-02T02:58:18.618172Z","shell.execute_reply.started":"2026-01-02T02:58:18.611582Z","shell.execute_reply":"2026-01-02T02:58:18.617502Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\ndef list_files(folder: str, exts: Tuple[str, ...]) -> List[str]:\n    if not os.path.isdir(folder):\n        return []\n    out = []\n    for fn in os.listdir(folder):\n        if fn.lower().endswith(exts):\n            out.append(os.path.join(folder, fn))\n    return sorted(out)\n\ndef base_id(path: str) -> str:\n    return os.path.splitext(os.path.basename(path))[0]\n\ndef parse_label_from_train_name(bid: str) -> Optional[str]:\n    m = re.match(r\"^(Health|Rust|Other)_\", bid)\n    return m.group(1) if m else None\n\ndef build_index(root: str, split: str) -> Dict[str, Dict[str, str]]:\n    split_dir = os.path.join(root, split)\n    rgb_dir = os.path.join(split_dir, \"RGB\")\n    ms_dir  = os.path.join(split_dir, \"MS\")\n    hs_dir  = os.path.join(split_dir, \"HS\")\n\n    rgb_files = list_files(rgb_dir, (\".png\", \".jpg\", \".jpeg\"))\n    ms_files  = list_files(ms_dir, (\".tif\", \".tiff\"))\n    hs_files  = list_files(hs_dir, (\".tif\", \".tiff\"))\n\n    idx: Dict[str, Dict[str, str]] = {}\n    for p in rgb_files:\n        idx.setdefault(base_id(p), {})[\"rgb\"] = p\n    for p in ms_files:\n        idx.setdefault(base_id(p), {})[\"ms\"] = p\n    for p in hs_files:\n        idx.setdefault(base_id(p), {})[\"hs\"] = p\n    return idx\n    \ndef center_crop(x: torch.Tensor, size: int) -> torch.Tensor:\n    _, h, w = x.shape\n    top = (h - size) // 2\n    left = (w - size) // 2\n    return x[:, top:top+size, left:left+size]\n\ndef random_crop(x: torch.Tensor, size: int) -> torch.Tensor:\n    _, h, w = x.shape\n    if h == size and w == size:\n        return x\n    top = random.randint(0, h - size)\n    left = random.randint(0, w - size)\n    return x[:, top:top+size, left:left+size]\n    \ndef make_train_df(train_idx: Dict[str, Dict[str, str]]) -> pd.DataFrame:\n    rows = []\n    for bid, paths in train_idx.items():\n        lab = parse_label_from_train_name(bid)\n        if lab is None:\n            continue\n        rows.append({\n            \"base_id\": bid,\n            \"label\": lab,\n            \"rgb\": paths.get(\"rgb\"),\n            \"ms\":  paths.get(\"ms\"),\n            \"hs\":  paths.get(\"hs\"),\n        })\n    return pd.DataFrame(rows)\n\ndef make_val_df(val_idx: Dict[str, Dict[str, str]]) -> pd.DataFrame:\n    rows = []\n    for bid, paths in val_idx.items():\n        rows.append({\n            \"base_id\": bid,\n            \"rgb\": paths.get(\"rgb\"),\n            \"ms\":  paths.get(\"ms\"),\n            \"hs\":  paths.get(\"hs\"),\n        })\n    return pd.DataFrame(rows)\n\ndef stratified_holdout(df: pd.DataFrame, frac: float = 0.1, seed: int = 42) -> Tuple[pd.DataFrame, pd.DataFrame]:\n    df = df.sample(frac=1.0, random_state=seed).reset_index(drop=True)\n    parts = []\n    for lab, g in df.groupby(\"label\"):\n        n = max(1, int(len(g) * frac))\n        parts.append(g.iloc[:n])\n    df_va = pd.concat(parts).drop_duplicates(\"base_id\")\n    df_tr = df[~df[\"base_id\"].isin(df_va[\"base_id\"])].reset_index(drop=True)\n    df_va = df_va.reset_index(drop=True)\n    return df_tr, df_va","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.619075Z","iopub.execute_input":"2026-01-02T02:58:18.619255Z","iopub.status.idle":"2026-01-02T02:58:18.639693Z","shell.execute_reply.started":"2026-01-02T02:58:18.619238Z","shell.execute_reply":"2026-01-02T02:58:18.639136Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)\nIMAGENET_STD  = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)\n\ndef fix_channels(x: torch.Tensor, target_c: int) -> torch.Tensor:\n    # x: (C,H,W) -> (target_c,H,W) by crop or zero-pad\n    c, h, w = x.shape\n    if c == target_c:\n        return x\n    if c > target_c:\n        return x[:target_c]\n    pad = torch.zeros((target_c - c, h, w), dtype=x.dtype)\n    return torch.cat([x, pad], dim=0)\n    \ndef read_rgb(path: str) -> torch.Tensor:\n    img = cv2.imread(path, cv2.IMREAD_COLOR)\n    if img is None:\n        raise FileNotFoundError(path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n    x = torch.from_numpy(img).permute(2, 0, 1)  # (3,H,W)\n    x = (x - IMAGENET_MEAN) / IMAGENET_STD      \n    return x\n\ndef read_tiff_multiband(path: str) -> np.ndarray:\n    arr = tiff.imread(path)  # (H,W,C) or (C,H,W)\n    if arr.ndim != 3:\n        raise ValueError(f\"Expected 3D TIFF, got {arr.shape} for {path}\")\n    if arr.shape[0] < arr.shape[1] and arr.shape[0] < arr.shape[2]:\n        arr = np.transpose(arr, (1, 2, 0))  # -> (H,W,C)\n    return arr\n\ndef normalize_per_band_minmax(x: np.ndarray, eps: float = 1e-6) -> np.ndarray:\n    x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)\n    H, W, C = x.shape\n    flat = x.reshape(-1, C)\n    mn = flat.min(axis=0)\n    mx = flat.max(axis=0)\n    denom = (mx - mn)\n    denom[denom < eps] = 1.0\n    x = (x - mn.reshape(1, 1, C)) / denom.reshape(1, 1, C)\n    return np.clip(x, 0.0, 1.0)\n\ndef read_ms(path: str) -> torch.Tensor:\n    arr = read_tiff_multiband(path)           # (H,W,5)\n    arr = normalize_per_band_minmax(arr)\n    return torch.from_numpy(arr).permute(2, 0, 1)  # (5,H,W)\n\ndef read_hs(path: str, drop_first: int, drop_last: int) -> torch.Tensor:\n    arr = read_tiff_multiband(path)           # (H,W,B)\n    B = arr.shape[2]\n    if B > (drop_first + drop_last + 1):\n        arr = arr[:, :, drop_first:B - drop_last]\n    arr = normalize_per_band_minmax(arr)\n    return torch.from_numpy(arr).permute(2, 0, 1)  # (B',H,W)\n\ndef resize_tensor(x: torch.Tensor, size: int) -> torch.Tensor:\n    # x: (C,H,W) -> (C,size,size)\n    return F.interpolate(x.unsqueeze(0), size=(size, size), mode=\"bilinear\", align_corners=False).squeeze(0)\n\ndef apply_joint_aug(x_rgb, x_ms, x_hs):\n    k = random.randint(0, 3)\n    do_h = random.random() < 0.5\n    do_v = random.random() < 0.5\n\n    def _tf(x):\n        if x is None:\n            return None\n        if k:\n            x = torch.rot90(x, k, dims=(1, 2))\n        if do_h:\n            x = torch.flip(x, dims=(2,))\n        if do_v:\n            x = torch.flip(x, dims=(1,))\n        return x\n\n    return _tf(x_rgb), _tf(x_ms), _tf(x_hs)\n\ndef infer_hs_in_ch(df_train: pd.DataFrame, df_val: pd.DataFrame, cfg: CFG) -> int:\n    for df in (df_train, df_val):\n        if \"hs\" in df.columns:\n            for p in df[\"hs\"].dropna().tolist():\n                if p and os.path.exists(p):\n                    x = read_hs(p, cfg.HS_DROP_FIRST, cfg.HS_DROP_LAST)\n                    return int(x.shape[0])\n    return 101","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.640629Z","iopub.execute_input":"2026-01-02T02:58:18.641Z","iopub.status.idle":"2026-01-02T02:58:18.664168Z","shell.execute_reply.started":"2026-01-02T02:58:18.640971Z","shell.execute_reply":"2026-01-02T02:58:18.663526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WheatMultiModalDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, cfg: CFG, hs_in_ch: int, train: bool):\n        self.df = df.reset_index(drop=True)\n        self.cfg = cfg\n        self.hs_in_ch = hs_in_ch\n        self.train = train\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, i: int):\n        row = self.df.iloc[i]\n        bid = row[\"base_id\"]\n\n        x_rgb = x_ms = x_hs = None\n        m_rgb = m_ms = m_hs = 0.0\n\n        if self.cfg.USE_RGB and row.get(\"rgb\") is not None:\n            x_rgb = read_rgb(row[\"rgb\"])\n            x_rgb = resize_tensor(x_rgb, self.cfg.LOAD_SIZE)\n            x_rgb = random_crop(x_rgb, self.cfg.IMG_SIZE) if self.train else center_crop(x_rgb, self.cfg.IMG_SIZE)\n            m_rgb = 1.0\n\n        if self.cfg.USE_MS and row.get(\"ms\") is not None:\n            x_ms = read_ms(row[\"ms\"])\n            x_ms = resize_tensor(x_ms, self.cfg.LOAD_SIZE)\n            x_ms = random_crop(x_ms, self.cfg.IMG_SIZE) if self.train else center_crop(x_ms, self.cfg.IMG_SIZE)\n            m_ms = 1.0\n\n\n        if self.cfg.USE_HS and isinstance(row.get(\"hs\"), str) and row[\"hs\"]:\n            x_hs = read_hs(row[\"hs\"], self.cfg.HS_DROP_FIRST, self.cfg.HS_DROP_LAST)\n            x_hs = fix_channels(x_hs, self.hs_in_ch) \n            x_hs = resize_tensor(x_hs, self.cfg.LOAD_SIZE)\n            x_hs = random_crop(x_hs, self.cfg.IMG_SIZE) if self.train else center_crop(x_hs, self.cfg.IMG_SIZE)\n            m_hs = 1.0\n    \n        if self.train:\n            x_rgb, x_ms, x_hs = apply_joint_aug(x_rgb, x_ms, x_hs)\n\n        if self.cfg.USE_RGB and x_rgb is None:\n            x_rgb = torch.zeros(3, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE, dtype=torch.float32)\n        if self.cfg.USE_MS and x_ms is None:\n            x_ms = torch.zeros(5, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE, dtype=torch.float32)\n        if self.cfg.USE_HS and x_hs is None:\n            x_hs = torch.zeros(self.hs_in_ch, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE, dtype=torch.float32)\n\n        mask = torch.tensor([m_rgb, m_ms, m_hs], dtype=torch.float32)\n\n        if \"label\" in row:\n            y = LBL2ID[row[\"label\"]]\n            return {\"id\": bid, \"rgb\": x_rgb, \"ms\": x_ms, \"hs\": x_hs, \"mask\": mask, \"y\": torch.tensor(y, dtype=torch.long)}\n        else:\n            return {\"id\": bid, \"rgb\": x_rgb, \"ms\": x_ms, \"hs\": x_hs, \"mask\": mask}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.665601Z","iopub.execute_input":"2026-01-02T02:58:18.666022Z","iopub.status.idle":"2026-01-02T02:58:18.679086Z","shell.execute_reply.started":"2026-01-02T02:58:18.665999Z","shell.execute_reply":"2026-01-02T02:58:18.678313Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EnhancedSpectralEncoder(nn.Module):\n    \"\"\"改进的光谱编码器 - 引入光谱注意力和残差连接\"\"\"\n    \n    def __init__(self, in_ch: int, emb_dim: int = 256, use_spectral_attention=True):\n        super().__init__()\n        self.use_spectral_attention = use_spectral_attention\n        \n        # 光谱降维层（特别针对高光谱）\n        if in_ch > 32:\n            self.spectral_reduce = nn.Sequential(\n                nn.Conv2d(in_ch, 64, kernel_size=1),\n                nn.BatchNorm2d(64),\n                nn.GELU(),\n                nn.Conv2d(64, 32, kernel_size=1),\n                nn.BatchNorm2d(32),\n                nn.GELU()\n            )\n            in_ch_reduced = 32\n        else:\n            self.spectral_reduce = nn.Identity()\n            in_ch_reduced = in_ch\n        \n        # 光谱注意力机制\n        if use_spectral_attention:\n            self.spectral_attention = SpectralAttentionBlock(in_ch_reduced)\n        \n        # 残差块序列\n        self.residual_blocks = nn.ModuleList([\n            ResidualBlock(in_ch_reduced, 64, stride=1),\n            ResidualBlock(64, 128, stride=2),\n            ResidualBlock(128, 256, stride=2),\n        ])\n        \n        # 全局上下文模块\n        self.global_context = GlobalContextBlock(256)\n        \n        # 输出层\n        self.head = nn.Sequential(\n            nn.AdaptiveAvgPool2d(1),\n            nn.Flatten(),\n            nn.Linear(256, emb_dim),\n            nn.LayerNorm(emb_dim),\n            nn.GELU(),\n            nn.Dropout(0.2)\n        )\n        \n    def forward(self, x):\n        # 光谱降维（针对高光谱）\n        x = self.spectral_reduce(x)\n        \n        # 光谱注意力\n        if self.use_spectral_attention:\n            x = self.spectral_attention(x)\n        \n        # 残差块\n        for block in self.residual_blocks:\n            x = block(x)\n        \n        # 全局上下文\n        x = self.global_context(x)\n        \n        # 输出\n        return self.head(x)\n\n\nclass SpectralAttentionBlock(nn.Module):\n    \"\"\"光谱注意力：学习不同波段的重要性\"\"\"\n    \n    def __init__(self, in_channels):\n        super().__init__()\n        self.gap = nn.AdaptiveAvgPool2d(1)\n        \n        self.attention_net = nn.Sequential(\n            nn.Linear(in_channels, max(4, in_channels // 8)),\n            nn.GELU(),\n            nn.Linear(max(4, in_channels // 8), in_channels),\n            nn.Sigmoid()\n        )\n        \n    def forward(self, x):\n        B, C, H, W = x.shape\n        # 全局平均池化获取每个波段的全局统计\n        pooled = self.gap(x).view(B, C)  # [B, C]\n        # 计算注意力权重\n        weights = self.attention_net(pooled).view(B, C, 1, 1)  # [B, C, 1, 1]\n        # 应用注意力\n        return x * weights\n\n\nclass ResidualBlock(nn.Module):\n    \"\"\"残差块\"\"\"\n    \n    def __init__(self, in_channels, out_channels, stride=1):\n        super().__init__()\n        \n        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, \n                              stride=stride, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(out_channels)\n        \n        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,\n                              padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(out_channels)\n        \n        self.activation = nn.GELU()\n        \n        # 下采样连接\n        self.downsample = None\n        if stride != 1 or in_channels != out_channels:\n            self.downsample = nn.Sequential(\n                nn.Conv2d(in_channels, out_channels, kernel_size=1,\n                         stride=stride, bias=False),\n                nn.BatchNorm2d(out_channels)\n            )\n    \n    def forward(self, x):\n        identity = x\n        \n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.activation(out)\n        \n        out = self.conv2(out)\n        out = self.bn2(out)\n        \n        if self.downsample is not None:\n            identity = self.downsample(x)\n        \n        out += identity\n        out = self.activation(out)\n        \n        return out\n\n\nclass GlobalContextBlock(nn.Module):\n    \"\"\"全局上下文模块\"\"\"\n    \n    def __init__(self, channels, reduction=16):\n        super().__init__()\n        self.channels = channels\n        \n        self.gap = nn.AdaptiveAvgPool2d(1)\n        \n        self.fc = nn.Sequential(\n            nn.Linear(channels, channels // reduction, bias=False),\n            nn.GELU(),\n            nn.Linear(channels // reduction, channels, bias=False),\n            nn.Sigmoid()\n        )\n    \n    def forward(self, x):\n        B, C, H, W = x.shape\n        y = self.gap(x).view(B, C)  # [B, C]\n        y = self.fc(y).view(B, C, 1, 1)  # [B, C, 1, 1]\n        return x * y.expand_as(x)\n\n\nclass EnhancedMultiModalNet(nn.Module):\n    \"\"\"改进的多模态网络 - 引入跨模态注意力和渐进式融合\"\"\"\n    \n    def __init__(self, cfg: CFG, hs_in_ch: int, n_classes: int = 3):\n        super().__init__()\n        self.use_rgb = cfg.USE_RGB\n        self.use_ms  = cfg.USE_MS\n        self.use_hs  = cfg.USE_HS\n        \n        # 1. 各模态编码器（增强版）\n        if self.use_rgb:\n            # 使用更强的骨干网络\n            self.rgb_enc = timm.create_model(\n                cfg.RGB_BACKBONE, \n                pretrained=True, \n                num_classes=0, \n                global_pool=\"avg\"\n            )\n            rgb_dim = self.rgb_enc.num_features  # convnext_base: 1024\n            \n            # RGB特征增强\n            self.rgb_enhancer = nn.Sequential(\n                nn.Linear(rgb_dim, 512),\n                nn.LayerNorm(512),\n                nn.GELU(),\n                nn.Dropout(0.1),\n                nn.Linear(512, 256),\n            )\n            rgb_out_dim = 256\n        else:\n            self.rgb_enc = None\n        \n        if self.use_ms:\n            # 多光谱编码器（增强）\n            self.ms_enc = EnhancedSpectralEncoder(in_ch=5, emb_dim=256, use_spectral_attention=True)\n            ms_out_dim = 256\n        else:\n            self.ms_enc = None\n        \n        if self.use_hs:\n            # 高光谱编码器（增强）\n            self.hs_enc = EnhancedSpectralEncoder(in_ch=hs_in_ch, emb_dim=256, use_spectral_attention=True)\n            hs_out_dim = 256\n        else:\n            self.hs_enc = None\n        \n        # 2. 跨模态注意力融合模块\n        fusion_dims = []\n        if self.use_rgb:\n            fusion_dims.append(rgb_out_dim)\n        if self.use_ms:\n            fusion_dims.append(ms_out_dim)\n        if self.use_hs:\n            fusion_dims.append(hs_out_dim)\n        \n        total_fusion_dim = sum(fusion_dims)\n        \n        # 跨模态注意力\n        self.cross_modal_attention = CrossModalAttention(\n            rgb_dim=256 if self.use_rgb else 0,\n            ms_dim=256 if self.use_ms else 0,\n            hs_dim=256 if self.use_hs else 0\n        )\n        \n        # 3. 渐进式特征融合\n        self.fusion_net = nn.Sequential(\n            nn.Linear(total_fusion_dim, 512),\n            nn.LayerNorm(512),\n            nn.GELU(),\n            nn.Dropout(0.3),\n            \n            nn.Linear(512, 256),\n            nn.LayerNorm(256),\n            nn.GELU(),\n            nn.Dropout(0.2),\n        )\n        \n        # 4. 分类头（带注意力）\n        self.classifier = nn.Sequential(\n            nn.Linear(256, 128),\n            nn.LayerNorm(128),\n            nn.GELU(),\n            nn.Dropout(0.1),\n            nn.Linear(128, n_classes)\n        )\n        \n        # 5. 模态权重学习（自适应权重）\n        self.modal_weight_net = nn.Sequential(\n            nn.Linear(total_fusion_dim, 64),\n            nn.GELU(),\n            nn.Linear(64, 3 if self.use_rgb and self.use_ms and self.use_hs else 2),\n            nn.Softmax(dim=-1)\n        )\n    \n    def forward(self, rgb, ms, hs, mask):\n        features = []\n        modal_features = []\n        \n        # RGB特征提取\n        if self.use_rgb:\n            rgb_feat = self.rgb_enc(rgb)  # [B, 1024]\n            rgb_feat = self.rgb_enhancer(rgb_feat) * mask[:, 0:1]  # [B, 256]\n            features.append(rgb_feat)\n            modal_features.append(rgb_feat.unsqueeze(1))  # [B, 1, 256]\n        \n        # 多光谱特征提取\n        if self.use_ms:\n            ms_feat = self.ms_enc(ms) * mask[:, 1:2]  # [B, 256]\n            features.append(ms_feat)\n            modal_features.append(ms_feat.unsqueeze(1))  # [B, 1, 256]\n        \n        # 高光谱特征提取\n        if self.use_hs:\n            hs_feat = self.hs_enc(hs) * mask[:, 2:3]  # [B, 256]\n            features.append(hs_feat)\n            modal_features.append(hs_feat.unsqueeze(1))  # [B, 1, 256]\n        \n        # 早期融合：拼接特征\n        early_fused = torch.cat(features, dim=1)  # [B, total_dim]\n        \n        # 学习模态权重\n        modal_weights = self.modal_weight_net(early_fused)  # [B, 3]\n        \n        # 加权融合\n        weighted_features = []\n        start_idx = 0\n        for i, feat in enumerate(features):\n            weight = modal_weights[:, i:i+1]  # [B, 1]\n            weighted_features.append(feat * weight)\n        \n        # 中期融合：跨模态注意力（如果多个模态）\n        if len(modal_features) > 1:\n            # 拼接模态特征 [B, num_modalities, 256]\n            stacked_modal = torch.cat(modal_features, dim=1)\n            # 跨模态注意力\n            attended = self.cross_modal_attention(stacked_modal)  # [B, 256]\n        else:\n            attended = features[0]\n        \n        # 晚期融合：与加权特征结合\n        if len(weighted_features) > 0:\n            weighted_sum = torch.stack(weighted_features, dim=1).sum(dim=1)\n            combined = attended + weighted_sum\n        else:\n            combined = attended\n        \n        # 渐进式融合网络\n        fused = self.fusion_net(early_fused)\n        \n        # 最终融合\n        final_features = combined + fused\n        \n        # 分类\n        logits = self.classifier(final_features)\n        \n        return logits\n\n\nclass CrossModalAttention(nn.Module):\n    \"\"\"跨模态注意力模块\"\"\"\n    \n    def __init__(self, rgb_dim=0, ms_dim=0, hs_dim=0):\n        super().__init__()\n        \n        # 确定模态数量和维度\n        self.modalities = []\n        dims = []\n        if rgb_dim > 0:\n            self.modalities.append('rgb')\n            dims.append(rgb_dim)\n        if ms_dim > 0:\n            self.modalities.append('ms')\n            dims.append(ms_dim)\n        if hs_dim > 0:\n            self.modalities.append('hs')\n            dims.append(hs_dim)\n        \n        self.num_modalities = len(self.modalities)\n        \n        if self.num_modalities < 2:\n            # 如果少于2个模态，不需要跨模态注意力\n            self.attention = None\n            return\n        \n        # 统一的特征维度\n        self.unified_dim = 256\n        \n        # 投影层（将不同模态投影到统一维度）\n        self.projection = nn.ModuleDict()\n        for modality, dim in zip(self.modalities, dims):\n            self.projection[modality] = nn.Linear(dim, self.unified_dim)\n        \n        # 多头自注意力\n        self.attention = nn.MultiheadAttention(\n            embed_dim=self.unified_dim,\n            num_heads=4,\n            batch_first=True,\n            dropout=0.1\n        )\n        \n        # 层归一化\n        self.norm = nn.LayerNorm(self.unified_dim)\n        \n        # 前馈网络\n        self.ffn = nn.Sequential(\n            nn.Linear(self.unified_dim, self.unified_dim * 4),\n            nn.GELU(),\n            nn.Dropout(0.1),\n            nn.Linear(self.unified_dim * 4, self.unified_dim),\n            nn.Dropout(0.1)\n        )\n        \n        # 输出投影\n        self.output_proj = nn.Linear(self.unified_dim, self.unified_dim)\n    \n    def forward(self, modal_features):\n        \"\"\"\n        输入: [B, num_modalities, feature_dim]\n        输出: [B, unified_dim]\n        \"\"\"\n        if self.attention is None or self.num_modalities < 2:\n            # 如果是单模态，直接返回平均\n            return modal_features.mean(dim=1)\n        \n        B, M, D = modal_features.shape\n        \n        # 投影到统一维度\n        projected = []\n        for i in range(M):\n            feat = modal_features[:, i, :]\n            modality = self.modalities[i]\n            proj_feat = self.projection[modality](feat)\n            projected.append(proj_feat)\n        \n        # [B, M, unified_dim]\n        projected = torch.stack(projected, dim=1)\n        \n        # 自注意力\n        attended, _ = self.attention(projected, projected, projected)\n        \n        # 残差连接和层归一化\n        x = self.norm(projected + attended)\n        \n        # 前馈网络（应用于每个模态）\n        ffn_out = self.ffn(x)\n        x = self.norm(x + ffn_out)\n        \n        # 跨模态聚合（平均所有模态）\n        aggregated = x.mean(dim=1)  # [B, unified_dim]\n        \n        # 输出投影\n        output = self.output_proj(aggregated)\n        \n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.679875Z","iopub.execute_input":"2026-01-02T02:58:18.680152Z","iopub.status.idle":"2026-01-02T02:58:18.696324Z","shell.execute_reply.started":"2026-01-02T02:58:18.680132Z","shell.execute_reply":"2026-01-02T02:58:18.695666Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@torch.no_grad()\ndef evaluate(model, loader, device):\n    model.eval()\n    total, correct = 0, 0\n    conf = np.zeros((3, 3), dtype=np.int64)\n\n    for batch in loader:\n        rgb  = batch[\"rgb\"].to(device)\n        ms   = batch[\"ms\"].to(device)\n        hs   = batch[\"hs\"].to(device)\n        mask = batch[\"mask\"].to(device)\n        y    = batch[\"y\"].to(device)\n\n        logits = model(rgb, ms, hs, mask)\n        pred = logits.argmax(dim=1)\n\n        total += y.size(0)\n        correct += (pred == y).sum().item()\n\n        yt = y.cpu().numpy()\n        yp = pred.cpu().numpy()\n        for t, p in zip(yt, yp):\n            conf[t, p] += 1\n\n    acc = correct / max(1, total)\n\n    f1s = []\n    for c in range(3):\n        tp = conf[c, c]\n        fp = conf[:, c].sum() - tp\n        fn = conf[c, :].sum() - tp\n        prec = tp / max(1, (tp + fp))\n        rec  = tp / max(1, (tp + fn))\n        f1 = 0.0 if (prec + rec) == 0 else (2 * prec * rec / (prec + rec))\n        f1s.append(f1)\n\n    return {\"acc\": float(acc), \"macro_f1\": float(np.mean(f1s))}\n\ndef train_one_epoch_enhanced(model, loader, optimizer, scaler, device, loss_fn=None):\n    \"\"\"改进的训练循环\"\"\"\n    model.train()\n    total_loss, n = 0.0, 0\n    \n    for batch in loader:\n        rgb = batch[\"rgb\"].to(device)\n        ms = batch[\"ms\"].to(device)\n        hs = batch[\"hs\"].to(device)\n        mask = batch[\"mask\"].to(device)\n        y = batch[\"y\"].to(device)\n        \n        optimizer.zero_grad(set_to_none=True)\n        \n        with torch.cuda.amp.autocast(enabled=(scaler is not None)):\n            # 前向传播\n            if hasattr(model, 'return_features') and model.return_features:\n                logits, features, aux_output = model(rgb, ms, hs, mask, return_features=True)\n                \n                # 获取原始高光谱数据用于重建损失\n                if aux_output is not None:\n                    # 假设hs是原始高光谱数据\n                    original_spectra = hs.mean(dim=[2, 3])  # [B, C_hs]\n                    \n                    # 计算复合损失\n                    if loss_fn is not None:\n                        loss, _ = loss_fn(logits, y, features, aux_output, original_spectra)\n                    else:\n                        loss = F.cross_entropy(logits, y)\n                else:\n                    loss = F.cross_entropy(logits, y)\n            else:\n                logits = model(rgb, ms, hs, mask)\n                loss = F.cross_entropy(logits, y)\n        \n        if scaler is not None:\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            loss.backward()\n            optimizer.step()\n        \n        bs = y.size(0)\n        total_loss += loss.item() * bs\n        n += bs\n    \n    return total_loss / max(1, n)\n\n@torch.no_grad()\ndef predict(model, loader, device):\n    model.eval()\n    preds = []\n    ids = []\n    for batch in loader:\n        rgb  = batch[\"rgb\"].to(device)\n        ms   = batch[\"ms\"].to(device)\n        hs   = batch[\"hs\"].to(device)\n        mask = batch[\"mask\"].to(device)\n\n        logits = model(rgb, ms, hs, mask)\n        p = logits.argmax(dim=1).cpu().numpy().tolist()\n        preds.extend([ID2LBL[x] for x in p])\n        ids.extend(batch[\"id\"])\n    return ids, preds\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.697193Z","iopub.execute_input":"2026-01-02T02:58:18.697452Z","iopub.status.idle":"2026-01-02T02:58:18.714081Z","shell.execute_reply.started":"2026-01-02T02:58:18.697422Z","shell.execute_reply":"2026-01-02T02:58:18.713429Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cfg = CFG()\nseed_everything(cfg.SEED)\nos.makedirs(cfg.OUT_DIR, exist_ok=True)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Device:\", device)\nprint(\"Backbone:\", cfg.RGB_BACKBONE)\n\ntrain_idx = build_index(cfg.ROOT, cfg.TRAIN_DIR)\nval_idx   = build_index(cfg.ROOT, cfg.VAL_DIR)\n\ntrain_df = make_train_df(train_idx)\nval_df   = make_val_df(val_idx)\n\nprint(f\"Indexed train IDs: {len(train_idx)} | usable labeled train rows: {len(train_df)}\")\nprint(f\"Indexed val IDs:   {len(val_idx)}   | val rows: {len(val_df)}\")\n\nif len(train_df) == 0:\n    raise RuntimeError(\"No training samples found. Check ROOT/train and filename label pattern (Health_/Rust_/Other_).\")\n\nhs_in_ch = infer_hs_in_ch(train_df, val_df, cfg)\nprint(\"Inferred HS channels after trimming:\", hs_in_ch)\n\ndf_tr, df_va = stratified_holdout(train_df, frac=0.1, seed=cfg.SEED)\nprint(f\"Train split: {len(df_tr)} | Holdout split: {len(df_va)}\")\n\nds_tr = WheatMultiModalDataset(df_tr, cfg, hs_in_ch=hs_in_ch, train=True)\nds_va = WheatMultiModalDataset(df_va, cfg, hs_in_ch=hs_in_ch, train=False)\nds_te = WheatMultiModalDataset(val_df, cfg, hs_in_ch=hs_in_ch, train=False)\n\ndl_tr = DataLoader(ds_tr, batch_size=cfg.BATCH_SIZE, shuffle=True,\n                   num_workers=4, pin_memory=True, drop_last=True,\n                   persistent_workers=True, prefetch_factor=4)\n\ndl_va = DataLoader(ds_va, batch_size=cfg.BATCH_SIZE, shuffle=False,\n                   num_workers=4, pin_memory=True,\n                   persistent_workers=True, prefetch_factor=4)\n\ndl_te = DataLoader(ds_te, batch_size=cfg.BATCH_SIZE, shuffle=False,\n                   num_workers=cfg.NUM_WORKERS, pin_memory=True)\n\nmodel = EnhancedMultiModalNet(cfg, hs_in_ch=hs_in_ch, n_classes=3).to(device)\n\nbackbone, head = [], []\nfor n, p in model.named_parameters():\n    if not p.requires_grad:\n        continue\n    (backbone if n.startswith(\"rgb_enc\") else head).append(p)\n\noptimizer = torch.optim.AdamW(\n    [\n        {\"params\": backbone, \"lr\": cfg.LR * 0.3}, \n        {\"params\": head,     \"lr\": cfg.LR},\n    ],\n    weight_decay=cfg.WD\n)\n\nscaler = torch.cuda.amp.GradScaler(enabled=(cfg.AMP and device.type == \"cuda\"))\n\nbest_f1 = -1.0\nbest_path = os.path.join(cfg.OUT_DIR, cfg.BEST_CKPT)\n\nfor ep in range(1, cfg.EPOCHS + 1):\n    tr_loss = train_one_epoch_enhanced(model, dl_tr, optimizer, scaler if scaler.is_enabled() else None, device)\n    metrics = evaluate(model, dl_va, device)\n    print(f\"Epoch {ep:02d} | loss={tr_loss:.4f} | val_acc={metrics['acc']:.4f} | val_macroF1={metrics['macro_f1']:.4f}\")\n\n    if metrics[\"macro_f1\"] > best_f1:\n        best_f1 = metrics[\"macro_f1\"]\n        torch.save({\"model\": model.state_dict(), \"hs_in_ch\": hs_in_ch, \"cfg\": cfg.__dict__}, best_path)\n        print(f\"  -> saved best to {best_path}\")\n\nckpt = torch.load(best_path, map_location=device)\nmodel.load_state_dict(ckpt[\"model\"], strict=True)\n\npred_ids, pred_labels = predict(model, dl_te, device)\n\n\nsub_ids = []\nfor _, r in val_df.iterrows():\n    if isinstance(r.get(\"hs\"), str) and r.get(\"hs\"):\n        sub_ids.append(os.path.basename(r[\"hs\"]))\n    elif isinstance(r.get(\"ms\"), str) and r.get(\"ms\"):\n        sub_ids.append(os.path.basename(r[\"ms\"]))\n    else:\n        sub_ids.append(os.path.basename(r[\"rgb\"]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T02:58:18.714932Z","iopub.execute_input":"2026-01-02T02:58:18.715287Z","execution_failed":"2026-01-02T03:06:33.24Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub = pd.DataFrame({\"Id\": sub_ids, \"Category\": pred_labels})\nout_csv = os.path.join(cfg.OUT_DIR, \"submission.csv\")\nsub.to_csv(out_csv, index=False)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-02T03:06:33.24Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub.head(3)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-02T03:06:33.24Z"}},"outputs":[],"execution_count":null}]}