{"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":[{"sourceType":"competition","sourceId":126119,"databundleVersionId":14953781},{"sourceType":"datasetVersion","sourceId":14760483,"datasetId":9434488,"databundleVersionId":15611578},{"sourceType":"datasetVersion","sourceId":14760489,"datasetId":9434493,"databundleVersionId":15611583}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\", module=\"pydantic\")  \n\nimport os\nimport random\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport cv2\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import StratifiedKFold\n\ntry:\n    import rasterio\nexcept ModuleNotFoundError:\n    !pip -q install rasterio\n    import rasterio\n\nfrom rasterio.errors import NotGeoreferencedWarning\nwarnings.filterwarnings(\"ignore\", category=NotGeoreferencedWarning)\n\ntry:\n    import timm\nexcept ModuleNotFoundError:\n    !pip -q install timm\n    import timm\n\n# reproducibility package\nSEED = 42\n\ndef set_all_seeds(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Device:\", device)\n\nset_all_seeds(SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:50:29.823907Z","iopub.execute_input":"2026-02-28T15:50:29.824240Z","iopub.status.idle":"2026-02-28T15:50:50.652221Z","shell.execute_reply.started":"2026-02-28T15:50:29.824215Z","shell.execute_reply":"2026-02-28T15:50:50.651542Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**I got the data [here](https://www.kaggle.com/code/yaroslav53/bvs-ai-agri-2026-dataoverview-and-eda)** ","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv('/kaggle/input/d/yaroslav53/bvs-train/train_df (1).csv')\nval_df = pd.read_csv('//kaggle/input/d/yaroslav53/bvs-val/val_df (1).csv')\n\n# list of objects whose pixels are zero\nid_empty_objects_train = ['Other_hyper_26', 'Other_hyper_64', 'Other_hyper_113', 'Other_hyper_163', 'Other_hyper_174',\n                          'Health_hyper_34', 'Health_hyper_167', 'Health_hyper_38', 'Other_hyper_50', 'Health_hyper_67',\n                          'Other_hyper_160', 'Other_hyper_149', 'Other_hyper_31', 'Health_hyper_23', 'Other_hyper_102',\n                          'Health_hyper_153', 'Health_hyper_26', 'Other_hyper_122', 'Other_hyper_155', 'Health_hyper_12',\n                          'Health_hyper_76', 'Other_hyper_121', 'Other_hyper_22']\nid_empty_objects_val =   ['val_5adcdc2e', 'val_6b6509fd', 'val_ceaa5449', 'val_5a361ff8', 'val_6cef0158', 'val_8e35e5d5',\n                          'val_147643a2', 'val_d768878f', 'val_05cee914', 'val_f1a4afff', 'val_da9148e1']\n\n\ntrain_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:50:56.796281Z","iopub.execute_input":"2026-02-28T15:50:56.796818Z","iopub.status.idle":"2026-02-28T15:50:56.865602Z","shell.execute_reply.started":"2026-02-28T15:50:56.796773Z","shell.execute_reply":"2026-02-28T15:50:56.864888Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset/DataLoader","metadata":{}},{"cell_type":"code","source":"CLASSES = [\"Health\", \"Rust\", \"Other\"]\ncat2idx = {c: i for i, c in enumerate(CLASSES)}\nidx2cat = {i: c for c, i in cat2idx.items()}\n\nprint(cat2idx, idx2cat)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:50:59.835493Z","iopub.execute_input":"2026-02-28T15:50:59.836112Z","iopub.status.idle":"2026-02-28T15:50:59.840934Z","shell.execute_reply.started":"2026-02-28T15:50:59.836083Z","shell.execute_reply":"2026-02-28T15:50:59.840046Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# remove constants\ntrain_df = train_df[~train_df[\"Id\"].isin(id_empty_objects_train)].reset_index(drop=True)\ntrain_df.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:22:02.717043Z","iopub.execute_input":"2026-02-28T16:22:02.717735Z","iopub.status.idle":"2026-02-28T16:22:02.725432Z","shell.execute_reply.started":"2026-02-28T16:22:02.717704Z","shell.execute_reply":"2026-02-28T16:22:02.724710Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"*Multispectral and hyperspectral images store raw sensor measurements with higher bit depth rather than visual color intensities. Therefore, we apply per-sample min–max normalization to make the input scale-invariant and robust to illumination differences.*\n\n- RGB global max: 255.0\n- MS global max: 11406.0\n- HS global max: 65535.0","metadata":{}},{"cell_type":"markdown","source":"# MS","metadata":{}},{"cell_type":"markdown","source":"**Core baseline**\n\nA solid baseline was implemented and locked-in using ConvNeXt-Tiny (timm) with a simple MS stem (mapping 5 → 3 channels) and an MLP head:\n\n- MS stem: Conv(5→W, 1x1) → BN → GELU → Conv(W→3, 1x1) → BN\n-Backbone: convnext_tiny pretrained, num_classes=0\n- Head: LayerNorm → Dropout → Linear → GELU → Dropout → Linear\n-Additionally: a set of augmentations was selected via multiple runs\n- MS normalization: baseline minmax_norm\n\nOutcome: this baseline remained the best local performer across most tested variants and delivered the best public leaderboard score among all explored configurations.\n\n**What was tested (and did not help)**\n\nDozens of unique configurations were evaluated, including:\n\n- Other model families: EfficientNetV2, DINO ViT (timm)\n- Multiple sizes/variants within each family (e.g., EfficientNetV2 M/L)\n- Alternative MS input strategies:\n    - “stronger” stem blocks\n    - stems with spatial mixing (including depthwise 3×3)\n- Different MS normalizations:\n    - minmax_norm\n    - make_per_channel_minmax_tf\n    - div_msmax\n\nIn most cases, these alternatives were worse than the baseline locally, so pushing them to public evaluation was not justified.\n\n**The only local win — and why it was rejected**\n\nOne setup improved local CV:\n\n- backbone: vit_small_patch8_224.dino\n- img_size=64 (to resample pos_embed when loading pretrained weights)\n- normalization: make_per_channel_minmax_tf\n\nHowever, it performed significantly worse on the public leaderboard (roughly 0.66 vs 0.72 for the baseline), so it was not considered for ensembling.","metadata":{}},{"cell_type":"code","source":"ms_train_tf = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n\n    A.MultiplicativeNoise(\n        multiplier=(0.97, 1.03),\n        per_channel=False,\n        p=0.15\n    ),\n], p=1.0)\n\nms_val_tf = A.Compose([], p=1.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:22:04.651469Z","iopub.execute_input":"2026-02-28T16:22:04.652126Z","iopub.status.idle":"2026-02-28T16:22:04.659132Z","shell.execute_reply.started":"2026-02-28T16:22:04.652094Z","shell.execute_reply":"2026-02-28T16:22:04.658294Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def minmax_norm(x: np.ndarray, eps: float = 1e-6) -> np.ndarray:\n    \"\"\"per-sample normalization\"\"\"\n    # x: (C,H,W) float32\n    x_min = x.min()\n    x_max = x.max()\n    if (x_max - x_min) < eps:\n        return np.zeros_like(x, dtype=np.float32)\n    return (x - x_min) / (x_max - x_min)\n\npixel_minmax_tf = minmax_norm","metadata":{"execution":{"iopub.status.busy":"2026-02-28T16:22:06.425388Z","iopub.execute_input":"2026-02-28T16:22:06.426018Z","iopub.status.idle":"2026-02-28T16:22:06.430781Z","shell.execute_reply.started":"2026-02-28T16:22:06.425990Z","shell.execute_reply":"2026-02-28T16:22:06.429872Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WheatMSDataset(Dataset):\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        is_train: bool = True,\n        pixel_tf=None,        \n        aug_tf=None,            \n        check_shape: bool = False\n    ):\n        df = df.reset_index(drop=True)\n        self.is_train = is_train\n        self.pixel_tf = pixel_tf\n        self.aug_tf = aug_tf\n        self.check_shape = check_shape\n\n        self.ms_paths = df[\"ms_path\"].astype(str).tolist()\n        self.ids = df[\"Id\"].astype(str).tolist()\n\n        if is_train:\n            self.y = df[\"Category\"].map(cat2idx).to_numpy(np.int64)\n\n    def __len__(self):\n        return len(self.ms_paths)\n\n    def _read_ms(self, path) -> np.ndarray:\n        with rasterio.open(path) as src:\n            x = src.read().astype(np.float32)  # (5,64,64) CHW\n        return x\n\n    def __getitem__(self, idx):\n        ms = self._read_ms(self.ms_paths[idx])  # (C,H,W)\n\n        if self.check_shape and ms.shape != (5, 64, 64):\n            raise ValueError(f\"Unexpected MS shape {ms.shape} for {self.ms_paths[idx]}\")\n\n        # Albumentations is awaiting HWC\n        if self.aug_tf is not None:\n            ms_hwc = np.transpose(ms, (1, 2, 0))                 # (H,W,C)\n            ms_hwc = self.aug_tf(image=ms_hwc)[\"image\"]          # (H,W,C)\n            ms = np.transpose(ms_hwc, (2, 0, 1)).astype(np.float32)  # (C,H,W)\n\n        # normalization of pixels \n        if self.pixel_tf is not None:\n            ms = self.pixel_tf(ms)\n\n        x = torch.from_numpy(ms)  # float32 CHW\n\n        if self.is_train:\n            y = torch.tensor(self.y[idx], dtype=torch.long)\n            return x, y, self.ids[idx]\n        else:\n            return x, self.ids[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:22:08.300163Z","iopub.execute_input":"2026-02-28T16:22:08.300804Z","iopub.status.idle":"2026-02-28T16:22:08.310097Z","shell.execute_reply.started":"2026-02-28T16:22:08.300773Z","shell.execute_reply":"2026-02-28T16:22:08.309417Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MSStemConvNeXt(nn.Module):\n    def __init__(self, backbone_name=\"convnext_tiny\", num_classes=3, drop=0.1, stem_width=16):\n        super().__init__()\n\n        # Slightly stronger spectral stem\n        self.stem = nn.Sequential(\n            nn.Conv2d(5, stem_width, kernel_size=1, bias=False),\n            nn.BatchNorm2d(stem_width),\n            nn.GELU(),\n            nn.Conv2d(stem_width, 3, kernel_size=1, bias=False),\n            nn.BatchNorm2d(3),\n        )\n\n        self.backbone = timm.create_model(backbone_name, pretrained=True, num_classes=0)\n        feat_dim = self.backbone.num_features\n\n        hidden = feat_dim // 2\n        self.head = nn.Sequential(\n            nn.LayerNorm(feat_dim),\n            nn.Dropout(drop),\n            nn.Linear(feat_dim, hidden),\n            nn.GELU(),\n            nn.Dropout(drop),\n            nn.Linear(hidden, num_classes),\n        )\n\n    def forward(self, x):\n        x = self.stem(x)\n        feats = self.backbone(x)\n        return self.head(feats)","metadata":{"execution":{"iopub.status.busy":"2026-02-28T16:22:14.176988Z","iopub.execute_input":"2026-02-28T16:22:14.177633Z","iopub.status.idle":"2026-02-28T16:22:14.185170Z","shell.execute_reply.started":"2026-02-28T16:22:14.177599Z","shell.execute_reply":"2026-02-28T16:22:14.184197Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# configs\nBATCH_SIZE = 64\nNUM_CLASSES = len(CLASSES)\nN_FOLDS = 5\nEPOCHS = 20\nlr_backbone = 1e-4\nlr_head = 3e-4\nWEIGHT_DECAY = 1e-2\nBACKBONE = \"convnext_tiny\"\n\nuse_amp = False    #torch.cuda.is_available()\npin = torch.cuda.is_available()\n\ntrain_df_ms = train_df.copy().reset_index(drop=True)\n\n# OOF containers\noof_logits_ms = np.zeros((len(train_df_ms), NUM_CLASSES), dtype=np.float32)\noof_true   = train_df_ms[\"Category\"].map(cat2idx).values.astype(np.int64)\n\n# Store model weights per fold\nfold_states_ms = []\nfold_best_acc = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:27:53.973946Z","iopub.execute_input":"2026-02-28T16:27:53.974760Z","iopub.status.idle":"2026-02-28T16:27:53.989123Z","shell.execute_reply.started":"2026-02-28T16:27:53.974722Z","shell.execute_reply":"2026-02-28T16:27:53.988327Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_worker(worker_id):\n    worker_seed = (SEED + worker_id) % 2**32\n    np.random.seed(worker_seed)\n    random.seed(worker_seed)\n    torch.manual_seed(worker_seed)\n\ndef make_loaders(\n    df_tr, df_va,\n    batch_size=BATCH_SIZE,\n    pixel_tf=None,\n    aug_tf_train=None,\n    aug_tf_val=None,\n    check_shape=False\n):\n    tr_ds = WheatMSDataset(df_tr, is_train=True, pixel_tf=pixel_tf, aug_tf=aug_tf_train, check_shape=check_shape)\n    va_ds = WheatMSDataset(df_va, is_train=True, pixel_tf=pixel_tf, aug_tf=aug_tf_val,   check_shape=check_shape)\n\n    g = torch.Generator()\n    g.manual_seed(SEED)\n\n    tr_loader = DataLoader(\n        tr_ds,\n        batch_size=batch_size,\n        shuffle=True,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        generator=g,\n        worker_init_fn=seed_worker\n    )\n    va_loader = DataLoader(\n        va_ds,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        worker_init_fn=seed_worker,\n        generator=g\n    )\n    return tr_loader, va_loader\n\n@torch.no_grad()\ndef accuracy_from_logits(logits: torch.Tensor, y: torch.Tensor) -> float:\n    preds = torch.argmax(logits, dim=1)\n    return (preds == y).float().mean().item()\n\ndef train_one_epoch(model, loader, optimizer, criterion, scaler):\n    model.train()\n    total_loss = 0.0\n    total_correct = 0\n    total = 0\n\n    for xb, yb, _ in loader:\n        xb = xb.to(device, non_blocking=True)\n        yb = yb.to(device, non_blocking=True)\n\n        optimizer.zero_grad(set_to_none=True)\n\n        with torch.amp.autocast('cuda', enabled=use_amp):\n            logits = model(xb)\n            loss = criterion(logits, yb)\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        bs = xb.size(0)\n        total_loss += loss.item() * bs\n        total_correct += (logits.argmax(1) == yb).sum().item()\n        total += bs\n\n    return total_loss / total, total_correct / total\n\n@torch.no_grad()\ndef validate_one_epoch(model, loader, criterion):\n    model.eval()\n    total_loss = 0.0\n    total_correct = 0\n    total = 0\n\n    for xb, yb, _ in loader:\n        xb = xb.to(device, non_blocking=True)\n        yb = yb.to(device, non_blocking=True)\n\n        logits = model(xb)\n        loss = criterion(logits, yb)\n\n        bs = xb.size(0)\n        total_loss += loss.item() * bs\n        total_correct += (logits.argmax(1) == yb).sum().item()\n        total += bs\n\n    return total_loss / total, total_correct / total\n\n@torch.no_grad()\ndef predict_logits(model, loader):\n    model.eval()\n    all_logits = []\n\n    for xb, yb, _ in loader:\n        xb = xb.to(device, non_blocking=True)\n        logits = model(xb)\n        all_logits.append(logits.detach().cpu().numpy())\n\n    return np.concatenate(all_logits, axis=0)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:22:30.401245Z","iopub.execute_input":"2026-02-28T16:22:30.402089Z","iopub.status.idle":"2026-02-28T16:22:30.416947Z","shell.execute_reply.started":"2026-02-28T16:22:30.402056Z","shell.execute_reply":"2026-02-28T16:22:30.415914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\nFREEZE_EPOCHS = 2\nPATIENCE = 7\n\nskf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\n\nfor fold, (tr_idx, va_idx) in enumerate(skf.split(train_df_ms, train_df_ms[\"Category\"])):\n    df_tr = train_df_ms.iloc[tr_idx].reset_index(drop=True)\n    df_va = train_df_ms.iloc[va_idx].reset_index(drop=True)\n\n    tr_loader, va_loader = make_loaders(\n        df_tr, df_va,\n        batch_size=BATCH_SIZE,\n        pixel_tf=pixel_minmax_tf,\n        aug_tf_train=ms_train_tf,\n        aug_tf_val=None,         \n        check_shape=True\n    )\n\n    set_all_seeds(SEED + fold)    # To ensure determinism\n    model = MSStemConvNeXt(backbone_name=BACKBONE, num_classes=NUM_CLASSES, drop=0.1).to(device)\n    criterion = nn.CrossEntropyLoss()\n    scaler = torch.amp.GradScaler('cuda', enabled=use_amp)\n\n    # freeze backbone initially\n    for p in model.backbone.parameters():\n        p.requires_grad = False\n\n    # optimizer only for trainable params initially (stem+head)\n    optimizer = torch.optim.AdamW(\n        [\n            {\"params\": model.stem.parameters(), \"lr\": lr_head},\n            {\"params\": model.head.parameters(), \"lr\": lr_head},\n        ],\n        weight_decay=WEIGHT_DECAY\n    )\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)\n\n    best_acc = -1.0\n    best_state = None\n    bad_epochs = 0\n    backbone_unfrozen = False\n\n    print(f\"\\n===== FOLD {fold + 1} / {N_FOLDS} | train={len(df_tr)} valid={len(df_va)} =====\")\n\n    val_acc_hist = []\n    val_loss_hist = []\n\n    for epoch in range(1, EPOCHS + 1):\n\n        # unfreeze backbone after FREEZE_EPOCHS\n        if (not backbone_unfrozen) and (epoch == FREEZE_EPOCHS + 1):\n            for p in model.backbone.parameters():\n                p.requires_grad = True\n            backbone_unfrozen = True\n\n            # Rebuild optimizer to include backbone params\n            optimizer = torch.optim.AdamW(\n                [\n                    {\"params\": model.backbone.parameters(), \"lr\": lr_backbone},\n                    {\"params\": model.stem.parameters(), \"lr\": lr_head},\n                    {\"params\": model.head.parameters(), \"lr\": lr_head},\n                ],\n                weight_decay=WEIGHT_DECAY\n            )\n            # Rebuild scheduler for remaining epochs\n            remaining = EPOCHS - epoch + 1\n            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1, remaining))\n\n            print(f\"Unfroze backbone at epoch {epoch}. Rebuilt optimizer/scheduler (remaining={remaining}).\")\n\n        tr_loss, tr_acc = train_one_epoch(model, tr_loader, optimizer, criterion, scaler)\n        va_loss, va_acc = validate_one_epoch(model, va_loader, criterion)\n        val_loss_hist.append(float(va_loss))\n        val_acc_hist.append(float(va_acc))\n\n        scheduler.step()\n\n        improved = va_acc > best_acc\n        if improved:\n            best_acc = va_acc\n            best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}\n            bad_epochs = 0\n        else:\n            bad_epochs += 1\n\n        current_lrs = [pg[\"lr\"] for pg in optimizer.param_groups]\n        lrs_str = \"[\" + \", \".join([f\"{lr:.2e}\" for lr in current_lrs]) + \"]\"\n        \n        if epoch in (1, 2, 3, 5, 10, 15, 20) or epoch % 5 == 0:\n            print(f\"Epoch {epoch:02d}/{EPOCHS} | \"\n                  f\"train loss {tr_loss:.4f} acc {tr_acc:.4f} | \"\n                  f\"valid loss {va_loss:.4f} acc {va_acc:.4f} | \"\n                  f\"best accuracy {best_acc:.4f} | \"\n                  f\"patience {bad_epochs}/{PATIENCE}\"\"\\n\"\n                  f\"lrs {lrs_str}\")\n\n        # early stopping\n        if bad_epochs >= PATIENCE:\n            print(f\"Early stopping on epoch {epoch} (no val_acc improvement for {PATIENCE} epochs).\")\n            break\n\n    # plots per fold (val metrics vs epoch index)\n    epochs_ran = np.arange(1, len(val_acc_hist) + 1)\n    \n    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n    \n    # Accuracy\n    axes[0].plot(epochs_ran, val_acc_hist)\n    axes[0].set_title(f\"Fold {fold+1}: Val Accuracy\")\n    axes[0].set_xlabel(\"epoch\")\n    axes[0].set_ylabel(\"val_acc\")\n    \n    # Loss\n    axes[1].plot(epochs_ran, val_loss_hist)\n    axes[1].set_title(f\"Fold {fold+1}: Val Loss\")\n    axes[1].set_xlabel(\"epoch\")\n    axes[1].set_ylabel(\"val_loss\")\n    \n    plt.tight_layout()\n    plt.show()\n\n    # Save best model of fold\n    fold_states_ms.append(best_state)\n    fold_best_acc.append(best_acc)\n\n    # OOF logits from best model\n    model.load_state_dict(best_state)\n    fold_val_logits = predict_logits(model, va_loader)\n    oof_logits_ms[va_idx] = fold_val_logits\n\nprint(\"\\nFold best acc:\", fold_best_acc)\nprint(\"CV mean:\", float(np.mean(fold_best_acc)), \"std:\", float(np.std(fold_best_acc)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:22:39.829629Z","iopub.execute_input":"2026-02-28T16:22:39.830460Z","iopub.status.idle":"2026-02-28T16:27:12.919612Z","shell.execute_reply.started":"2026-02-28T16:22:39.830429Z","shell.execute_reply":"2026-02-28T16:27:12.918595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"oof_pred_ms = oof_logits_ms.argmax(axis=1)\noof_acc_ms = (oof_pred_ms == oof_true).mean()\nprint(\"OOF accuracy MS:\", float(oof_acc_ms))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:33:46.898460Z","iopub.execute_input":"2026-02-28T16:33:46.898828Z","iopub.status.idle":"2026-02-28T16:33:46.904397Z","shell.execute_reply.started":"2026-02-28T16:33:46.898794Z","shell.execute_reply":"2026-02-28T16:33:46.903703Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# RGB","metadata":{}},{"cell_type":"code","source":"rgb_train_tf = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n    A.ToFloat(max_value=255.0),\n    ToTensorV2(),\n])\n\nrgb_val_tf = A.Compose([\n    A.ToFloat(max_value=255.0),\n    ToTensorV2(),\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:55:42.713259Z","iopub.execute_input":"2026-02-28T15:55:42.713505Z","iopub.status.idle":"2026-02-28T15:55:42.727068Z","shell.execute_reply.started":"2026-02-28T15:55:42.713482Z","shell.execute_reply":"2026-02-28T15:55:42.726183Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WheatRGBDataset(Dataset):\n    def __init__(self, df, is_train: bool = True, aug_tf=None, check_shape: bool = False):\n        df = df.reset_index(drop=True)\n        self.is_train = is_train\n        self.aug_tf = aug_tf\n        self.check_shape = check_shape\n\n        assert self.aug_tf is not None, \"For RGB baseline, pass aug_tf with Normalize + ToTensorV2.\"\n\n        self.rgb_paths = df[\"rgb_path\"].astype(str).tolist()\n        self.ids = df[\"Id\"].astype(str).tolist()\n\n        if is_train:\n            self.y = df[\"Category\"].map(cat2idx).to_numpy(np.int64)\n\n    def __len__(self):\n        return len(self.rgb_paths)\n\n    def _read_rgb(self, path) -> np.ndarray:\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)  # HWC uint8\n        return img\n\n    def __getitem__(self, idx):\n        img = self._read_rgb(self.rgb_paths[idx])\n\n        if self.check_shape:\n            if img.ndim != 3 or img.shape[2] != 3:\n                raise ValueError(f\"Unexpected RGB shape {img.shape} for {self.rgb_paths[idx]}\")\n\n        out = self.aug_tf(image=img)\n        x = out[\"image\"]  # torch.Tensor CHW float (после ToTensorV2)\n\n        if self.is_train:\n            y = torch.tensor(self.y[idx], dtype=torch.long)\n            return x, y, self.ids[idx]\n        else:\n            return x, self.ids[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:55:42.728945Z","iopub.execute_input":"2026-02-28T15:55:42.729261Z","iopub.status.idle":"2026-02-28T15:55:42.737406Z","shell.execute_reply.started":"2026-02-28T15:55:42.729238Z","shell.execute_reply":"2026-02-28T15:55:42.736803Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class RGBBackboneModel(nn.Module):\n    def __init__(self, backbone_name=\"resnet18\", num_classes=3, drop=0.2, pretrained=True):\n        super().__init__()\n\n        self.backbone = timm.create_model(\n            backbone_name,\n            pretrained=pretrained,\n            num_classes=0,   \n            global_pool=\"avg\"   \n        )\n\n        feat_dim = self.backbone.num_features\n        hidden = feat_dim // 2\n\n        self.head = nn.Sequential(\n            nn.Dropout(drop),\n            nn.Linear(feat_dim, hidden),\n            nn.ReLU(inplace=True),   # лучше для ResNet\n            nn.Dropout(drop),\n            nn.Linear(hidden, num_classes),\n        )\n\n    def forward(self, x):\n        feats = self.backbone(x)\n        return self.head(feats)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:55:42.738312Z","iopub.execute_input":"2026-02-28T15:55:42.738542Z","iopub.status.idle":"2026-02-28T15:55:42.750934Z","shell.execute_reply.started":"2026-02-28T15:55:42.738520Z","shell.execute_reply":"2026-02-28T15:55:42.750271Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# configs\nBATCH_SIZE = 32\nNUM_CLASSES = len(CLASSES)\nN_FOLDS = 5\nEPOCHS = 30\nlr_backbone = 1e-5\nlr_head = 3e-4\nWEIGHT_DECAY = 1e-2\nBACKBONE = \"resnet18\"\n\nuse_amp = False    #torch.cuda.is_available()\npin = torch.cuda.is_available()\n\ntrain_df_rgb = train_df.copy().reset_index(drop=True)\n\n# OOF containers\noof_logits_rgb = np.zeros((len(train_df_rgb), NUM_CLASSES), dtype=np.float32)\noof_true   = train_df_rgb[\"Category\"].map(cat2idx).values.astype(np.int64)\n\n# Store model weights per fold\nfold_states_rgb = []\nfold_best_acc = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:55:42.751615Z","iopub.execute_input":"2026-02-28T15:55:42.751908Z","iopub.status.idle":"2026-02-28T15:55:42.765705Z","shell.execute_reply.started":"2026-02-28T15:55:42.751887Z","shell.execute_reply":"2026-02-28T15:55:42.765024Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_loaders(\n    df_tr, df_va,\n    batch_size=BATCH_SIZE,\n    aug_tf_train=None,\n    aug_tf_val=None,\n    check_shape=False\n    ):\n    \n    tr_ds = WheatRGBDataset(\n        df_tr,\n        is_train=True,\n        aug_tf=aug_tf_train,\n        check_shape=check_shape\n        )\n    va_ds = WheatRGBDataset(\n        df_va,\n        is_train=True,\n        aug_tf=aug_tf_val,\n        check_shape=check_shape\n        )\n\n    g = torch.Generator()\n    g.manual_seed(SEED)\n\n    tr_loader = DataLoader(\n        tr_ds,\n        batch_size=batch_size,\n        shuffle=True,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        generator=g,\n        worker_init_fn=seed_worker\n    )\n    va_loader = DataLoader(\n        va_ds,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        worker_init_fn=seed_worker,\n        generator=g\n    )\n    return tr_loader, va_loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:55:42.766746Z","iopub.execute_input":"2026-02-28T15:55:42.767008Z","iopub.status.idle":"2026-02-28T15:55:42.778250Z","shell.execute_reply.started":"2026-02-28T15:55:42.766986Z","shell.execute_reply":"2026-02-28T15:55:42.777543Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\nPATIENCE = 15\n\nskf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\n\nfor fold, (tr_idx, va_idx) in enumerate(\n    skf.split(train_df_rgb, train_df_rgb[\"Category\"])\n):\n    df_tr = train_df_rgb.iloc[tr_idx].reset_index(drop=True)\n    df_va = train_df_rgb.iloc[va_idx].reset_index(drop=True)\n\n    tr_loader, va_loader = make_loaders(\n        df_tr, df_va,\n        batch_size=BATCH_SIZE,\n        aug_tf_train=rgb_train_tf,\n        aug_tf_val=rgb_val_tf,\n        check_shape=True\n    )\n\n    set_all_seeds(SEED + fold)    # To ensure determinism\n    model = RGBBackboneModel(\n        backbone_name=BACKBONE,\n        num_classes=NUM_CLASSES,\n        drop=0.2\n    ).to(device)\n\n    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)\n    scaler = torch.amp.GradScaler('cuda', enabled=use_amp)\n\n    optimizer = torch.optim.AdamW(\n        [\n            {\"params\": model.backbone.parameters(), \"lr\": lr_backbone},\n            {\"params\": model.head.parameters(),     \"lr\": lr_head},\n        ],\n        weight_decay=WEIGHT_DECAY\n    )\n\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer,\n        T_max=EPOCHS\n    )\n\n    best_acc = -1.0\n    best_state = None\n    bad_epochs = 0\n\n    print(f\"\\n===== FOLD {fold + 1} / {N_FOLDS} | \"\n          f\"train={len(df_tr)} valid={len(df_va)} =====\")\n\n    val_acc_hist = []\n    val_loss_hist = []\n\n    for epoch in range(1, EPOCHS + 1):\n\n        tr_loss, tr_acc = train_one_epoch(\n            model, tr_loader, optimizer, criterion, scaler\n        )\n\n        va_loss, va_acc = validate_one_epoch(\n            model, va_loader, criterion\n        )\n\n        val_loss_hist.append(float(va_loss))\n        val_acc_hist.append(float(va_acc))\n\n        scheduler.step()\n\n        if va_acc > best_acc:\n            best_acc = va_acc\n            best_state = {\n                k: v.detach().cpu()\n                for k, v in model.state_dict().items()\n            }\n            bad_epochs = 0\n        else:\n            bad_epochs += 1\n\n        current_lrs = [pg[\"lr\"] for pg in optimizer.param_groups]\n        lrs_str = \"[\" + \", \".join([f\"{lr:.2e}\" for lr in current_lrs]) + \"]\"\n\n        if epoch in (1, 2, 3, 5, 10, 15, 20) or epoch % 5 == 0:\n            print(\n                f\"Epoch {epoch:02d}/{EPOCHS} | \"\n                f\"train loss {tr_loss:.4f} acc {tr_acc:.4f} | \"\n                f\"valid loss {va_loss:.4f} acc {va_acc:.4f} | \"\n                f\"best accuracy {best_acc:.4f} | \"\n                f\"patience {bad_epochs}/{PATIENCE}\\n\"\n                f\"lrs {lrs_str}\"\n            )\n\n        if bad_epochs >= PATIENCE:\n            print(\n                f\"Early stopping on epoch {epoch} \"\n                f\"(no val_acc improvement for {PATIENCE} epochs).\"\n            )\n            break\n\n    epochs_ran = np.arange(1, len(val_acc_hist) + 1)\n\n    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n\n    axes[0].plot(epochs_ran, val_acc_hist)\n    axes[0].set_title(f\"Fold {fold+1}: Val Accuracy\")\n    axes[0].set_xlabel(\"epoch\")\n    axes[0].set_ylabel(\"val_acc\")\n\n    axes[1].plot(epochs_ran, val_loss_hist)\n    axes[1].set_title(f\"Fold {fold+1}: Val Loss\")\n    axes[1].set_xlabel(\"epoch\")\n    axes[1].set_ylabel(\"val_loss\")\n\n    plt.tight_layout()\n    plt.show()\n\n    fold_states_rgb.append(best_state)\n    fold_best_acc.append(best_acc)\n\n    model.load_state_dict(best_state)\n    fold_val_logits = predict_logits(model, va_loader)\n    oof_logits_rgb[va_idx] = fold_val_logits\n\n\nprint(\"\\nFold best acc:\", fold_best_acc)\nprint(\"CV mean:\", float(np.mean(fold_best_acc)),\n      \"std:\", float(np.std(fold_best_acc)))","metadata":{"execution":{"iopub.status.busy":"2026-02-28T15:55:42.779239Z","iopub.execute_input":"2026-02-28T15:55:42.779468Z","iopub.status.idle":"2026-02-28T15:57:46.302369Z","shell.execute_reply.started":"2026-02-28T15:55:42.779447Z","shell.execute_reply":"2026-02-28T15:57:46.301495Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"oof_pred_rgb = oof_logits_rgb.argmax(axis=1)\noof_acc_rgb = (oof_pred_rgb == oof_true).mean()\nprint(\"OOF accuracy RGB:\", float(oof_acc_rgb))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.303739Z","iopub.execute_input":"2026-02-28T15:57:46.304086Z","iopub.status.idle":"2026-02-28T15:57:46.309068Z","shell.execute_reply.started":"2026-02-28T15:57:46.304055Z","shell.execute_reply":"2026-02-28T15:57:46.308320Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# HS","metadata":{}},{"cell_type":"code","source":"hs_train_tf = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n\n    A.MultiplicativeNoise(multiplier=(0.95, 1.05), per_channel=False, p=0.10),\n ], p=1.0)\n\nhs_val_tf = A.Compose([], p=1.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.311351Z","iopub.execute_input":"2026-02-28T15:57:46.311635Z","iopub.status.idle":"2026-02-28T15:57:46.322764Z","shell.execute_reply.started":"2026-02-28T15:57:46.311600Z","shell.execute_reply":"2026-02-28T15:57:46.321946Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_hs_mean_std(\n    hs_paths,\n    drop_first=10,\n    drop_last=14,\n    keep_channels=125,\n    hs_max=65535.0,\n):\n\n    channel_sum = None\n    channel_sq_sum = None\n    total_pixels = 0\n\n    for path in tqdm(hs_paths):\n        with rasterio.open(path) as src:\n            x = src.read().astype(np.float32)  # (C,H,W)\n\n        # 126 -> 125\n        if x.shape[0] > keep_channels:\n            x = x[:keep_channels]\n\n        # trim noisy spectral ends\n        x = x[drop_first: keep_channels - drop_last]  # (101,H,W)\n\n        x = np.clip(x, 0.0, hs_max) / hs_max\n\n        C, H, W = x.shape\n        x = x.reshape(C, -1)  # (101, H*W)\n\n        if channel_sum is None:\n            channel_sum = x.sum(axis=1)\n            channel_sq_sum = (x ** 2).sum(axis=1)\n        else:\n            channel_sum += x.sum(axis=1)\n            channel_sq_sum += (x ** 2).sum(axis=1)\n\n        total_pixels += H * W\n\n    mean = channel_sum / total_pixels\n    var = channel_sq_sum / total_pixels - mean ** 2\n    std = np.sqrt(np.maximum(var, 1e-12))\n\n    return mean.astype(np.float32), std.astype(np.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.323747Z","iopub.execute_input":"2026-02-28T15:57:46.324064Z","iopub.status.idle":"2026-02-28T15:57:46.335623Z","shell.execute_reply.started":"2026-02-28T15:57:46.324039Z","shell.execute_reply":"2026-02-28T15:57:46.334990Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_hs_standardize_tf(mean, std):\n\n    mean = mean.reshape(-1, 1, 1)\n    std = std.reshape(-1, 1, 1)\n\n    def transform(x: np.ndarray) -> np.ndarray:\n        return (x - mean) / (std + 1e-8)\n\n    return transform","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.336495Z","iopub.execute_input":"2026-02-28T15:57:46.337104Z","iopub.status.idle":"2026-02-28T15:57:46.349768Z","shell.execute_reply.started":"2026-02-28T15:57:46.337072Z","shell.execute_reply":"2026-02-28T15:57:46.349122Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WheatHSDataset(Dataset):\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        is_train: bool = True,\n        pixel_tf=None,\n        aug_tf=None,\n        check_shape: bool = False,\n        drop_first: int = 10,\n        drop_last: int = 14,\n        keep_channels: int = 125,\n    ):\n        df = df.reset_index(drop=True)\n        self.is_train = is_train\n        self.pixel_tf = pixel_tf\n        self.aug_tf = aug_tf\n        self.check_shape = check_shape\n\n        self.drop_first = drop_first\n        self.drop_last = drop_last\n        self.keep_channels = keep_channels\n\n        self.hs_paths = df[\"hs_path\"].astype(str).tolist()\n        self.ids = df[\"Id\"].astype(str).tolist()\n\n        if is_train:\n            self.y = df[\"Category\"].map(cat2idx).to_numpy(np.int64)\n\n    def __len__(self):\n        return len(self.hs_paths)\n\n    def _read_hs(self, path) -> np.ndarray:\n        with rasterio.open(path) as src:\n            x = src.read().astype(np.float32)  # (C,H,W)\n        return x\n\n    def _fix_and_trim(self, x: np.ndarray) -> np.ndarray:\n        # fix 126 -> 125\n        if x.shape[0] > self.keep_channels:\n            x = x[:self.keep_channels]\n\n        # optional strict shape check\n        if self.check_shape:\n            if x.shape[1:] != (32, 32):\n                raise ValueError(f\"Unexpected HS spatial shape {x.shape} for {path}\")\n            if x.shape[0] not in (self.keep_channels, self.keep_channels + 1):\n                pass\n\n        # drop noisy spectral ends\n        if self.drop_last > 0:\n            x = x[self.drop_first: self.keep_channels - self.drop_last]\n        else:\n            x = x[self.drop_first:]\n\n        return x\n\n    def __getitem__(self, idx):\n        hs = self._read_hs(self.hs_paths[idx])  # (C,H,W)\n\n        # fix channels + drop edges\n        # (125 or 126) -> 125 -> 101\n        if hs.shape[0] == 126:\n            hs = hs[:125]\n        hs = hs[10:125-14]  # 101 channels\n\n        if self.check_shape and hs.shape[1:] != (32, 32):\n            raise ValueError(f\"Unexpected HS shape {hs.shape} for {self.hs_paths[idx]}\")\n\n        # Albumentations expects HWC\n        if self.aug_tf is not None:\n            hs_hwc = np.transpose(hs, (1, 2, 0))                 # (H,W,C)\n            hs_hwc = self.aug_tf(image=hs_hwc)[\"image\"]          # (H,W,C)\n            hs = np.transpose(hs_hwc, (2, 0, 1)).astype(np.float32)  # (C,H,W)\n\n        if self.pixel_tf is not None:\n            hs = self.pixel_tf(hs)\n\n        x = torch.from_numpy(hs)  # float32 CHW\n\n        if self.is_train:\n            y = torch.tensor(self.y[idx], dtype=torch.long)\n            return x, y, self.ids[idx]\n        else:\n            return x, self.ids[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.350670Z","iopub.execute_input":"2026-02-28T15:57:46.350919Z","iopub.status.idle":"2026-02-28T15:57:46.365622Z","shell.execute_reply.started":"2026-02-28T15:57:46.350896Z","shell.execute_reply":"2026-02-28T15:57:46.365048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def adapt_first_conv_in_chans(model: nn.Module, new_in_chans: int) -> nn.Module:\n    \"\"\"\n    Replace the very first Conv2d in a timm model to accept `new_in_chans`.\n    Init rule: take pretrained weights (out, 3, k, k) -> mean over input channels -> repeat to new_in_chans.\n    Also scales by (old_in / new_in) to keep activation magnitude roughly comparable.\n    \"\"\"\n    first_conv = None\n    first_conv_name = None\n\n    # find first Conv2d\n    for name, m in model.named_modules():\n        if isinstance(m, nn.Conv2d):\n            first_conv = m\n            first_conv_name = name\n            break\n\n    if first_conv is None:\n        raise RuntimeError(\"No Conv2d found in model.\")\n\n    old_in = first_conv.in_channels\n    if old_in == new_in_chans:\n        return model\n\n    if old_in != 3:\n        # still works, but you asked to average from existing channels (usually 3)\n        print(f\"[warn] First conv has in_channels={old_in}, not 3. Will still average from old_in.\")\n\n    # build new conv with same hyperparams\n    new_conv = nn.Conv2d(\n        in_channels=new_in_chans,\n        out_channels=first_conv.out_channels,\n        kernel_size=first_conv.kernel_size,\n        stride=first_conv.stride,\n        padding=first_conv.padding,\n        dilation=first_conv.dilation,\n        groups=first_conv.groups,   # convnext stem is usually groups=1\n        bias=(first_conv.bias is not None),\n        padding_mode=first_conv.padding_mode,\n    )\n\n    # init weights from pretrained\n    with torch.no_grad():\n        w = first_conv.weight  # (out, old_in, k, k)\n        w_mean = w.mean(dim=1, keepdim=True)                 # (out, 1, k, k)\n        w_new = w_mean.repeat(1, new_in_chans, 1, 1)         # (out, new_in, k, k)\n\n        # optional but usually helpful: keep activation scale similar\n        w_new *= (old_in / float(new_in_chans))\n\n        new_conv.weight.copy_(w_new)\n\n        if first_conv.bias is not None:\n            new_conv.bias.copy_(first_conv.bias)\n\n    # set it back into the model by name\n    parent = model\n    parts = first_conv_name.split(\".\")\n    for p in parts[:-1]:\n        parent = getattr(parent, p)\n    setattr(parent, parts[-1], new_conv)\n\n    return model\n\n\nclass HSStemConvNeXt(nn.Module):\n    \"\"\"\n    HS input: (B, 101, 32, 32)\n    Stem compresses 101 -> 16, then convnext_tiny adapted to in_chans=16.\n    \"\"\"\n    def __init__(self, backbone_name=\"convnext_tiny\", num_classes=3, drop=0.1, stem_width=64, out_chans=16):\n        super().__init__()\n\n        self.stem = nn.Sequential(\n            nn.Conv2d(101, stem_width, kernel_size=1, bias=False),\n            nn.BatchNorm2d(stem_width),\n            nn.GELU(),\n            nn.Conv2d(stem_width, out_chans, kernel_size=1, bias=False),\n            nn.BatchNorm2d(out_chans),\n            nn.GELU(),\n        )\n\n        # create 3-ch pretrained backbone, then adapt first conv to 16 input chans\n        self.backbone = timm.create_model(backbone_name, pretrained=True, num_classes=0)\n        self.backbone = adapt_first_conv_in_chans(self.backbone, new_in_chans=out_chans)\n\n        feat_dim = self.backbone.num_features\n        hidden = feat_dim // 2\n\n        self.head = nn.Sequential(\n            nn.LayerNorm(feat_dim),\n            nn.Dropout(drop),\n            nn.Linear(feat_dim, hidden),\n            nn.GELU(),\n            nn.Dropout(drop),\n            nn.Linear(hidden, num_classes),\n        )\n\n    def forward(self, x):\n        x = self.stem(x)          # (B, 16, H, W)\n        feats = self.backbone(x)  # (B, feat_dim)\n        return self.head(feats)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.366446Z","iopub.execute_input":"2026-02-28T15:57:46.366757Z","iopub.status.idle":"2026-02-28T15:57:46.382747Z","shell.execute_reply.started":"2026-02-28T15:57:46.366734Z","shell.execute_reply":"2026-02-28T15:57:46.382028Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_loaders_hs(\n    df_tr, df_va,\n    batch_size=BATCH_SIZE,\n    pixel_tf=None,\n    aug_tf_train=None,\n    aug_tf_val=None,\n    check_shape=False\n):\n    tr_ds = WheatHSDataset(df_tr, is_train=True, pixel_tf=pixel_tf, aug_tf=aug_tf_train, check_shape=check_shape)\n    va_ds = WheatHSDataset(df_va, is_train=True, pixel_tf=pixel_tf, aug_tf=aug_tf_val,   check_shape=check_shape)\n\n    g = torch.Generator()\n    g.manual_seed(SEED)\n\n    tr_loader = DataLoader(\n        tr_ds,\n        batch_size=batch_size,\n        shuffle=True,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        generator=g,\n        worker_init_fn=seed_worker\n    )\n    va_loader = DataLoader(\n        va_ds,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=3,\n        pin_memory=True,\n        persistent_workers=False,\n        prefetch_factor=2,\n        worker_init_fn=seed_worker,\n        generator=g\n    )\n    return tr_loader, va_loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T15:57:46.383544Z","iopub.execute_input":"2026-02-28T15:57:46.383835Z","iopub.status.idle":"2026-02-28T15:57:46.396275Z","shell.execute_reply.started":"2026-02-28T15:57:46.383801Z","shell.execute_reply":"2026-02-28T15:57:46.395708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# configs\nBATCH_SIZE = 64\nNUM_CLASSES = len(CLASSES)\nN_FOLDS = 5\nEPOCHS = 25\nlr_backbone = 2e-5\nlr_head = 1e-4\nWEIGHT_DECAY = 5e-2\nBACKBONE = \"convnext_tiny\"\n\nuse_amp = False    #torch.cuda.is_available()\npin = torch.cuda.is_available()\n\ntrain_df_hs = train_df.copy().reset_index(drop=True)\n\n# OOF containers\noof_logits_hs = np.zeros((len(train_df_hs), NUM_CLASSES), dtype=np.float32)\noof_true   = train_df_hs[\"Category\"].map(cat2idx).values.astype(np.int64)\n\n# Store model weights per fold\nfold_states_hs = []\nfold_best_acc = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:13:17.659793Z","iopub.execute_input":"2026-02-28T16:13:17.660129Z","iopub.status.idle":"2026-02-28T16:13:17.673103Z","shell.execute_reply.started":"2026-02-28T16:13:17.660095Z","shell.execute_reply":"2026-02-28T16:13:17.672256Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mean_hs, std_hs = compute_hs_mean_std(train_df_hs[\"hs_path\"].tolist())\npixel_hsmax_tf = make_hs_standardize_tf(mean_hs, std_hs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:13:24.291324Z","iopub.execute_input":"2026-02-28T16:13:24.291629Z","iopub.status.idle":"2026-02-28T16:13:29.404593Z","shell.execute_reply.started":"2026-02-28T16:13:24.291603Z","shell.execute_reply":"2026-02-28T16:13:29.403910Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\nPATIENCE = 10\n\nskf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\n\nfor fold, (tr_idx, va_idx) in enumerate(skf.split(train_df_hs, train_df_hs[\"Category\"])):\n    df_tr = train_df_hs.iloc[tr_idx].reset_index(drop=True)\n    df_va = train_df_hs.iloc[va_idx].reset_index(drop=True)\n        \n    tr_loader, va_loader = make_loaders_hs(\n        df_tr, df_va,\n        batch_size=BATCH_SIZE,\n        pixel_tf=pixel_hsmax_tf,\n        aug_tf_train=hs_train_tf,\n        aug_tf_val=None,\n        check_shape=True\n    )\n\n    set_all_seeds(SEED + fold)    # To ensure determinism\n    model = HSStemConvNeXt(\n        backbone_name=BACKBONE,\n        num_classes=NUM_CLASSES,\n        drop=0.1\n    ).to(device)\n\n    criterion = nn.CrossEntropyLoss()\n    scaler = torch.amp.GradScaler('cuda', enabled=use_amp)\n\n    optimizer = torch.optim.AdamW(\n        [\n            {\"params\": model.backbone.parameters(), \"lr\": lr_backbone},\n            {\"params\": model.stem.parameters(),     \"lr\": lr_head},\n            {\"params\": model.head.parameters(),     \"lr\": lr_head},\n        ],\n        weight_decay=WEIGHT_DECAY\n    )\n\n    # scheduler: reduce LR when val_acc plateaus\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer,\n        mode=\"max\",      \n        factor=0.5,\n        patience=2,\n        threshold=1e-3,\n        min_lr=1e-6\n    )\n\n    best_acc = -1.0\n    best_state = None\n    bad_epochs = 0\n\n    print(f\"\\n===== FOLD {fold + 1} / {N_FOLDS} | train={len(df_tr)} valid={len(df_va)} =====\")\n\n    val_acc_hist = []\n    val_loss_hist = []\n\n    for epoch in range(1, EPOCHS + 1):\n        tr_loss, tr_acc = train_one_epoch(model, tr_loader, optimizer, criterion, scaler)\n        va_loss, va_acc = validate_one_epoch(model, va_loader, criterion)\n\n        val_loss_hist.append(float(va_loss))\n        val_acc_hist.append(float(va_acc))\n\n        scheduler.step(va_acc)\n\n        if va_acc > best_acc:\n            best_acc = va_acc\n            best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}\n            bad_epochs = 0\n        else:\n            bad_epochs += 1\n\n        current_lrs = [pg[\"lr\"] for pg in optimizer.param_groups]\n        lrs_str = \"[\" + \", \".join([f\"{lr:.2e}\" for lr in current_lrs]) + \"]\"\n\n        if epoch in (1, 2, 3, 5, 10, 15, 20) or epoch % 5 == 0:\n            print(f\"Epoch {epoch:02d}/{EPOCHS} | \"\n                  f\"train loss {tr_loss:.4f} acc {tr_acc:.4f} | \"\n                  f\"valid loss {va_loss:.4f} acc {va_acc:.4f} | \"\n                  f\"best accuracy {best_acc:.4f} | \"\n                  f\"patience {bad_epochs}/{PATIENCE}\\n\"\n                  f\"lrs {lrs_str}\")\n\n        if bad_epochs >= PATIENCE:\n            print(f\"Early stopping on epoch {epoch} (no val_acc improvement for {PATIENCE} epochs).\")\n            break\n\n    # plots\n    epochs_ran = np.arange(1, len(val_acc_hist) + 1)\n\n    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n    axes[0].plot(epochs_ran, val_acc_hist)\n    axes[0].set_title(f\"Fold {fold+1}: Val Accuracy\")\n    axes[0].set_xlabel(\"epoch\")\n    axes[0].set_ylabel(\"val_acc\")\n\n    axes[1].plot(epochs_ran, val_loss_hist)\n    axes[1].set_title(f\"Fold {fold+1}: Val Loss\")\n    axes[1].set_xlabel(\"epoch\")\n    axes[1].set_ylabel(\"val_loss\")\n\n    plt.tight_layout()\n    plt.show()\n\n    fold_states_hs.append(best_state)\n    fold_best_acc.append(best_acc)\n\n    model.load_state_dict(best_state)\n    fold_val_logits = predict_logits(model, va_loader)\n    oof_logits_hs[va_idx] = fold_val_logits\n\nprint(\"\\nFold best acc:\", fold_best_acc)\nprint(\"CV mean:\", float(np.mean(fold_best_acc)), \"std:\", float(np.std(fold_best_acc)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:13:32.441714Z","iopub.execute_input":"2026-02-28T16:13:32.442470Z","iopub.status.idle":"2026-02-28T16:20:49.196338Z","shell.execute_reply.started":"2026-02-28T16:13:32.442439Z","shell.execute_reply":"2026-02-28T16:20:49.195365Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"oof_pred_hs = oof_logits_hs.argmax(axis=1)\noof_acc_hs = (oof_pred_hs == oof_true).mean()\nprint(\"OOF accuracy HS:\", float(oof_acc_hs))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T16:34:06.647019Z","iopub.execute_input":"2026-02-28T16:34:06.647599Z","iopub.status.idle":"2026-02-28T16:34:06.652326Z","shell.execute_reply.started":"2026-02-28T16:34:06.647567Z","shell.execute_reply":"2026-02-28T16:34:06.651492Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"def find_best_blend_weights(oof_logits_ms, oof_logits_rgb, oof_logits_hs, y_true, grid_n=21):\n    def acc_from_logits(logits, y):\n        return (logits.argmax(axis=1) == y).mean()\n\n    best_acc = -1.0\n    best_w = None\n\n    grid = np.linspace(0, 1, grid_n)\n    for w_ms in grid:\n        for w_rgb in grid:\n            if w_ms + w_rgb > 1.0:\n                continue\n            w_hs = 1.0 - w_ms - w_rgb\n\n            mix = w_ms * oof_logits_ms + w_rgb * oof_logits_rgb + w_hs * oof_logits_hs\n            acc = acc_from_logits(mix, y_true)\n\n            if acc > best_acc:\n                best_acc = acc\n                best_w = (float(w_ms), float(w_rgb), float(w_hs))\n\n    return best_acc, best_w","metadata":{"execution":{"iopub.status.busy":"2026-02-28T16:45:08.926528Z","iopub.execute_input":"2026-02-28T16:45:08.927165Z","iopub.status.idle":"2026-02-28T16:45:08.932540Z","shell.execute_reply.started":"2026-02-28T16:45:08.927136Z","shell.execute_reply":"2026-02-28T16:45:08.931704Z"}}},{"cell_type":"markdown","source":"best_acc, best_w = find_best_blend_weights(oof_logits_ms, oof_logits_rgb, oof_logits_hs, oof_true, grid_n=21)\nprint(\"Best ensemble OOF acc:\", best_acc)\nprint(\"Best weights (MS, RGB, HS):\", best_w)\n\nw_ms, w_rgb, w_hs = best_w","metadata":{"execution":{"iopub.status.busy":"2026-02-28T16:45:19.064157Z","iopub.execute_input":"2026-02-28T16:45:19.064902Z","iopub.status.idle":"2026-02-28T16:45:19.081372Z","shell.execute_reply.started":"2026-02-28T16:45:19.064863Z","shell.execute_reply":"2026-02-28T16:45:19.080592Z"}}},{"cell_type":"code","source":"w_ms, w_rgb, w_hs = 0.6, 0.0, 0.4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:13:40.550751Z","iopub.execute_input":"2026-02-28T17:13:40.551046Z","iopub.status.idle":"2026-02-28T17:13:40.554780Z","shell.execute_reply.started":"2026-02-28T17:13:40.551021Z","shell.execute_reply":"2026-02-28T17:13:40.554071Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"@torch.no_grad()\ndef predict_logits_cv_states(model_ctor, states, test_loader):\n    logits_sum = None\n\n    for st in states:\n        model = model_ctor().to(device)\n        model.load_state_dict(st)\n        model.eval()\n\n        fold_logits = []\n        for xb, ids in test_loader:\n            xb = xb.to(device, non_blocking=True)\n            logits = model(xb)\n            fold_logits.append(logits.cpu().numpy())\n\n        fold_logits = np.concatenate(fold_logits, axis=0)\n        logits_sum = fold_logits if logits_sum is None else (logits_sum + fold_logits)\n\n    return logits_sum / len(states)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:13:44.089829Z","iopub.execute_input":"2026-02-28T17:13:44.090551Z","iopub.status.idle":"2026-02-28T17:13:44.095365Z","shell.execute_reply.started":"2026-02-28T17:13:44.090519Z","shell.execute_reply":"2026-02-28T17:13:44.094799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_test_loader_ms(df_test, batch_size):\n    ds = WheatMSDataset(df_test, is_train=False, pixel_tf=pixel_minmax_tf, aug_tf=None, check_shape=True)\n    return DataLoader(ds, batch_size=batch_size, shuffle=False,\n                      num_workers=3, pin_memory=True, persistent_workers=False,\n                      prefetch_factor=2, worker_init_fn=seed_worker)\n\ndef make_test_loader_rgb(df_test, batch_size):\n    ds = WheatRGBDataset(df_test, is_train=False, aug_tf=rgb_val_tf, check_shape=True)\n    return DataLoader(ds, batch_size=batch_size, shuffle=False,\n                      num_workers=3, pin_memory=True, persistent_workers=False,\n                      prefetch_factor=2, worker_init_fn=seed_worker)\n\ndef make_test_loader_hs(df_test, batch_size):\n    ds = WheatHSDataset(df_test, is_train=False, pixel_tf=pixel_hsmax_tf, aug_tf=None, check_shape=True)\n    return DataLoader(\n        ds, batch_size=batch_size, shuffle=False,\n        num_workers=3, pin_memory=True, persistent_workers=False,\n        prefetch_factor=2, worker_init_fn=seed_worker\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:13:45.969928Z","iopub.execute_input":"2026-02-28T17:13:45.970513Z","iopub.status.idle":"2026-02-28T17:13:45.976027Z","shell.execute_reply.started":"2026-02-28T17:13:45.970483Z","shell.execute_reply":"2026-02-28T17:13:45.975388Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@torch.no_grad()\ndef predict_proba_ms_rgb_hs_ensemble(\n    df_test,\n    ms_states,\n    rgb_states,\n    hs_states,\n    w_ms=0.6,\n    w_rgb=0.0,\n    w_hs=0.4,\n    normalize_weights=True,\n    batch_size_ms=64,\n    batch_size_rgb=32,\n    batch_size_hs=64,\n):\n    ms_loader  = make_test_loader_ms(df_test, batch_size_ms)\n    rgb_loader = make_test_loader_rgb(df_test, batch_size_rgb)\n    hs_loader  = make_test_loader_hs(df_test, batch_size_hs)\n\n    def ms_ctor():\n        return MSStemConvNeXt(\n            backbone_name=\"convnext_tiny\",\n            num_classes=NUM_CLASSES,\n            drop=0.1\n        )\n\n    def rgb_ctor():\n        return RGBBackboneModel(\n            backbone_name=\"resnet18\",\n            num_classes=NUM_CLASSES,\n            drop=0.2\n        )\n\n    def hs_ctor():\n        return HSStemConvNeXt(\n            backbone_name=\"convnext_tiny\",\n            num_classes=NUM_CLASSES,\n            drop=0.1  \n        )\n\n        # logits per modality\n    logits_ms  = predict_logits_cv_states(ms_ctor,  ms_states,  ms_loader)\n    logits_rgb = predict_logits_cv_states(rgb_ctor, rgb_states, rgb_loader)\n    logits_hs  = predict_logits_cv_states(hs_ctor,  hs_states,  hs_loader)\n    \n    # blend logits \n    logits_mix = (\n            w_ms * logits_ms +\n            w_rgb * logits_rgb +\n            w_hs * logits_hs\n        )\n    \n    proba_mix = torch.softmax(torch.tensor(logits_mix), dim=1).numpy()\n    \n    pred_idx = proba_mix.argmax(axis=1)\n    pred_cat = [idx2cat[i] for i in pred_idx]\n    \n    return proba_mix, pred_cat","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:13:47.670953Z","iopub.execute_input":"2026-02-28T17:13:47.671217Z","iopub.status.idle":"2026-02-28T17:13:47.678400Z","shell.execute_reply.started":"2026-02-28T17:13:47.671195Z","shell.execute_reply":"2026-02-28T17:13:47.677738Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_proba, test_pred = predict_proba_ms_rgb_hs_ensemble(\n    df_test=val_df,              \n    ms_states=fold_states_ms,\n    rgb_states=fold_states_rgb,\n    hs_states=fold_states_hs,\n    w_ms=w_ms, w_rgb=w_rgb, w_hs=w_hs,\n    batch_size_ms=64,\n    batch_size_rgb=32,\n    batch_size_hs=64\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:13:52.501174Z","iopub.execute_input":"2026-02-28T17:13:52.501884Z","iopub.status.idle":"2026-02-28T17:14:16.058466Z","shell.execute_reply.started":"2026-02-28T17:13:52.501853Z","shell.execute_reply":"2026-02-28T17:14:16.057676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# calculation of dummy variables for a hidden test\ndef find_constant_triplet_ids(df: pd.DataFrame, eps: float = 1e-6) -> set[str]:\n    def is_const(path: str) -> bool:\n        with rasterio.open(path) as src:\n            x = src.read().astype(np.float32)\n        return float(x.std()) < eps\n\n    ids = set()\n    for _, row in df.iterrows():\n        if is_const(row[\"rgb_path\"]) and is_const(row[\"ms_path\"]) and is_const(row[\"hs_path\"]):\n            ids.add(str(row[\"Id\"]))\n    return ids\n\nempty_ids = find_constant_triplet_ids(val_df, eps=1e-6)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:14:16.060372Z","iopub.execute_input":"2026-02-28T17:14:16.061016Z","iopub.status.idle":"2026-02-28T17:14:17.093231Z","shell.execute_reply.started":"2026-02-28T17:14:16.060981Z","shell.execute_reply":"2026-02-28T17:14:17.092678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = pd.DataFrame({\n    \"Id\": val_df[\"Id\"].values,\n    \"Category\": test_pred\n})\nsubmission.loc[submission[\"Id\"].isin(empty_ids), \"Category\"] = \"Other\"\nsubmission[\"Id\"] = submission[\"Id\"].astype(str) + \".tif\"\nsubmission.to_csv(\"submission.csv\", index=False)\n\nsubmission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T17:14:17.094063Z","iopub.execute_input":"2026-02-28T17:14:17.094294Z","iopub.status.idle":"2026-02-28T17:14:17.106483Z","shell.execute_reply.started":"2026-02-28T17:14:17.094273Z","shell.execute_reply":"2026-02-28T17:14:17.105845Z"}},"outputs":[],"execution_count":null}]}