{"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\nimport re\nimport random\nimport time\nfrom dataclasses import dataclass\nfrom typing import Dict, Optional, List, Tuple\n\nimport numpy as np\nimport pandas as pd\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\n# model pretrained for RGB\nimport timm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"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    # MaxViT yêu cầu size chuẩn 224\n    LOAD_SIZE: int = 256\n    IMG_SIZE: int = 224 \n\n    # --- HYPERPARAMETERS ---\n    BATCH_SIZE: int = 16    # Giữ thấp vì MaxViT nặng\n    EPOCHS: int = 50        # Tăng lên để Contrastive Learning có tác dụng\n    \n    # Learning Rate Scheduler Params\n    LR: float = 1e-4        # Max Learning Rate\n    MIN_LR: float = 1e-6    # Min Learning Rate (kết thúc Cosine)\n    WD: float = 0.05        # Weight Decay\n    \n    # Loss Params\n    LABEL_SMOOTHING: float = 0.1\n    SUPCON_WEIGHT: float = 0.5  # Trọng số cho Contrastive Loss\n\n    FUSION_DIM: int = 512\n    NUM_WORKERS: int = 2\n    SEED: int = 42\n\n    RGB_BACKBONE: str = \"maxvit_tiny_tf_224.in1k\" \n    AMP: bool = True\n\n    HS_DROP_FIRST: int = 10\n    HS_DROP_LAST: int = 14\n    OUT_DIR: str = \"./outputs/\"\n    BEST_CKPT: str = \"best_model.pt\"\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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2. UTILS & DATA PREPARATION\n\ndef 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 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, \"label\": lab,\n            \"rgb\": paths.get(\"rgb\"), \"ms\":  paths.get(\"ms\"), \"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\"), \"ms\":  paths.get(\"ms\"), \"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    # Chia train/val nhưng đảm bảo tỉ lệ các class cân bằng\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    return df_tr, df_va.reset_index(drop=True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. IMAGE LOADING & TRANSFORMS\nIMAGENET_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    c, h, w = x.shape\n    if c == target_c: return x\n    if c > target_c: 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: 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)\n    return (x - IMAGENET_MEAN) / IMAGENET_STD\n\ndef read_tiff_multiband(path: str) -> np.ndarray:\n    arr = tiff.imread(path)\n    if arr.ndim != 3: raise ValueError(f\"Expected 3D TIFF, got {arr.shape}\")\n    # Đảm bảo shape là (H, W, C)\n    if arr.shape[0] < arr.shape[1] and arr.shape[0] < arr.shape[2]:\n        arr = np.transpose(arr, (1, 2, 0))\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).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)\n    arr = normalize_per_band_minmax(arr)\n    return torch.from_numpy(arr).permute(2, 0, 1)\n\ndef read_hs(path: str, drop_first: int, drop_last: int) -> torch.Tensor:\n    arr = read_tiff_multiband(path)\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)\n\n# Helper functions for crop/resize\ndef resize_tensor(x: torch.Tensor, size: int) -> torch.Tensor:\n    return F.interpolate(x.unsqueeze(0), size=(size, size), mode=\"bilinear\", align_corners=False).squeeze(0)\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: 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 apply_joint_aug(x_rgb, x_ms, x_hs):\n    # Augmentation đồng bộ cho cả 3 modal\n    k = random.randint(0, 3)\n    do_h = random.random() < 0.5\n    do_v = random.random() < 0.5\n    def _tf(x):\n        if x is None: return None\n        if k: x = torch.rot90(x, k, dims=(1, 2))\n        if do_h: x = torch.flip(x, dims=(2,))\n        if do_v: x = torch.flip(x, dims=(1,))\n        return x\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    # Tự động xác định số kênh HS\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 # Giá trị mặc định nếu không tìm thấy file HS nào","metadata":{"trusted":true},"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        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\"):\n            try:\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            except: pass\n\n        if self.cfg.USE_MS and row.get(\"ms\"):\n            try:\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            except: pass\n\n        if self.cfg.USE_HS and row.get(\"hs\"):\n            try:\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            except: pass\n    \n        if self.train:\n            x_rgb, x_ms, x_hs = apply_joint_aug(x_rgb, x_ms, x_hs)\n\n        if x_rgb is None: x_rgb = torch.zeros(3, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE)\n        if x_ms is None:  x_ms  = torch.zeros(5, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE)\n        if x_hs is None:  x_hs  = torch.zeros(self.hs_in_ch, self.cfg.IMG_SIZE, self.cfg.IMG_SIZE)\n\n        mask = torch.tensor([m_rgb, m_ms, m_hs], dtype=torch.float32)\n        data = {\"id\": bid, \"rgb\": x_rgb, \"ms\": x_ms, \"hs\": x_hs, \"mask\": mask}\n        if \"label\" in row:\n            data[\"y\"] = torch.tensor(LBL2ID[row[\"label\"]], dtype=torch.long)\n        return data","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 5. MODEL ARCHITECTURE (WITH ATTENTION FUSION)\nclass SupervisedContrastiveLoss(nn.Module):\n    def __init__(self, temperature=0.07):\n        super().__init__()\n        self.temperature = temperature\n\n    def forward(self, features, labels):\n        # features: (B, D), labels: (B)\n        device = features.device\n        batch_size = features.shape[0]\n        labels = labels.contiguous().view(-1, 1)\n        mask = torch.eq(labels, labels.T).float().to(device)\n\n        contrast_feature = torch.matmul(features, features.T) \n        logits = contrast_feature / self.temperature\n        \n        logits_max, _ = torch.max(logits, dim=1, keepdim=True)\n        logits = logits - logits_max.detach()\n\n        logits_mask = torch.scatter(torch.ones_like(mask), 1, torch.arange(batch_size).view(-1, 1).to(device), 0)\n        mask = mask * logits_mask\n\n        exp_logits = torch.exp(logits) * logits_mask\n        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)\n\n        mean_log_prob_pos = (mask * log_prob).sum(1) / (mask.sum(1) + 1e-6)\n        loss = - mean_log_prob_pos.mean()\n        return loss\n\nclass SmallSpectralEncoder(nn.Module):\n    \"\"\"Mạng CNN nhẹ cho dữ liệu MS và HS\"\"\"\n    def __init__(self, in_ch: int, emb_dim: int = 256):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Conv2d(in_ch, 32, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(32), nn.ReLU(inplace=True),\n        )\n        self.block = nn.Sequential(\n            nn.Conv2d(32, 64, 3, stride=2, padding=1, bias=False), # Downsample\n            nn.BatchNorm2d(64), nn.ReLU(inplace=True),\n            nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False), # Downsample\n            nn.BatchNorm2d(128), nn.ReLU(inplace=True),\n            nn.Conv2d(128, emb_dim, 3, stride=2, padding=1, bias=False), # Downsample\n            nn.BatchNorm2d(emb_dim), nn.ReLU(inplace=True),\n        )\n        self.pool = nn.AdaptiveAvgPool2d(1)\n\n    def forward(self, x):\n        x = self.stem(x)\n        x = self.block(x)\n        x = self.pool(x)\n        return x.flatten(1)\n\nclass TransformerFusion(nn.Module):\n    def __init__(self, dim: int, num_heads: int = 4, dropout: float = 0.1):\n        super().__init__()\n        # Token đặc biệt [CLS] để tổng hợp thông tin (Learnable)\n        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))\n        \n        # Transformer Encoder Layer đơn giản\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=dim, \n            nhead=num_heads, \n            dim_feedforward=dim * 2, \n            dropout=dropout,\n            batch_first=True,\n            norm_first=True # Quan trọng cho sự ổn định\n        )\n        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)\n\n    def forward(self, x, mask):\n        \"\"\"\n        x: (Batch, 3, Dim) - Đặc trưng của 3 modal\n        mask: (Batch, 3) - 1 nếu có dữ liệu, 0 nếu thiếu\n        \"\"\"\n        b = x.shape[0]\n        \n        # 1. Thêm CLS token vào đầu chuỗi -> (Batch, 4, Dim)\n        cls_tokens = self.cls_token.expand(b, -1, -1)\n        x = torch.cat((cls_tokens, x), dim=1)\n\n        # 2. Tạo Mask cho Transformer\n        # Transformer của PyTorch yêu cầu mask True ở những vị trí CẦN BỎ QUA (Padding)\n        # mask đầu vào: 1=giữ, 0=bỏ.\n        # Ta cần thêm 1 cho vị trí CLS (luôn giữ)\n        mask_cls = torch.ones((b, 1), device=x.device)\n        full_mask = torch.cat((mask_cls, mask), dim=1) # (B, 4)\n        \n        # Đảo ngược: 0 -> True (Ignore), 1 -> False (Keep)\n        src_key_padding_mask = (full_mask == 0)\n\n        # 3. Đi qua Transformer\n        out = self.encoder(x, src_key_padding_mask=src_key_padding_mask)\n\n        # 4. Chỉ lấy vector của CLS token làm đầu ra\n        return out[:, 0]\n\nclass MultiModalNet(nn.Module):\n    def __init__(self, cfg: CFG, hs_in_ch: int, n_classes: int = 3):\n        super().__init__()\n        self.cfg = cfg\n        fusion_dim = cfg.FUSION_DIM\n        \n        if cfg.USE_RGB:\n            self.rgb_enc = timm.create_model(cfg.RGB_BACKBONE, pretrained=True, num_classes=0, global_pool=\"avg\")\n            self.rgb_proj = nn.Sequential(nn.Linear(self.rgb_enc.num_features, fusion_dim), nn.LayerNorm(fusion_dim), nn.GELU())\n        if cfg.USE_MS:\n            self.ms_enc = SmallSpectralEncoder(5, 256)\n            self.ms_proj = nn.Sequential(nn.Linear(256, fusion_dim), nn.LayerNorm(fusion_dim), nn.GELU())\n        if cfg.USE_HS:\n            self.hs_enc = SmallSpectralEncoder(hs_in_ch, 256)\n            self.hs_proj = nn.Sequential(nn.Linear(256, fusion_dim), nn.LayerNorm(fusion_dim), nn.GELU())\n\n        self.fusion = TransformerFusion(dim=fusion_dim, num_heads=8)\n        self.classifier = nn.Sequential(nn.Dropout(0.2), nn.Linear(fusion_dim, n_classes))\n        self._init_weights()\n\n    def _init_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Linear):\n                nn.init.trunc_normal_(m.weight, std=0.02)\n                if m.bias is not None: nn.init.constant_(m.bias, 0)\n\n    def forward(self, rgb, ms, hs, mask):\n        bs = mask.size(0)\n        device = mask.device\n        \n        f_rgb = self.rgb_proj(self.rgb_enc(rgb)) if self.cfg.USE_RGB else torch.zeros(bs, self.cfg.FUSION_DIM, device=device)\n        f_ms = self.ms_proj(self.ms_enc(ms)) if self.cfg.USE_MS else torch.zeros(bs, self.cfg.FUSION_DIM, device=device)\n        f_hs = self.hs_proj(self.hs_enc(hs)) if self.cfg.USE_HS else torch.zeros(bs, self.cfg.FUSION_DIM, device=device)\n\n        stacked = torch.stack([f_rgb, f_ms, f_hs], dim=1)\n        fused = self.fusion(stacked, mask)\n        logits = self.classifier(fused)\n        \n        # Trả về cả Logits và Normalized Features (cho contrastive learning)\n        return logits, F.normalize(fused, dim=1)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 6. TRAINING ENGINE\ndef train_one_epoch(model, loader, optimizer, scaler, device, criterion_ce, criterion_supcon, cfg):\n    model.train()\n    total_loss, n = 0.0, 0\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        with torch.amp.autocast('cuda', enabled=(scaler is not None)):\n            logits, features = model(rgb, ms, hs, mask)\n            loss_ce = criterion_ce(logits, y)\n            loss_sc = criterion_supcon(features, y)\n            loss = loss_ce + (cfg.SUPCON_WEIGHT * loss_sc)\n\n        if scaler:\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            loss.backward()\n            optimizer.step()\n\n        total_loss += loss.item() * y.size(0)\n        n += y.size(0)\n    return total_loss / max(1, n)\n\n@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        output = model(rgb, ms, hs, mask)\n        # Handle tuple return (logits, features)\n        logits = output[0] if isinstance(output, tuple) else output\n\n        pred = logits.argmax(dim=1)\n        total += y.size(0)\n        correct += (pred == y).sum().item()\n        for t, p in zip(y.cpu().numpy(), pred.cpu().numpy()):\n            conf[t, p] += 1\n\n    acc = correct / max(1, total)\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 if (prec + rec) == 0 else (2 * prec * rec / (prec + rec))\n        f1s.append(f1)\n\n    return {\"acc\": acc, \"macro_f1\": np.mean(f1s)}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 7. INFERENCE UTILS\n@torch.no_grad()\ndef predict(model, loader, device):\n    model.eval()\n    preds = []\n    print(\"Running inference...\")\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        output = model(rgb, ms, hs, mask)\n        logits = output[0] if isinstance(output, tuple) else output\n        \n        p = logits.argmax(dim=1).cpu().numpy().tolist()\n        preds.extend([ID2LBL[x] for x in p])\n    return preds","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 8. MAIN EXECUTION\nif __name__ == \"__main__\":\n    cfg = CFG()\n    seed_everything(cfg.SEED)\n    os.makedirs(cfg.OUT_DIR, exist_ok=True)\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Device: {device}\")\n\n    # --- Data ---\n    print(\"Indexing data...\")\n    train_idx = build_index(cfg.ROOT, cfg.TRAIN_DIR)\n    val_idx   = build_index(cfg.ROOT, cfg.VAL_DIR)\n    train_df = make_train_df(train_idx)\n    hs_in_ch = infer_hs_in_ch(train_df, pd.DataFrame(), cfg)\n    print(f\"HS Channels: {hs_in_ch}\")\n\n    df_tr, df_va = stratified_holdout(train_df, frac=0.1, seed=cfg.SEED)\n    ds_tr = WheatMultiModalDataset(df_tr, cfg, hs_in_ch, train=True)\n    ds_va = WheatMultiModalDataset(df_va, cfg, hs_in_ch, train=False)\n\n    dl_tr = DataLoader(ds_tr, batch_size=cfg.BATCH_SIZE, shuffle=True, \n                       num_workers=cfg.NUM_WORKERS, pin_memory=True, drop_last=True)\n    dl_va = DataLoader(ds_va, batch_size=cfg.BATCH_SIZE, shuffle=False, \n                       num_workers=cfg.NUM_WORKERS, pin_memory=True)\n\n    # --- Model & Training Setup ---\n    model = MultiModalNet(cfg, hs_in_ch=hs_in_ch).to(device)\n    \n    # Param Groups for optimizer\n    backbone_params = []\n    head_params = []\n    for n, p in model.named_parameters():\n        if \"rgb_enc\" in n: backbone_params.append(p)\n        else: head_params.append(p)\n\n    optimizer = torch.optim.AdamW([\n        {\"params\": backbone_params, \"lr\": cfg.LR * 0.1},\n        {\"params\": head_params,     \"lr\": cfg.LR}\n    ], weight_decay=cfg.WD)\n\n    # --- IMPROVEMENT: Scheduler ---\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, T_max=cfg.EPOCHS, eta_min=cfg.MIN_LR\n    )\n\n    scaler = torch.amp.GradScaler('cuda', enabled=(cfg.AMP and device.type == \"cuda\"))\n    \n    # --- IMPROVEMENT: Losses ---\n    criterion_ce = nn.CrossEntropyLoss(label_smoothing=cfg.LABEL_SMOOTHING)\n    criterion_supcon = SupervisedContrastiveLoss(temperature=0.07)\n\n    # --- Training Loop ---\n    best_f1 = -1.0\n    best_path = os.path.join(cfg.OUT_DIR, cfg.BEST_CKPT)\n\n    print(f\"Start Training (MaxViT + Transformer Fusion + Contrastive Learning)...\")\n    for ep in range(1, cfg.EPOCHS + 1):\n        t0 = time.time()\n        loss = train_one_epoch(model, dl_tr, optimizer, scaler, device, criterion_ce, criterion_supcon, cfg)\n        \n        # Step Scheduler\n        scheduler.step()\n        curr_lr = optimizer.param_groups[-1]['lr']\n\n        met = evaluate(model, dl_va, device)\n        dt = time.time() - t0\n\n        print(f\"Ep {ep:02d} | LR: {curr_lr:.2e} | Loss: {loss:.4f} | Acc: {met['acc']:.4f} | F1: {met['macro_f1']:.4f}\")\n\n        if met['macro_f1'] > best_f1:\n            best_f1 = met['macro_f1']\n            torch.save(model.state_dict(), best_path)\n            print(f\"  >>> New Best F1: {best_f1:.4f} (Saved)\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 9. SUBMISSION GENERATION\nprint(f\"Loading best model from {best_path}...\")\n    \ntry:\n    # Thử load theo cách thông thường\n    checkpoint = torch.load(best_path, map_location=device)\nexcept:\n    # Nếu lỗi (do version PyTorch mới), thử load với weights_only=True\n    checkpoint = torch.load(best_path, map_location=device, weights_only=True)\n\nmodel.load_state_dict(checkpoint)\nval_df_submission = make_val_df(val_idx)\n\nif len(val_df_submission) > 0:\n    ds_submission = WheatMultiModalDataset(val_df_submission, cfg, hs_in_ch, train=False)\n    dl_submission = DataLoader(ds_submission, batch_size=cfg.BATCH_SIZE, shuffle=False, \n                               num_workers=cfg.NUM_WORKERS, pin_memory=True)\n    pred_labels = predict(model, dl_submission, device)\n    # Reconstruct IDs with extensions\n    real_ids = []\n    for _, row in val_df_submission.iterrows():\n        if row.get(\"ms\") and isinstance(row[\"ms\"], str):\n            filename = os.path.basename(row[\"ms\"])\n        elif row.get(\"hs\") and isinstance(row[\"hs\"], str):\n            filename = os.path.basename(row[\"hs\"])\n        elif row.get(\"rgb\") and isinstance(row[\"rgb\"], str):\n            filename = os.path.basename(row[\"rgb\"])\n        else:\n            filename = row[\"base_id\"] + \".tif\"\n        real_ids.append(filename)\n    sub = pd.DataFrame({\"Id\": real_ids, \"Category\": pred_labels})\n    out_csv = os.path.join(cfg.OUT_DIR, \"submission.csv\")\n    sub.to_csv(out_csv, index=False)\n    print(f\"Submission saved to: {out_csv}\")\n    print(sub.head(3))\nelse:\n    print(\"No validation files found for submission.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-04T13:14:50.735929Z","iopub.execute_input":"2026-01-04T13:14:50.736187Z","iopub.status.idle":"2026-01-04T13:14:50.741309Z","shell.execute_reply.started":"2026-01-04T13:14:50.736155Z","shell.execute_reply":"2026-01-04T13:14:50.740430Z"}},"outputs":[],"execution_count":null}]}