{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":660029,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":499177,"modelId":514426}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n#for dirname, _, filenames in os.walk('/kaggle/input'):\n #   for filename in filenames:\n   #     print(os.path.join(dirname, filename))\n#\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:21:41.588992Z","iopub.execute_input":"2025-11-27T16:21:41.589238Z","iopub.status.idle":"2025-11-27T16:21:44.783686Z","shell.execute_reply.started":"2025-11-27T16:21:41.589217Z","shell.execute_reply":"2025-11-27T16:21:44.782595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# full_train_infer.py\n# End-to-end train + inference pipeline\n# Replace TRAIN_CSV, IMAGE_ROOT, OUT_DIR with your paths\n\nimport os\nimport math\nimport time\nimport random\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.optim.swa_utils import AveragedModel, SWALR\n\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import train_test_split\n\n# ---------------------------\n# Config\n# ---------------------------\nclass CFG:\n    TRAIN_CSV = \"/kaggle/input/csiro-biomass/train.csv\"     \n    test_csv= \"/kaggle/input/csiro-biomass/test.csv\"  \n    # path to train CSV (image_path + targets)\n    IMAGE_ROOT = \"/kaggle/input/csiro-biomass/\"  # base folder for images\n    OUT_DIR = \"./outputs\"\n    os.makedirs(OUT_DIR, exist_ok=True)\n    base_model =\"/kaggle/input/convnext-moe/pytorch/default/1/outputs/best_convnext_moe.pth\"\n    swa_model = \"/kaggle/input/convnext-moe/pytorch/default/1/outputs/swa_convnext_moe.pth\"\n    MODEL_NAME = \"convnext_base\"  # timm model\n    IMG_H = 512\n    IMG_W = 1024\n    BATCH = 4\n    EPOCHS = 60\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    SEED = 42\n\n    # MoE head\n    MOE_EXPERTS = 3\n    MOE_HIDDEN = 1024\n    DROPOUT = 0.25\n\n    # optimizer / scheduler\n    BACKBONE_LR = 2e-5\n    HEAD_LR = 1e-3\n    WEIGHT_DECAY = 1e-2\n    WARMUP_EPOCHS = 5\n    SWA_START = 20   # when to start SWA (set > EPOCHS to skip)\n    SWA_FREQ = 1\n\n    # weights for weighted MSE (order: Dry_Clover, Dry_Dead, Dry_Green, Dry_Total, GDM)\n    LOSS_WEIGHTS = torch.tensor([0.1, 0.1, 0.1, 0.5, 0.2], dtype=torch.float32).to(DEVICE)\n\n    # val selection: fraction of total used as val\n    VAL_FRAC = 0.2\n\n    # clip grads\n    CLIP_GRAD = 1.0\n    TARGETS = [\"Dry_Clover_g\",\"Dry_Dead_g\",\"Dry_Green_g\",\"Dry_Total_g\",\"GDM_g\"]\n\n    \n    # save\n    BEST_MODEL_PATH = os.path.join(OUT_DIR, \"best_convnext_moe.pth\")\n    SWA_MODEL_PATH = os.path.join(OUT_DIR, \"swa_convnext_moe.pth\")\n    EMA_DECAY = 0.999\n\n  #  @staticmethod\n    \n    train_transform = A.Compose([\n      A.Resize(IMG_H, IMG_W),\n\n    # geometry\n      A.HorizontalFlip(p=0.5),\n      A.VerticalFlip(p=0.1),\n      A.ShiftScaleRotate(\n        shift_limit=0.02,\n        scale_limit=0.0,\n        rotate_limit=10,\n        border_mode=cv2.BORDER_REFLECT_101,\n        p=0.4,\n      ),\n\n    # color\n      A.RandomBrightnessContrast(p=0.3),\n      A.HueSaturationValue(p=0.2),\n      A.RGBShift(p=0.1),\n\n    # noise\n      A.GaussNoise(p=0.2),\n      A.MotionBlur(blur_limit=3, p=0.1),\n\n    # normalization + tensor\n      A.Normalize(\n        mean=(0.485, 0.456, 0.406),\n        std=(0.229, 0.224, 0.225)\n      ),\n      ToTensorV2(),\n      ])\n\n    \n    \n\n    valid_transform = A.Compose([\n      A.Resize(IMG_H, IMG_W),\n      A.Normalize(mean=(0.485,0.456,0.406), std=(0.229,0.224,0.225)),\n            ToTensorV2()\n      ])\n\n\n# reproducibility\ndef seed_everything(seed=CFG.SEED):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\nseed_everything()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:21:44.785891Z","iopub.execute_input":"2025-11-27T16:21:44.786383Z","iopub.status.idle":"2025-11-27T16:22:46.252037Z","shell.execute_reply.started":"2025-11-27T16:21:44.786357Z","shell.execute_reply":"2025-11-27T16:22:46.250547Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# Data, transforms, dataset\n# ---------------------------\n\n\nclass WholeImageDataset(Dataset):\n    def __init__(self, df, image_root, transform=None, return_meta=False):\n        self.df = df.reset_index(drop=True)\n        self.root = Path(image_root)\n        self.transform = transform\n        self.return_meta = return_meta\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.root / row[\"image_path\"]\n        img = cv2.imread(str(img_path))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        if self.transform:\n            img = self.transform(image=img)[\"image\"]\n        if set(CFG.TARGETS).issubset(self.df.columns):\n            target = torch.tensor([row[t] for t in CFG.TARGETS], dtype=torch.float32)\n            if self.return_meta:\n                return img, target, row\n            else:\n                return img, target\n        else:\n            return img\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.253178Z","iopub.execute_input":"2025-11-27T16:22:46.254050Z","iopub.status.idle":"2025-11-27T16:22:46.264837Z","shell.execute_reply.started":"2025-11-27T16:22:46.254017Z","shell.execute_reply":"2025-11-27T16:22:46.263395Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# Losses and metrics\n# ---------------------------\nclass WeightedMSELoss(nn.Module):\n    def __init__(self, weights):\n        super().__init__()\n        self.weights = weights\n    def forward(self, pred, target):\n        se = (pred - target)**2  # (B,5)\n        se = se * self.weights   # broadcast\n        return se.mean()\n\ndef per_target_rmse(preds, gts):\n    se = (preds - gts)**2\n    mse = se.mean(axis=0)\n    rmse = np.sqrt(mse)\n    return rmse\n\ndef weighted_r2(gt, pred, weights_vec):\n    # flatten to rows (image-target pairs)\n    # gt, pred shape (N,5)\n    w = np.repeat(weights_vec.cpu().numpy()[None, :], gt.shape[0], axis=0)  # (N,5)\n    gt_flat = gt.reshape(-1)\n    pred_flat = pred.reshape(-1)\n    w_flat = w.reshape(-1)\n    # weighted mean of gt\n    mean_gt = np.sum(w_flat * gt_flat) / np.sum(w_flat)\n    rss = np.sum(w_flat * (gt_flat - pred_flat)**2)\n    tss = np.sum(w_flat * (gt_flat - mean_gt)**2)\n    return 1.0 - (rss / tss)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.266205Z","iopub.execute_input":"2025-11-27T16:22:46.266620Z","iopub.status.idle":"2025-11-27T16:22:46.289374Z","shell.execute_reply.started":"2025-11-27T16:22:46.266587Z","shell.execute_reply":"2025-11-27T16:22:46.287871Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# Model: Backbone + MoE Head\n# ---------------------------\nclass MetaMoE(nn.Module):\n    def __init__(self, in_dim, hidden=CFG.MOE_HIDDEN, num_experts=CFG.MOE_EXPERTS, out_dim=5, dropout=CFG.DROPOUT):\n        super().__init__()\n        self.in_dim = in_dim\n        self.num_experts = num_experts\n        self.gate = nn.Sequential(\n            nn.LayerNorm(in_dim),\n            nn.Linear(in_dim, 256),\n            nn.ReLU(),\n            nn.Linear(256, num_experts)\n        )\n        self.experts = nn.ModuleList([\n            nn.Sequential(\n                nn.LayerNorm(in_dim),\n                nn.Linear(in_dim, hidden),\n                nn.ReLU(),\n                nn.Dropout(dropout),\n                nn.Linear(hidden, 256),\n                nn.ReLU(),\n                nn.Linear(256, out_dim)\n            ) for _ in range(num_experts)\n        ])\n    def forward(self, x):\n        gate_logits = self.gate(x)               # (B, E)\n        gate_weights = F.softmax(gate_logits, dim=1).unsqueeze(-1)  # (B, E, 1)\n        expert_outs = torch.stack([e(x) for e in self.experts], dim=1)  # (B, E, out_dim)\n        out = (expert_outs * gate_weights).sum(dim=1)  # (B, out_dim)\n        return out\n\nclass ConvNextWithMoE(nn.Module):\n    def __init__(self, model_name=CFG.MODEL_NAME, pretrained=False):\n        super().__init__()\n        backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool=\"avg\")\n        self.backbone = backbone\n        self.feat_dim = backbone.num_features\n        self.head = MetaMoE(in_dim=self.feat_dim, hidden=CFG.MOE_HIDDEN, num_experts=CFG.MOE_EXPERTS, out_dim=len(CFG.TARGETS))\n    def forward(self, x):\n        feat = self.backbone(x)    # (B, feat_dim)\n        out = self.head(feat)\n        return out\n    def forward_features(self, x):\n        return self.backbone(x)\n\n\n# ---------------------------\n# EMA helper\n# ---------------------------\nclass ModelEMA:\n    def __init__(self, model, decay=CFG.EMA_DECAY):\n        self.decay = decay\n        self.shadow = {}\n        self.device = next(model.parameters()).device\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.detach().cpu().clone()\n    def update(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                new = param.detach().cpu()\n                self.shadow[name].mul_(self.decay)\n                self.shadow[name].add_( (1.0 - self.decay) * new )\n    def apply_to(self, model):\n        # copy shadow to model (use for eval or saving)\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                param.data.copy_(self.shadow[name].to(param.device))\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.290425Z","iopub.execute_input":"2025-11-27T16:22:46.290764Z","iopub.status.idle":"2025-11-27T16:22:46.317776Z","shell.execute_reply.started":"2025-11-27T16:22:46.290741Z","shell.execute_reply":"2025-11-27T16:22:46.316751Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_ema_state_dict(model, ema):\n    \"\"\"\n    Build a state_dict containing the EMA weights,\n    matching the real model's state_dict structure.\n    \"\"\"\n    ema_sd = {}\n\n    model_sd = model.state_dict()\n    \n    for name, tensor in model_sd.items():\n        # Only replace actual trainable parameters tracked by EMA\n        if name in ema.shadow:\n            ema_sd[name] = ema.shadow[name].clone()\n        else:\n            # Buffers (running_mean, running_var, etc.)\n            ema_sd[name] = tensor.clone()\n    \n    return ema_sd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.318811Z","iopub.execute_input":"2025-11-27T16:22:46.319148Z","iopub.status.idle":"2025-11-27T16:22:46.348762Z","shell.execute_reply.started":"2025-11-27T16:22:46.319117Z","shell.execute_reply":"2025-11-27T16:22:46.347452Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# Utility: choose validation indices from dense region around Dry_Total mean\n# ---------------------------\ndef choose_dense_val_indices(df, targets=[\"Dry_Total_g\", \"GDM_g\",\n                                                   \"Dry_Green_g\", \"Dry_Dead_g\",\n                                                   \"Dry_Clover_g\"], \n                             val_frac=CFG.VAL_FRAC, q=10, mid_quantile=0.4):\n    \"\"\"\n    Select validation samples that are in the middle-density region\n    for *all* targets simultaneously.\n\n    Ensures:\n        - All extremes located in train only\n        - Val contains only mid-range dense rows\n        - More stable than single-target stratification\n    \"\"\"\n\n    df = df.copy()\n    n = len(df)\n    target_val_count = int(np.ceil(val_frac * n))\n\n    # ---------------------------------------------------\n    # STEP 1: For each target create q quantile bins\n    # ---------------------------------------------------\n    for t in targets:\n        df[f\"{t}_bin\"] = pd.qcut(df[t], q=q, duplicates=\"drop\", labels=False)\n\n    # ---------------------------------------------------\n    # STEP 2: Define mid-range bins for each target\n    #   e.g. for q=10 → allowed bins are [3,4,5,6]\n    # ---------------------------------------------------\n    low = int((1 - mid_quantile) / 2 * q)      # ≈ 3\n    high = int((1 + mid_quantile) / 2 * q)     # ≈ 6\n    # clamp\n    low = max(0, low)\n    high = min(q - 1, high)\n\n    # ---------------------------------------------------\n    # STEP 3: Keep only rows whose bins lie inside mid-range for ALL targets\n    # ---------------------------------------------------\n    cond = np.ones(len(df), dtype=bool)\n    for t in targets:\n        cond &= (df[f\"{t}_bin\"] >= low) & (df[f\"{t}_bin\"] <= high)\n\n    mid_indices = df[cond].index.tolist()\n\n    # ---------------------------------------------------\n    # STEP 4: Sample val_frac from the mid indices\n    # ---------------------------------------------------\n    if len(mid_indices) >= target_val_count:\n        rng = np.random.default_rng(42)\n        val_idx = rng.choice(mid_indices, size=target_val_count, replace=False)\n    else:\n        # if middle area too small — fallback: pick all mids\n        val_idx = mid_indices\n\n    val_idx = sorted(val_idx)\n    train_idx = sorted(list(set(range(n)) - set(val_idx)))\n\n    return train_idx, val_idx\n# ---------------------------\n# Scheduler helper: warmup + cosine\n# ---------------------------\ndef get_scheduler_with_warmup(optimizer, total_epochs, warmup_epochs):\n    def lr_lambda(epoch):\n        if epoch < warmup_epochs:\n            return float(epoch + 1) / float(max(1, warmup_epochs))\n        # cosine after warmup\n        progress = float(epoch - warmup_epochs) / float(max(1, total_epochs - warmup_epochs))\n        return 0.5 * (1.0 + math.cos(math.pi * progress))\n    return LambdaLR(optimizer, lr_lambda=lr_lambda)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.351629Z","iopub.execute_input":"2025-11-27T16:22:46.351919Z","iopub.status.idle":"2025-11-27T16:22:46.372858Z","shell.execute_reply.started":"2025-11-27T16:22:46.351897Z","shell.execute_reply":"2025-11-27T16:22:46.371554Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MoECollapseLoss:\n    def __init__(self, ent_weight=1e-3, kl_weight=1e-3, collapse_thresh=0.70):\n        self.ent_weight = ent_weight\n        self.kl_weight = kl_weight\n        self.collapse_thresh = collapse_thresh\n        self.reset_counter = 0\n\n    def entropy(self, gate):\n        # gate: (B,E)\n        return (-gate * torch.log(gate + 1e-9)).sum(dim=1).mean()\n\n    def kl_uniform(self, gate):\n        # KL(p||U), U = 1/E\n        E = gate.shape[1]\n        uniform = 1.0 / E\n        return (gate * torch.log((gate + 1e-9) / uniform)).sum(dim=1).mean()\n\n    def collapse_detected(self, gate):\n        # % samples with max gate > 0.9\n        collapse_pct = (gate.max(dim=1).values > 0.9).float().mean().item()\n        return collapse_pct > self.collapse_thresh, collapse_pct\n\n    def compute(self, gate):\n        ent = self.entropy(gate)\n        kl  = self.kl_uniform(gate)\n        return ent, kl\n\ndef reset_gate_head(gate_module):\n    for name, p in gate_module.named_parameters():\n        if \"weight\" in name:\n            nn.init.xavier_uniform_(p)\n        elif \"bias\" in name:\n            nn.init.zeros_(p)\n    print(\"⚠️ Gate RESET triggered — collapse protection\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.373857Z","iopub.execute_input":"2025-11-27T16:22:46.374143Z","iopub.status.idle":"2025-11-27T16:22:46.404672Z","shell.execute_reply.started":"2025-11-27T16:22:46.374120Z","shell.execute_reply":"2025-11-27T16:22:46.403311Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_one_fold(train_df, val_df):\n    device = CFG.DEVICE\n\n    train_ds = WholeImageDataset(train_df, CFG.IMAGE_ROOT,\n                                 transform=CFG.train_transform, return_meta=False)\n    val_ds = WholeImageDataset(val_df, CFG.IMAGE_ROOT,\n                               transform=CFG.valid_transform, return_meta=False)\n\n    train_loader = DataLoader(train_ds, batch_size=CFG.BATCH,\n                              shuffle=True, num_workers=4, pin_memory=True)\n    val_loader = DataLoader(val_ds, batch_size=CFG.BATCH,\n                            shuffle=False, num_workers=4, pin_memory=True)\n\n    # Model\n    model = ConvNextWithMoE(model_name=CFG.MODEL_NAME).to(device)\n\n    # Optim\n    opt = torch.optim.AdamW([\n        {\"params\": model.backbone.parameters(), \"lr\": CFG.BACKBONE_LR},\n        {\"params\": model.head.parameters(),     \"lr\": CFG.HEAD_LR},\n    ], weight_decay=CFG.WEIGHT_DECAY)\n\n    scheduler = get_scheduler_with_warmup(opt, \n                                          total_epochs=CFG.EPOCHS,\n                                          warmup_epochs=CFG.WARMUP_EPOCHS)\n\n    loss_fn = WeightedMSELoss(CFG.LOSS_WEIGHTS)\n\n    # EMA\n    ema = ModelEMA(model, decay=CFG.EMA_DECAY)\n\n    # SWA\n    use_swa = CFG.SWA_START < CFG.EPOCHS\n    swa_model = AveragedModel(model) if use_swa else None\n\n    # Best tracking\n    best_val_loss = float('inf')\n    best_train_loss = float('inf')\n    best_epoch = -1\n\n    gen_bad_epochs = 0  # <-- FIXED\n\n    # ----------------------- TRAIN LOOP -----------------------\n    for epoch in range(CFG.EPOCHS):\n        model.train()\n        running_loss = 0.0\n        batches = 0\n\n        pbar = tqdm(train_loader, desc=f\"Train E{epoch+1}/{CFG.EPOCHS}\")\n\n        for imgs, tgts in pbar:\n            imgs = imgs.to(device)\n            tgts = tgts.to(device)\n\n            opt.zero_grad()\n            preds = model(imgs)\n            loss = loss_fn(preds, tgts)\n            loss.backward()\n\n            torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.CLIP_GRAD)\n            opt.step()\n\n            # EMA update AFTER optimizer step\n            ema.update(model)\n\n            running_loss += loss.item()\n            batches += 1\n            pbar.set_postfix(train_loss=running_loss / batches)\n\n            del imgs, tgts, preds\n\n        train_loss = running_loss / max(1, batches)\n        torch.cuda.empty_cache()\n\n        # ----------------------- VALIDATION -----------------------\n        model.eval()\n        val_losses = []\n\n        with torch.no_grad():\n            for imgs, tgts in tqdm(val_loader, desc=\"Val\"):\n                imgs = imgs.to(device)\n                tgts = tgts.to(device)\n                preds = model(imgs)\n\n                l = loss_fn(preds, tgts)\n                val_losses.append(l.item())\n            \n        val_loss = float(np.mean(val_losses))\n        torch.cuda.empty_cache()\n\n        print(f\"Val Loss: {val_loss:.5f} | Train Loss: {train_loss:.5f}\")\n\n        gen_ratio = val_loss / (train_loss + 1e-12)\n\n        # --------------------- SAVE BEST MODEL ---------------------\n        if train_loss < best_train_loss:\n\n            best_val_loss = val_loss\n            best_train_loss = train_loss\n            best_epoch = epoch + 1\n\n            # Save EMA shadow weights WITHOUT touching the model\n            ema_state = get_ema_state_dict(model, ema)\n\n            torch.save({\n                \"ema_state_dict\": ema_state,\n                \"model_state_dict\": model.state_dict(),\n                \"epoch\": best_epoch,\n                \"opt_state\": opt.state_dict(),\n            }, CFG.BEST_MODEL_PATH)\n\n            if gen_ratio < 1.2:\n                torch.save({\n                    \"ema_state_dict\": ema_state,\n                    \"model_state_dict\": model.state_dict(),\n                    \"epoch\": best_epoch,\n                    \"opt_state\": opt.state_dict(),\n                    }, \"gen_best_model.pth\")\n            print(f\"✔ Saved BEST EMA model at epoch {best_epoch}\")\n\n        # --------------------- SWA UPDATE ---------------------\n        if use_swa and epoch >= CFG.SWA_START:\n            swa_model.update_parameters(model)\n\n        # --------------------- SCHEDULER STEP ---------------------\n        scheduler.step()\n\n        # --------------------- EARLY STOPPING ---------------------\n        if gen_ratio > 1.5:\n            gen_bad_epochs += 1\n        else: \n            gen_bad_epochs = 0\n\n        if gen_bad_epochs > 5:\n            print(\"High gen_ratio:\", gen_ratio)\n           # break\n\n    # ------------------------ SWA FINALIZE ------------------------\n    if use_swa:\n        torch.save({\n            \"model_state_dict\": swa_model.module.state_dict(),\n            \"epoch\": epoch + 1,\n        }, CFG.SWA_MODEL_PATH)\n        print(\"✔ Saved SWA model.\")\n\n    return CFG.BEST_MODEL_PATH\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.406003Z","iopub.execute_input":"2025-11-27T16:22:46.406287Z","iopub.status.idle":"2025-11-27T16:22:46.441410Z","shell.execute_reply.started":"2025-11-27T16:22:46.406264Z","shell.execute_reply":"2025-11-27T16:22:46.440253Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# Inference (single or ensemble)\n# ---------------------------\ndef load_model_checkpoint(path, device=CFG.DEVICE, use_ema=True):\n    ck = torch.load(path, map_location=device)\n    model = ConvNextWithMoE(model_name=CFG.MODEL_NAME).to(device)\n\n    if use_ema and \"ema_state_dict\" in ck:\n        print(\">> Loading EMA weights\")\n        model.load_state_dict(ck[\"ema_state_dict\"])\n    else:\n        print(\">> Loading RAW model weights\")\n        model.load_state_dict(ck[\"model_state_dict\"])\n\n    model.eval()\n    return model\n\n\n@torch.no_grad()\ndef predict_image_single(model, img_path, image_root=CFG.IMAGE_ROOT):\n    tfm = CFG.valid_transform\n    img = cv2.imread(str(Path(image_root)/img_path))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    inp = tfm(image=img)[\"image\"].unsqueeze(0).to(CFG.DEVICE)\n    out = model(inp).cpu().numpy()[0]   # raw predictions\n    return out\n\n@torch.no_grad()\ndef predict_ensemble(model_paths, img_path, image_root=CFG.IMAGE_ROOT, use_ema=True):\n    preds = []\n    for p in model_paths:\n        model = load_model_checkpoint(p, use_ema=use_ema)\n        preds.append(predict_image_single(model, img_path, image_root))\n    preds = np.vstack(preds)\n    final = preds.mean(axis=0)\n    return final\n\ndef create_submission(model_paths, test_csv, out_csv=\"submission.csv\"):\n    \"\"\"\n    test_csv must have columns: sample_id, image_path, target_name\n    We'll predict per-image a 5-vector and then expand to one row per image-target pair.\n    \"\"\"\n    test_df = pd.read_csv(test_csv)\n    # unique images\n    images = test_df[\"image_path\"].unique().tolist()\n    image_preds = {}\n    for img in tqdm(images, desc=\"Predict images\"):\n        pred = predict_ensemble(model_paths, img)\n        image_preds[img] = pred\n\n    rows = []\n    for _, row in test_df.iterrows():\n        img = row[\"image_path\"]\n        target_name = row[\"target_name\"]\n        # map target positions\n        idx_map = { \"Dry_Clover_g\":0, \"Dry_Dead_g\":1, \"Dry_Green_g\":2, \"Dry_Total_g\":3, \"GDM_g\":4 }\n        val = float(max(0.0, image_preds[img][idx_map[target_name]]))\n        rows.append({\"sample_id\": row[\"sample_id\"], \"target\": val})\n    out = pd.DataFrame(rows)\n    out.to_csv(out_csv, index=False)\n    print(\"Saved submission:\", out_csv)\n    print(out.head())\n\n# ---------------------------\n# Entrypoint\n# ---------------------------\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.442896Z","iopub.execute_input":"2025-11-27T16:22:46.443264Z","iopub.status.idle":"2025-11-27T16:22:46.481977Z","shell.execute_reply.started":"2025-11-27T16:22:46.443240Z","shell.execute_reply":"2025-11-27T16:22:46.480939Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model_paths = [CFG.base_model]  # add more .pths to average if you have them\ncreate_submission(model_paths, CFG.test_csv, out_csv=\"submission.csv\")\nprint(\"Done.\")\n\n\n\"\"\"\n# ---------------------------\n# Entrypoint\n# ---------------------------\nif __name__ == \"__main__\":\n    # Load CSV and build train/val split where val is from dense region near Dry_Total \n    # Example inference: create submission using best saved model (single) or ensemble (list)\n  \n    \n    train = pd.read_csv(CFG.TRAIN_CSV)\n    # Ensure image_path is relative path under IMAGE_ROOT\n    # Pivot to wide format\n    df = train.pivot_table(\n     index=['image_path', 'Sampling_Date', 'State', 'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm'],\n     columns='target_name',\n     values='target'\n     ).reset_index()\n    \n    print(\"Loaded\", len(df), \"rows\")\n\n    train_idx, val_idx = choose_dense_val_indices(df)\n    train_df = df.loc[train_idx].reset_index(drop=True)\n    val_df = df.loc[val_idx].reset_index(drop=True)\n    print(f\"Train rows: {len(train_df)} Val rows (dense): {len(val_df)}\")\n    print(f\"Train rows: {train_df.describe()} Val rows (dense): {val_df.describe()}\")\n\n    # Train single fold with our procedure\n    best_path = train_one_fold(train_df, val_df)\n\n    model_paths = [base_model,swa_model]  # add more .pths to average if you have them\n    create_submission(model_paths, \"test.csv\", out_csv=\"submission.csv\")\n    print(\"Done.\")\n    \"\"\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:22:46.483732Z","iopub.execute_input":"2025-11-27T16:22:46.484144Z","iopub.status.idle":"2025-11-27T16:23:10.365399Z","shell.execute_reply.started":"2025-11-27T16:22:46.484121Z","shell.execute_reply":"2025-11-27T16:23:10.364667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"score: 0.62 submission\n                    sample_id     target\n0  ID1001187975__Dry_Clover_g   1.289032\n1    ID1001187975__Dry_Dead_g  26.302822\n2   ID1001187975__Dry_Green_g  24.061487\n3   ID1001187975__Dry_Total_g  52.621857\n4         ID1001187975__GDM_g  25.638224\n\nscore: 0.6 submission\n\n\"\"\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T16:23:10.366222Z","iopub.execute_input":"2025-11-27T16:23:10.366526Z","iopub.status.idle":"2025-11-27T16:23:10.372515Z","shell.execute_reply.started":"2025-11-27T16:23:10.366504Z","shell.execute_reply":"2025-11-27T16:23:10.371485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}