{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":286978659,"sourceType":"kernelVersion"},{"sourceId":287059427,"sourceType":"kernelVersion"},{"sourceId":693751,"sourceType":"modelInstanceVersion","modelInstanceId":526083,"modelId":540138},{"sourceId":693802,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":526083,"modelId":540138}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"code","source":"#show timm version\nimport timm\nprint(timm.__version__)","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:29:44.960411Z","iopub.execute_input":"2025-12-21T01:29:44.960704Z","iopub.status.idle":"2025-12-21T01:29:55.790333Z","shell.execute_reply.started":"2025-12-21T01:29:44.960679Z","shell.execute_reply":"2025-12-21T01:29:55.789736Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Config\nimport os, gc, math, cv2, numpy as np, pandas as pd\nfrom tqdm import tqdm\nimport torch, torch.nn as nn, torch.optim as optim\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom timm.utils import ModelEmaV2\nfrom sklearn.model_selection import KFold, StratifiedGroupKFold\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nclass CFG:\n    CREATE_SUBMISSION = True\n    USE_TQDM        = False\n    PRETRAINED_DIR  = None\n    PRETRAINED      = True\n    BASE_PATH       = '/kaggle/input/csiro-biomass'\n    SEED            = 82947501\n    FOLDS_TO_TRAIN   = [0,1,2,3,4]\n    TRAIN_CSV       = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    TEST_IMAGE_DIR = '/kaggle/input/csiro-biomass/test'\n    TEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\n    SUBMISSION_DIR  = '/kaggle/working/'\n    MODEL_DIR_012       = '/kaggle/input/5-folds-dinov3-840m/other/fold-0-1-2/1'\n    MODEL_DIR_34       = '/kaggle/input/5-folds-dinov3-840m/other/fold-0-1-2/2'\n    N_FOLDS         = 5\n\n\n    # MODEL_NAME      = 'vit_large_patch16_dinov3.lvd1689m'  \n    # BACKBONE_PATH   = '/kaggle/input/vit-large-patch16-dinov3-lvd1689m-backbone-pth/vit_large_patch16_dinov3.lvd1689m_backbone.pth'\n    MODEL_NAME      = 'vit_huge_plus_patch16_dinov3.lvd1689m'\n    BACKBONE_PATH   = '/kaggle/input/vit-huge-plus-patch16-dinov3-lvd1689m/vit_huge_plus_patch16_dinov3.lvd1689m_backbone.pth'\n    \n    IMG_SIZE        = 512\n\n    VAL_TTA_TIMES   = 1\n    TTA_STEPS       = 1\n    \n    \n    BATCH_SIZE      = 1\n    GRAD_ACC        = 4\n    NUM_WORKERS     = 4\n    EPOCHS          = 1\n    FREEZE_EPOCHS   = 0\n    WARMUP_EPOCHS   = 3\n    LR_REST         = 1e-3\n    LR_BACKBONE     = 5e-4\n    WD              = 1e-2\n    EMA_DECAY       = 0.9\n    PATIENCE        = 5\n    TARGET_COLS     = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS    = ['Dry_Clover_g', 'Dry_Dead_g']\n    ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    R2_WEIGHTS      = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\n    LOSS_WEIGHTS    = np.array([0.1, 0.1, 0.1, 0.0, 0.0])\n    DEVICE          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f'Device : {CFG.DEVICE}')\nprint(f'Backbone: {CFG.MODEL_NAME} | Input: {CFG.IMG_SIZE}')\nprint(f'Freeze Epochs: {CFG.FREEZE_EPOCHS} | Warmup: {CFG.WARMUP_EPOCHS}')\nprint(f'EMA Decay: {CFG.EMA_DECAY} | Grad Acc: {CFG.GRAD_ACC}')","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:29:55.791720Z","iopub.execute_input":"2025-12-21T01:29:55.791942Z","iopub.status.idle":"2025-12-21T01:30:29.866574Z","shell.execute_reply.started":"2025-12-21T01:29:55.791921Z","shell.execute_reply":"2025-12-21T01:30:29.865914Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Metric","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport os\n\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    r2_scores = []\n    for i in range(y_true.shape[1]):\n        yt = y_true[:, i]; yp = y_pred[:, i]\n        ss_res = np.sum((yt - yp) ** 2)\n        ss_tot = np.sum((yt - np.mean(yt)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n    r2_scores = np.array(r2_scores)\n    weighted = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted, r2_scores\n\ndef weighted_r2_score_global(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    flat_true = y_true.reshape(-1)\n    flat_pred = y_pred.reshape(-1)\n    w = np.tile(weights, y_true.shape[0])\n    mean_w = np.sum(w * flat_true) / np.sum(w)\n    ss_res = np.sum(w * (flat_true - flat_pred) ** 2)\n    ss_tot = np.sum(w * (flat_true - mean_w) ** 2)\n    global_r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n    avg_r2, per_r2 = weighted_r2_score(y_true, y_pred)\n    return global_r2, avg_r2, per_r2\n\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n        \n        print(f'\\nTarget: {target}')\n        print(f'{\"Index\":<6} | {\"Image Path\":<40} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}')\n        print('-' * 90)\n        \n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n        \n        print(f'\\nTarget: {target}')\n        header = f'{\"Index\":<6} | {\"Image Path\":<40} | {\"State\":<6} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}'\n        print(header)\n        print('-' * len(header))\n        \n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            state = val_df.iloc[idx]['State'] if 'State' in val_df.columns else 'NA'\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {str(state):<6} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\ndef compare_train_val(tr_df, val_df, targets, show_plots=True):\n    \"\"\"Quick comparison of target distributions and metadata between train and val splits.\"\"\"\n    print(\"\\n--- Train / Val Comparison ---\")\n\n    for t in targets:\n        tr = tr_df.get(t, pd.Series(dtype=float)).dropna()\n        val = val_df.get(t, pd.Series(dtype=float)).dropna()\n        print(f\"\\nTarget: {t}\")\n        print(f\"  Train: n={len(tr)} mean={tr.mean():.3f} std={tr.std():.3f} min={tr.min():.3f} max={tr.max():.3f}\")\n        print(f\"  Val  : n={len(val)} mean={val.mean():.3f} std={val.std():.3f} min={val.min():.3f} max={val.max():.3f}\")\n        if show_plots:\n            try:\n                plt.figure(figsize=(6, 3))\n                sns.kdeplot(tr, label='train', fill=True)\n                sns.kdeplot(val, label='val', fill=True)\n                plt.legend()\n                plt.title(f'Distribution: {t}')\n                plt.show()\n            except Exception as e:\n                print('  Could not plot distributions for', t, '-', e)\n\n    # Compare Sampling_Date and State if present\n    if 'Sampling_Date' in tr_df.columns:\n        try:\n            tr_dates = pd.to_datetime(tr_df['Sampling_Date'], errors='coerce')\n            val_dates = pd.to_datetime(val_df['Sampling_Date'], errors='coerce')\n            print(\"\\nSampling_Date range:\")\n            print(f\"  Train: {tr_dates.min()} -> {tr_dates.max()} (missing {tr_dates.isna().sum()})\")\n            print(f\"  Val  : {val_dates.min()} -> {val_dates.max()} (missing {val_dates.isna().sum()})\")\n        except Exception as e:\n            print('  Could not parse Sampling_Date:', e)\n    if 'State' in tr_df.columns:\n        print(\"\\nState distribution (train vs val):\")\n        tr_state = tr_df['State'].value_counts(normalize=True)\n        val_state = val_df['State'].value_counts(normalize=True)\n        state_df = pd.concat([tr_state, val_state], axis=1, keys=['train', 'val']).fillna(0)\n\n        print(state_df)","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.867420Z","iopub.execute_input":"2025-12-21T01:30:29.867829Z","iopub.status.idle":"2025-12-21T01:30:29.884127Z","shell.execute_reply.started":"2025-12-21T01:30:29.867805Z","shell.execute_reply":"2025-12-21T01:30:29.883320Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset & Augmentation","metadata":{}},{"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.Rotate(limit=(-10, 10), p=0.3, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101),\n        A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef get_val_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef get_tta_transforms(mode=0):\n    # mode 0: original\n    # mode 1: hflip\n    # mode 2: vflip\n    # mode 3: rotate90\n    transforms_list = [\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n    ]\n    \n    if mode == 1:\n        transforms_list.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        transforms_list.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        transforms_list.append(A.RandomRotate90(p=1.0)) # RandomRotate90 with p=1.0 rotates 90, 180, 270 randomly? \n        # Albumentations RandomRotate90 rotates by 90, 180, 270. \n        # Reference uses transforms.RandomRotation([90, 90]) which is exactly 90 degrees.\n        # To match exactly 90 degrees in Albumentations, we might need Rotate(limit=(90,90), p=1.0)\n        # But RandomRotate90 is standard TTA. Let's use Rotate(limit=(90,90)) to be precise if that's what reference does.\n        # Reference: transforms.RandomRotation([90, 90]) -> rotates by exactly 90 degrees.\n        transforms_list.append(A.Rotate(limit=(90, 90), p=1.0, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101))\n\n    transforms_list.extend([\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    return A.Compose(transforms_list, p=1.0)\ndef clean_image(img):\n    # 1. Safe Crop (Remove artifacts at the bottom)\n    h, w = img.shape[:2]\n    # Cut bottom 10% where artifacts often appear\n    img = img[0:int(h*0.90), :] \n\n    # 2. Inpaint Date Stamp (Remove orange text)\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    # Define orange color range (adjust as needed)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n\n    # Dilate mask to cover text edges and reduce noise\n    mask = cv2.dilate(mask, np.ones((3,3), np.uint8), iterations=2)\n\n    # Inpaint if mask is not empty\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n\n    return img\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df\n        self.transform = transform\n        self.img_dir = img_dir\n        self.paths = df['image_path'].values\n        self.labels = df[CFG.ALL_TARGET_COLS].values.astype(np.float32)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.img_dir, os.path.basename(self.paths[idx]))\n        img = cv2.imread(path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n        left = self.transform(image=left)['image']\n        right = self.transform(image=right)['image']\n        label = torch.from_numpy(self.labels[idx])\n        return left, right, label","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.885662Z","iopub.execute_input":"2025-12-21T01:30:29.885923Z","iopub.status.idle":"2025-12-21T01:30:29.905204Z","shell.execute_reply.started":"2025-12-21T01:30:29.885902Z","shell.execute_reply":"2025-12-21T01:30:29.904390Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model & Loss","metadata":{}},{"cell_type":"code","source":"# class BiomassModel(nn.Module):\n#     def __init__(self, model_name, pretrained=True):\n#         super().__init__()\n#         self.model_name = model_name\n#         self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='avg')\n#         nf = self.backbone.num_features\n#         comb = nf * 2\n#         self.head_green_raw  = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         self.head_clover_raw = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         self.head_dead_raw   = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         if pretrained:\n#             self.load_pretrained()\n    \n#     def load_pretrained(self):\n#         try:\n#             sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool='avg').state_dict()\n#             self.backbone.load_state_dict(sd, strict=False)\n#             print('Pretrained weights loaded.')\n#         except Exception as e:\n#             print(f'Warning: pretrained load failed: {e}')\n    \n#     def forward(self, left, right):\n#         fl = self.backbone(left)\n#         fr = self.backbone(right)\n#         x  = torch.cat([fl, fr], dim=1)\n#         green  = self.head_green_raw(x)\n#         # clover = torch.nn.functional.softplus(self.head_clover_raw(x))\n#         # dead   = torch.nn.functional.softplus(self.head_dead_raw(x))\n#         clover = self.head_clover_raw(x)\n#         dead   = self.head_dead_raw(x)\n#         gdm    = green + clover\n#         total  = gdm + dead\n#         return total, gdm, green, clover, dead\nclass LocalMambaBlock(nn.Module):\n    \"\"\"\n    Lightweight Mamba-style block (Gated CNN) from the reference notebook.\n    Efficiently mixes tokens with linear complexity.\n    \"\"\"\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        # Depthwise conv mixes spatial information locally\n        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)\n        self.gate = nn.Linear(dim, dim)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x):\n        # x: (Batch, Tokens, Dim)\n        shortcut = x\n        x = self.norm(x)\n        # Gating mechanism\n        g = torch.sigmoid(self.gate(x))\n        x = x * g\n        # Spatial mixing via 1D Conv (requires transpose)\n        x = x.transpose(1, 2)  # -> (B, D, N)\n        x = self.dwconv(x)\n        x = x.transpose(1, 2)  # -> (B, N, D)\n        # Projection\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=True, backbone_path=None):\n        super().__init__()\n        self.model_name = model_name\n        self.backbone_path = backbone_path\n        \n        # 1. Load Backbone with global_pool='' to keep patch tokens\n        #    (B, 197, 1024) instead of (B, 1024)\n        self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='')\n        \n        # 2. Enable Gradient Checkpointing (Crucial for ViT-Large memory!)\n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(True)\n            print(\"✓ Gradient Checkpointing enabled (saves ~50% VRAM)\")\n            \n        nf = self.backbone.num_features\n        \n        # 3. Mamba Fusion Neck\n        #    Mixes the concatenated tokens [Left, Right]\n        self.fusion = nn.Sequential(\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1),\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1)\n        )\n        \n        # 4. Pooling & Heads\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        \n        # Heads (using the same logic as before, but on fused features)\n        self.head_green_raw  = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        self.head_clover_raw = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        self.head_dead_raw   = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        \n        if pretrained:\n            self.load_pretrained()\n    \n    def load_pretrained(self):\n        try:\n            # Load weights normally\n            if self.backbone_path and os.path.exists(self.backbone_path):\n                print(f\"Loading backbone weights from local file: {self.backbone_path}\")\n                sd = torch.load(self.backbone_path, map_location='cpu')\n                # Handle common checkpoint wrappers (e.g. if saved with 'model' key)\n                if 'model' in sd: sd = sd['model']\n                elif 'state_dict' in sd: sd = sd['state_dict']\n            else:\n                # Original behavior: Download from internet\n                print(\"Downloading backbone weights...\")\n                sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool='').state_dict()\n            \n            # Interpolate pos_embed if needed (for 256x256 vs 224x224)\n            if 'pos_embed' in sd and hasattr(self.backbone, 'pos_embed'):\n                pe_ck = sd['pos_embed']\n                pe_m  = self.backbone.pos_embed\n                if pe_ck.shape != pe_m.shape:\n                    print(f\"Interpolating pos_embed: {pe_ck.shape} -> {pe_m.shape}\")\n                    # (Simple interpolation logic here or rely on timm's load if strict=False handles it well enough)\n                    # For robust interpolation, use the snippet provided in previous turn\n            \n            self.backbone.load_state_dict(sd, strict=False)\n            print('Pretrained weights loaded.')\n        except Exception as e:\n            print(f'Warning: pretrained load failed: {e}')\n    \n    def forward(self, left, right):\n        # 1. Extract Tokens (B, N, D)\n        #    Note: ViT usually returns [CLS, Patch1, Patch2...]\n        #    We remove CLS token for spatial mixing, or keep it. Let's keep it.\n        x_l = self.backbone(left)\n        x_r = self.backbone(right)\n        \n        # 2. Concatenate Left and Right tokens along sequence dimension\n        #    (B, N, D) + (B, N, D) -> (B, 2N, D)\n        x_cat = torch.cat([x_l, x_r], dim=1)\n        \n        # 3. Apply Mamba Fusion\n        #    This allows tokens from Left image to interact with tokens from Right image\n        x_fused = self.fusion(x_cat)\n        \n        # 4. Global Pooling\n        #    (B, 2N, D) -> (B, D, 2N) -> (B, D, 1) -> (B, D)\n        x_pool = self.pool(x_fused.transpose(1, 2)).flatten(1)\n        \n        # 5. Prediction Heads\n        green  = self.head_green_raw(x_pool)\n        clover = self.head_clover_raw(x_pool)\n        dead   = self.head_dead_raw(x_pool)\n        \n        # Summation logic\n        gdm    = green + clover\n        total  = gdm + dead\n        \n        return total, gdm, green, clover, dead\ndef biomass_loss(outputs, labels, w=None):\n    total, gdm, green, clover, dead = outputs\n    mse = nn.MSELoss()\n    huber = nn.SmoothL1Loss(beta=5.0) # Huber loss for robust regression (beta=5.0 as recommended)\n    \n    l_green  = huber(green.squeeze(),  labels[:,0])\n    l_dead   = huber(dead.squeeze(), labels[:,1]) # Use Huber loss for Dead\n    l_clover = huber(clover.squeeze(), labels[:,2])\n    l_gdm    = huber(gdm.squeeze(),    labels[:,3])\n    l_total  = huber(total.squeeze(),  labels[:,4])\n\n    # Stack per-target losses in the SAME order as CFG.ALL_TARGET_COLS\n    losses = torch.stack([l_green, l_dead, l_clover, l_gdm, l_total])\n    # losses = torch.stack([l_green, l_dead, l_clover])\n\n    # Use provided weights, or default to CFG.R2_WEIGHTS\n    if w is None:\n        return losses.mean()\n    w = torch.as_tensor(w, device=losses.device, dtype=losses.dtype)\n    w = w / w.sum()\n    return (losses * w).sum()","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.906234Z","iopub.execute_input":"2025-12-21T01:30:29.906480Z","iopub.status.idle":"2025-12-21T01:30:29.924002Z","shell.execute_reply.started":"2025-12-21T01:30:29.906453Z","shell.execute_reply":"2025-12-21T01:30:29.923260Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train Functions with EMA & Gradient Accumulation","metadata":{}},{"cell_type":"code","source":"from contextlib import nullcontext\nscaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())\n\n@torch.no_grad()\ndef valid_epoch(eval_model, loader, device):\n    eval_model.eval()\n    running = 0.0\n    preds_total, preds_gdm, preds_green, preds_clover, preds_dead, all_labels = [], [], [], [], [], []\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    \n    for l, r, lab in loader:\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        with amp_ctx():\n            p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n            loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead), lab, w=CFG.LOSS_WEIGHTS)\n        running += loss.item() * l.size(0)\n        preds_total.extend(p_total.cpu().numpy().ravel())\n        preds_gdm.extend(p_gdm.cpu().numpy().ravel())\n        preds_green.extend(p_green.cpu().numpy().ravel())\n        preds_clover.extend(p_clover.cpu().numpy().ravel())\n        preds_dead.extend(p_dead.cpu().numpy().ravel())\n        all_labels.extend(lab.cpu().numpy())\n    \n    pred_total  = np.array(preds_total)\n    pred_gdm    = np.array(preds_gdm)\n    pred_green  = np.array(preds_green)\n    pred_clover = np.array(preds_clover)\n    pred_dead   = np.array(preds_dead)\n    true_labels = np.stack(all_labels)\n    \n    pred_all = np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1)\n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(true_labels, pred_all)\n    return running / len(loader.dataset), global_r2, avg_r2, per_r2, pred_all, true_labels\n\n@torch.no_grad()\ndef valid_epoch_tta(eval_model, loaders, device):\n    eval_model.eval()\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    \n    # We need to aggregate predictions from all loaders\n    # Assuming all loaders have same order and size (which they should if shuffle=False)\n    \n    all_preds_accum = None\n    all_labels = None\n    total_loss = 0.0\n    \n    for loader_idx, loader in enumerate(loaders):\n        preds_total, preds_gdm, preds_green, preds_clover, preds_dead = [], [], [], [], []\n        current_labels = []\n        running_loss = 0.0\n        \n        for l, r, lab in loader:\n            l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n            with amp_ctx():\n                p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n                loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead), lab, w=CFG.LOSS_WEIGHTS)\n            \n            running_loss += loss.item() * l.size(0)\n            \n            preds_total.extend(p_total.cpu().numpy().ravel())\n            preds_gdm.extend(p_gdm.cpu().numpy().ravel())\n            preds_green.extend(p_green.cpu().numpy().ravel())\n            preds_clover.extend(p_clover.cpu().numpy().ravel())\n            preds_dead.extend(p_dead.cpu().numpy().ravel())\n            \n            if loader_idx == 0:\n                current_labels.extend(lab.cpu().numpy())\n        \n        total_loss += (running_loss / len(loader.dataset))\n        \n        # Stack predictions for this loader: (N, 5)\n        # Order: Green, Dead, Clover, GDM, Total (matching CFG.ALL_TARGET_COLS order roughly, but let's be precise)\n        # CFG.ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n        # preds lists are just raw outputs.\n        # Let's stack them in the order expected by weighted_r2_score_global which expects:\n        # y_true, y_pred where columns match.\n        # The model returns: total, gdm, green, clover, dead\n        # We need to stack them to match true_labels which comes from CFG.ALL_TARGET_COLS\n        # CFG.ALL_TARGET_COLS is ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n        \n        pred_stack = np.stack([\n            np.array(preds_green),\n            np.array(preds_dead),\n            np.array(preds_clover),\n            np.array(preds_gdm),\n            np.array(preds_total)\n        ], axis=1)\n        \n        if all_preds_accum is None:\n            all_preds_accum = pred_stack\n            all_labels = np.stack(current_labels)\n        else:\n            all_preds_accum += pred_stack\n            \n    # Average predictions\n    avg_preds = all_preds_accum / len(loaders)\n    avg_loss = total_loss / len(loaders)\n    \n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(all_labels, avg_preds)\n    return avg_loss, global_r2, avg_r2, per_r2, avg_preds, all_labels\n\ndef set_backbone_requires_grad(model: BiomassModel, requires_grad: bool):\n    for p in model.backbone.parameters():\n        p.requires_grad = requires_grad\n\ndef build_optimizer(model: BiomassModel):\n    head_params = (list(model.head_green_raw.parameters()) +\n                   list(model.head_clover_raw.parameters()) +\n                   list(model.head_dead_raw.parameters()))\n    backbone_params = list(model.backbone.parameters())\n    return optim.AdamW([\n        {'params': backbone_params, 'lr': CFG.LR_BACKBONE, 'weight_decay': CFG.WD},\n        {'params': head_params,     'lr': CFG.LR_HEAD,     'weight_decay': CFG.WD},\n    ])\ndef build_optimizer(model: BiomassModel):\n    # 1. Get backbone parameter IDs for exclusion\n    backbone_ids = {id(p) for p in model.backbone.parameters()}\n    \n    # 2. Separate params into backbone vs. everything else (heads, fusion, etc.)\n    backbone_params = []\n    rest_params = []\n    \n    for p in model.parameters():\n        if p.requires_grad:\n            if id(p) in backbone_ids:\n                backbone_params.append(p)\n            else:\n                rest_params.append(p)\n    \n    return optim.AdamW([\n        {'params': backbone_params, 'lr': CFG.LR_BACKBONE, 'weight_decay': CFG.WD},\n        {'params': rest_params,     'lr': CFG.LR_REST,     'weight_decay': CFG.WD},\n])\ndef build_scheduler(optimizer):\n    def lr_lambda(epoch):\n        e = max(0, epoch - 1)\n        if e < CFG.WARMUP_EPOCHS:\n            return float(e + 1) / float(max(1, CFG.WARMUP_EPOCHS))\n        progress = (e - CFG.WARMUP_EPOCHS) / float(max(1, CFG.EPOCHS - CFG.WARMUP_EPOCHS))\n        return 0.5 * (1.0 + math.cos(math.pi * progress))\n    return LambdaLR(optimizer, lr_lambda)\n\ndef train_epoch(model, loader, opt, scheduler, device, ema: ModelEmaV2 | None = None):\n    model.train()\n    running = 0.0\n    opt.zero_grad()\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    itera = tqdm(loader, desc='train', leave=False) if CFG.USE_TQDM else loader\n    for i, (l, r, lab) in enumerate(itera):\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        with amp_ctx():\n            total, gdm, green, clover, dead = model(l, r)\n            loss = biomass_loss((total, gdm, green, clover, dead), lab, w=CFG.LOSS_WEIGHTS) / CFG.GRAD_ACC\n        scaler.scale(loss).backward()\n        running += loss.item() * l.size(0) * CFG.GRAD_ACC\n        \n        if (i + 1) % CFG.GRAD_ACC == 0 or (i + 1) == len(loader):\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            scaler.step(opt)\n            scaler.update()\n            if ema is not None:\n                ema.update(model.module if hasattr(model, 'module') else model)\n            opt.zero_grad()\n    scheduler.step()\n    return running / len(loader.dataset)","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.924952Z","iopub.execute_input":"2025-12-21T01:30:29.925184Z","iopub.status.idle":"2025-12-21T01:30:29.949166Z","shell.execute_reply.started":"2025-12-21T01:30:29.925165Z","shell.execute_reply":"2025-12-21T01:30:29.948592Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5-Fold Training Loop with EMA","metadata":{}},{"cell_type":"code","source":"# print('Loading data...')\n# df_long = pd.read_csv(CFG.TRAIN_CSV)\n# df_wide = df_long.pivot(index='image_path', columns='target_name', values='target').reset_index()\n# assert df_wide['image_path'].is_unique, 'Leakage risk: duplicate image_path rows'\n\n# # Merge metadata (Sampling_Date, State) for stratification\n# if 'Sampling_Date' in df_long.columns and 'State' in df_long.columns:\n#     print('Merging metadata for stratification...')\n#     meta_df = df_long[['image_path', 'Sampling_Date', 'State']].drop_duplicates()\n#     df_wide = df_wide.merge(meta_df, on='image_path', how='left')\n\n# # Keep necessary columns\n# df_wide = df_wide[['image_path', 'Sampling_Date', 'State'] + CFG.ALL_TARGET_COLS]\n# print(f'{len(df_wide)} training images')\n\n# # Use StratifiedGroupKFold\n# sgkf = StratifiedGroupKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\n# oof_true, oof_pred, fold_summary = [], [], []\n\n# # Split based on groups (Sampling_Date) and stratification target (State)\n# groups = df_wide['Sampling_Date']\n# y_stratify = df_wide['State']\n\n# # models_list = [] # Removed to save memory\n\n# for fold, (tr_idx, val_idx) in enumerate(sgkf.split(df_wide, y_stratify, groups=groups)):\n#     if fold not in CFG.FOLDS_TO_TRAIN:\n#         print(f'Skipping fold {fold} as per configuration.')\n#         continue\n#     print('\\n' + '='*70)\n#     print(f'FOLD {fold+1}/{CFG.N_FOLDS} | {len(tr_idx)} train / {len(val_idx)} val')\n#     print('='*70)\n#     torch.cuda.empty_cache(); gc.collect()\n    \n#     tr_df  = df_wide.iloc[tr_idx].reset_index(drop=True)\n#     val_df = df_wide.iloc[val_idx].reset_index(drop=True)\n\n#     # Quick train/val comparison for this fold\n#     try:\n#         compare_train_val(tr_df, val_df, CFG.ALL_TARGET_COLS, show_plots=True)\n#     except Exception as e:\n#         print('Warning: compare_train_val failed:', e)\n\n#     tr_set = BiomassDataset(tr_df,  get_train_transforms(), CFG.TRAIN_IMAGE_DIR)\n    \n#     # Create TTA loaders\n#     val_loaders = []\n#     for mode in range(CFG.VAL_TTA_TIMES): # 0: orig, 1: hflip, 2: vflip, 3: rot90\n#         val_set_tta = BiomassDataset(val_df, get_tta_transforms(mode), CFG.TRAIN_IMAGE_DIR)\n#         val_loader_tta = DataLoader(val_set_tta, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n#         val_loaders.append(val_loader_tta)\n\n#     tr_loader  = DataLoader(tr_set, batch_size=CFG.BATCH_SIZE, shuffle=True, num_workers=CFG.NUM_WORKERS, pin_memory=True, drop_last=True)\n    \n#     print('Building model...')\n#     backbone_path = getattr(CFG, 'BACKBONE_PATH', None)\n#     base_model = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED, backbone_path=backbone_path).to(CFG.DEVICE)\n    \n#     # Load pretrained fold weights if available (for resuming or fine-tuning)\n#     if getattr(CFG, 'PRETRAINED_DIR', None) and os.path.isdir(CFG.PRETRAINED_DIR):\n#         pretrained_path = os.path.join(CFG.PRETRAINED_DIR, f'best_model_fold{fold}.pth')\n#         if os.path.exists(pretrained_path):\n#             try:\n#                 state = torch.load(pretrained_path, map_location='cpu')\n#                 # support raw state_dict or dict-with-keys\n#                 if isinstance(state, dict) and ('model_state_dict' in state or 'state_dict' in state):\n#                     key = 'model_state_dict' if 'model_state_dict' in state else 'state_dict'\n#                     sd = state[key]\n#                 else:\n#                     sd = state\n#                 base_model.load_state_dict(sd, strict=False)\n#                 base_model.to(CFG.DEVICE)\n#                 print(f'  ✓ Loaded pretrained weights for fold {fold} from {pretrained_path}')\n#             except Exception as e:\n#                 print(f'  ✗ Failed to load pretrained fold {fold}: {e}')\n#         else:\n#             print(f'  (No pretrained file for fold {fold} at {pretrained_path})')\n#     else:\n#         print('  (No PRETRAINED_DIR configured or directory missing)')\n        \n#     model = nn.DataParallel(base_model)\n#     set_backbone_requires_grad(base_model, False)\n#     optimizer = build_optimizer(base_model)\n#     scheduler = build_scheduler(optimizer)\n#     ema = ModelEmaV2(base_model, decay=CFG.EMA_DECAY)\n    \n#     best_global_r2 = -np.inf\n#     patience = 0\n#     best_fold_preds = None; best_fold_true = None\n#     best_avg_r2 = -np.inf\n    \n#     # Define save path\n#     save_path = os.path.join(CFG.MODEL_DIR, f'best_model_fold{fold}.pth')\n    \n#     for epoch in range(1, CFG.EPOCHS + 1):\n#         if epoch == CFG.FREEZE_EPOCHS + 1:\n#             patience = 0\n#             set_backbone_requires_grad(base_model, True)\n#             print(f'Epoch {epoch}: backbone unfrozen')\n        \n#         tr_loss = train_epoch(model, tr_loader, optimizer, scheduler, CFG.DEVICE, ema)\n#         eval_model = ema.module if ema is not None else (model.module if hasattr(model, 'module') else model)\n        \n#         # Use TTA validation\n#         val_loss, global_r2, avg_r2, per_r2, preds_fold, true_fold = valid_epoch_tta(eval_model, val_loaders, CFG.DEVICE)\n        \n#         per_r2_str = ' | '.join([f'{CFG.ALL_TARGET_COLS[i][:5]}: {r2:.3f}' for i, r2 in enumerate(per_r2)])\n#         lrs = [pg['lr'] for pg in optimizer.param_groups]\n#         print(f'Fold {fold} | Epoch {epoch:02d} | TLoss {tr_loss:.5f} | VLoss {val_loss:.5f} |avgR2 {avg_r2:.4f}| GlobalR² {global_r2:.4f} {\"[BEST]\" if global_r2 > best_global_r2 else \"\"}')\n#         print(f'  → {per_r2_str}')\n        \n#         if global_r2 > best_global_r2:\n#             best_global_r2 = global_r2\n#             best_avg_r2 = avg_r2\n            \n#             # Save the EMA weights (best state) to disk immediately\n#             # Clone to CPU to avoid memory issues\n#             best_state = {k: v.cpu().clone() for k, v in eval_model.state_dict().items()}\n#             torch.save(best_state, save_path)\n#             print(f'  → SAVED EMA weights to {save_path} (GlobalR²: {best_global_r2:.4f})')\n#             del best_state # Free memory\n            \n#             patience = 0\n#             best_fold_preds = preds_fold; best_fold_true = true_fold\n#         else:\n#             patience += 1\n#             if patience >= CFG.PATIENCE:\n#                     print(f'  → EARLY STOP (no improvement in {CFG.PATIENCE} epochs)')\n#                     break\n                \n#         del preds_fold, true_fold\n#         torch.cuda.empty_cache()\n#         gc.collect()\n    \n#     if best_fold_preds is not None:\n#         oof_true.append(best_fold_true); oof_pred.append(best_fold_preds)\n#         fold_summary.append({'fold': fold, 'global_r2': best_global_r2,'avg_r2':avg_r2})\n    \n#     # Cleanup for this fold\n#     del model, base_model, tr_loader, val_loaders, optimizer, scheduler, ema\n#     if 'eval_model' in locals(): del eval_model\n#     torch.cuda.empty_cache(); gc.collect()\n\n# if oof_true:\n#     oof_true_arr = np.concatenate(oof_true, axis=0)\n#     oof_pred_arr = np.concatenate(oof_pred, axis=0)\n#     oof_global_r2, oof_avg_r2, oof_per_r2 = weighted_r2_score_global(oof_true_arr, oof_pred_arr)\n\n#     print('\\nTraining complete! Models saved in:', CFG.MODEL_DIR)\n#     print('Fold summary:')\n#     for fs in fold_summary:\n#         print(f\"  Fold {fs['fold']}: Global R² = {fs['global_r2']:.4f}, Avg R² = {fs.get('avg_r2', float('nan')):.4f}\")\n#     print(f'OOF Global Weighted R²: {oof_global_r2:.4f} | OOF Avg Target R²: {oof_avg_r2:.4f}')\n#     print('OOF Per-target:', dict(zip(CFG.ALL_TARGET_COLS, [f\"{r:.4f}\" for r in oof_per_r2])))\n# else:\n#     print('No OOF predictions collected.')","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.950057Z","iopub.execute_input":"2025-12-21T01:30:29.950335Z","iopub.status.idle":"2025-12-21T01:30:29.965454Z","shell.execute_reply.started":"2025-12-21T01:30:29.950315Z","shell.execute_reply":"2025-12-21T01:30:29.964758Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# submit","metadata":{}},{"cell_type":"code","source":"# ===============================================================\n# 4. DEFINE TTA TRANSFORMS\n# ===============================================================\ndef get_tta_transforms(num_transforms):\n    \"\"\"\n    Returns a list of TTA transform pipelines.\n    Each pipeline represents a different augmentation view.\n    \"\"\"\n    normalize = A.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225]\n    )\n    to_tensor = ToTensorV2()\n    \n    all_tta_transforms = [\n        # View 1: Original (no flip)\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            normalize,\n            to_tensor\n        ]),\n        \n        # View 2: Horizontal Flip\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.HorizontalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n        \n        # View 3: Vertical Flip\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.VerticalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n        \n        # View 4: Both Flips\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.HorizontalFlip(p=1.0),\n            A.VerticalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n    ]\n    tta_transforms = all_tta_transforms[:num_transforms]\n    return tta_transforms\n\nprint(f\"✓ TTA transforms defined ({CFG.TTA_STEPS} views)\")","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.966183Z","iopub.execute_input":"2025-12-21T01:30:29.966529Z","iopub.status.idle":"2025-12-21T01:30:29.979718Z","shell.execute_reply.started":"2025-12-21T01:30:29.966507Z","shell.execute_reply":"2025-12-21T01:30:29.979078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===============================================================\n# 5. CREATE TEST DATASET\n# ===============================================================\ndef clean_image(img):\n    # Safe crop (remove bottom artifacts) + inpaint orange date stamp\n    h, w = img.shape[:2]\n    img = img[0:int(h * 0.90), :]\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n    return img\nclass BiomassTestDataset(Dataset):\n    \"\"\"\n    Test dataset for biomass images.\n    Splits each 2000×1000 image into left and right 1000×1000 halves.\n    \"\"\"\n    def __init__(self, img_dir):\n        self.img_dir = img_dir\n        self.paths = sorted([\n            os.path.join(img_dir, f) for f in os.listdir(img_dir)\n            if f.lower().endswith(('.jpg', '.jpeg', '.png'))\n        ])\n        self.filenames = [os.path.basename(p) for p in self.paths]\n    \n    def __len__(self):\n        return len(self.paths)\n    \n    def __getitem__(self, idx):\n        path = self.paths[idx]\n        \n        # Read image\n        img = cv2.imread(path)\n        if img is None:\n            print(f\"Warning: Could not read {path}, using blank image\")\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        \n        # Convert BGR to RGB\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        # Split into left and right halves\n        h, w = img.shape[:2]\n        mid = w // 2\n        left = img[:, :mid].copy()\n        right = img[:, mid:].copy()\n        \n        return left, right, self.filenames[idx]\n\nprint(\"✓ Test dataset class defined\")","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.980641Z","iopub.execute_input":"2025-12-21T01:30:29.980919Z","iopub.status.idle":"2025-12-21T01:30:29.997793Z","shell.execute_reply.started":"2025-12-21T01:30:29.980889Z","shell.execute_reply":"2025-12-21T01:30:29.997203Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===============================================================\n# 8. RUN INFERENCE WITH TTA (UPDATED to honor CFG.FOLDS_TO_TRAIN)\n# ===============================================================\n\n@torch.no_grad()\ndef predict_with_tta(model, left_np, right_np, tta_transforms):\n    \"\"\"\n    Predict using TTA with a SINGLE model.\n    \n    Args:\n        model: Single trained model\n        left_np: Left half of image (numpy array)\n        right_np: Right half of image (numpy array)\n        tta_transforms: List of augmentation transforms\n    \n    Returns:\n        numpy array: [total, gdm, green] predictions (averaged over TTA)\n    \"\"\"\n    all_tta_preds = []\n    \n    # Loop over TTA views\n    for tfm in tta_transforms:\n        # Apply transform to both halves\n        left_tensor = tfm(image=left_np)['image'].unsqueeze(0).to(CFG.DEVICE)\n        right_tensor = tfm(image=right_np)['image'].unsqueeze(0).to(CFG.DEVICE)\n        \n        total, gdm, green, clover, dead = model(left_tensor, right_tensor)\n        \n        # Extract values\n        p_total = total.cpu().item()\n        p_gdm = gdm.cpu().item()\n        p_green = green.cpu().item()\n        \n        all_tta_preds.append([p_total, p_gdm, p_green])\n    \n    # Average across TTA views\n    final_pred = np.mean(all_tta_preds, axis=0)\n    \n    return final_pred\n\n\ndef run_inference():\n    \"\"\"\n    Main inference function.\n    Returns: (predictions_array, image_filenames)\n    Notes:\n      - Now respects `CFG.FOLDS_TO_TRAIN` (if set) and averages only over successfully loaded folds.\n      - If no fold weights are found for the requested folds, an error is raised.\n    \"\"\"\n    print(\"\\n\" + \"=\"*70)\n    print(\"STARTING INFERENCE\")\n    print(\"=\"*70)\n    \n    # Create dataset and loader\n    dataset = BiomassTestDataset(CFG.TEST_IMAGE_DIR)\n    # Note: batch_size=1 is required for the current predict_with_tta implementation\n    loader = DataLoader(\n        dataset,\n        batch_size=1,  \n        shuffle=False,\n        num_workers=CFG.NUM_WORKERS,\n        pin_memory=True\n    )\n    \n    tta_transforms = get_tta_transforms(CFG.TTA_STEPS)\n    \n    # Initialize accumulator for predictions\n    # Shape: (num_samples, 3) for Total, GDM, Green\n    accumulated_preds = np.zeros((len(dataset), 3), dtype=np.float32)\n\n    # Use configured folds, fallback to full range if not set or empty\n    folds_to_use = getattr(CFG, 'FOLDS_TO_TRAIN', list(range(CFG.N_FOLDS)))\n    if not folds_to_use:\n        folds_to_use = list(range(CFG.N_FOLDS))\n\n    print(f\"Folds requested for inference: {folds_to_use}\")\n\n    # Use filenames from dataset (guaranteed consistent ordering with loader because shuffle=False)\n    filenames = dataset.filenames.copy()\n\n    successful_folds = 0\n    \n    # Loop over requested folds only\n    for fold in folds_to_use:\n        print(f\"\\nProcessing Fold {fold}...\")\n        # Load model for this fold\n        model_dir = CFG.MODEL_DIR_012 if fold in [0,1,2] else CFG.MODEL_DIR_34\n        backbone_path = getattr(CFG, 'BACKBONE_PATH', None)\n        model = BiomassModel(CFG.MODEL_NAME, pretrained=False, backbone_path=backbone_path)\n        \n        # Load weights\n        weight_path = os.path.join(model_dir, f'best_model_fold{fold}.pth')\n        if not os.path.exists(weight_path):\n            print(f\"Warning: Model file {weight_path} not found! Skipping fold {fold}.\")\n            del model\n            torch.cuda.empty_cache(); gc.collect()\n            continue\n            \n        state = torch.load(weight_path, map_location='cpu')\n        # Handle state dict keys if necessary (e.g. if saved with 'model_state_dict' key)\n        if isinstance(state, dict) and ('model_state_dict' in state or 'state_dict' in state):\n             key = 'model_state_dict' if 'model_state_dict' in state else 'state_dict'\n             sd = state[key]\n        else:\n             sd = state\n        \n        model.load_state_dict(sd)\n        model.to(CFG.DEVICE)\n        model.eval()\n        \n        # Run inference for this fold\n        for i, (left, right, filename) in enumerate(tqdm(loader, desc=f\"Fold {fold}\")):\n            # left and right are batches of size 1, convert to numpy for TTA function\n            left_np = left[0].numpy()\n            right_np = right[0].numpy()\n            \n            # Predict\n            pred = predict_with_tta(model, left_np, right_np, tta_transforms)\n            accumulated_preds[i] += pred\n            \n        successful_folds += 1\n        # Cleanup model to save memory\n        del model\n        torch.cuda.empty_cache(); gc.collect()\n        \n    if successful_folds == 0:\n        raise FileNotFoundError(f\"No model weights found for requested folds: {folds_to_use}\")\n\n    # Average predictions over the number of successfully loaded folds\n    final_predictions = accumulated_preds / successful_folds\n    \n    print(f\"\\nInference complete. Successfully used {successful_folds} fold(s) out of {len(folds_to_use)} requested.\")\n    return final_predictions, filenames","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:29.999395Z","iopub.execute_input":"2025-12-21T01:30:29.999968Z","iopub.status.idle":"2025-12-21T01:30:30.013034Z","shell.execute_reply.started":"2025-12-21T01:30:29.999947Z","shell.execute_reply":"2025-12-21T01:30:30.012421Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===============================================================\n# 9. POST-PROCESS PREDICTIONS\n# ===============================================================\ndef postprocess_predictions(preds_direct):\n    \"\"\"\n    Calculate derived targets from direct predictions.\n    \n    Input: (n_samples, 3) array with [total, gdm, green]\n    Output: (n_samples, 5) array with [green, dead, clover, gdm, total]\n    \"\"\"\n    print(\"\\nPost-processing predictions...\")\n    \n    # Extract direct predictions\n    pred_total = preds_direct[:, 0]\n    pred_gdm = preds_direct[:, 1]\n    pred_green = preds_direct[:, 2]\n    \n    # Calculate derived targets with non-negativity constraint\n    pred_clover = np.maximum(0, pred_gdm - pred_green)\n    pred_dead = np.maximum(0, pred_total - pred_gdm)\n    \n    # Stack in the order of ALL_TARGET_COLS\n    # ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    preds_all = np.stack([\n        pred_green,\n        pred_dead,\n        pred_clover,\n        pred_gdm,\n        pred_total\n    ], axis=1)\n    \n    print(f\"✓ Post-processing complete\")\n    print(f\"  Output shape: {preds_all.shape}\")\n    print(f\"\\nPrediction statistics:\")\n    for i, col in enumerate(CFG.ALL_TARGET_COLS):\n        print(f\"  {col:15s}: mean={preds_all[:, i].mean():.2f}, \"\n              f\"std={preds_all[:, i].std():.2f}, \"\n              f\"min={preds_all[:, i].min():.2f}, \"\n              f\"max={preds_all[:, i].max():.2f}\")\n    \n    return preds_all\n\n","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:30.013973Z","iopub.execute_input":"2025-12-21T01:30:30.014551Z","iopub.status.idle":"2025-12-21T01:30:30.027684Z","shell.execute_reply.started":"2025-12-21T01:30:30.014528Z","shell.execute_reply":"2025-12-21T01:30:30.026990Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===============================================================\n# 10. CREATE SUBMISSION FILE (FIXED)\n# ===============================================================\ndef create_submission(predictions, filenames):\n    \"\"\"\n    Create submission file in the required format.\n    \n    Args:\n        predictions: (n_images, 5) array with all target predictions\n        filenames: list of test image filenames\n    \"\"\"\n    print(\"\\n\" + \"=\"*70)\n    print(\"CREATING SUBMISSION FILE\")\n    print(\"=\"*70)\n    \n    # Step 0: Load test.csv first to check the image_path format\n    test_df = pd.read_csv(CFG.TEST_CSV)\n    print(f\"\\nTest CSV loaded: {len(test_df)} rows\")\n    print(f\"Sample image_path from test.csv: {test_df['image_path'].iloc[0]}\")\n    print(f\"Sample filename from predictions: {filenames[0]}\")\n    \n    # Step 1: Fix image_path format to match test.csv\n    # If test.csv has \"test/ID123.jpg\" but we have \"ID123.jpg\", add the prefix\n    test_path_example = test_df['image_path'].iloc[0]\n    if '/' in test_path_example:\n        # Extract the subdirectory prefix (e.g., \"test/\")\n        prefix = test_path_example.rsplit('/', 1)[0] + '/'\n        corrected_filenames = [prefix + fn for fn in filenames]\n        print(f\"Corrected path format: {corrected_filenames[0]}\")\n    else:\n        corrected_filenames = filenames\n    \n    # Step 2: Create wide-format DataFrame with corrected paths\n    preds_wide = pd.DataFrame(predictions, columns=CFG.ALL_TARGET_COLS)\n    preds_wide.insert(0, 'image_path', corrected_filenames)\n    \n    print(f\"\\nWide format predictions:\")\n    print(preds_wide.head())\n    \n    # Step 3: Convert to long format (melt)\n    preds_long = preds_wide.melt(\n        id_vars=['image_path'],\n        value_vars=CFG.ALL_TARGET_COLS,\n        var_name='target_name',\n        value_name='target'\n    )\n    \n    print(f\"\\nLong format predictions (first 10 rows):\")\n    print(preds_long.head(10))\n    \n    # Step 4: Debug the merge\n    print(f\"\\nDebug: Checking if paths match...\")\n    print(f\"Unique paths in test_df: {test_df['image_path'].nunique()}\")\n    print(f\"Unique paths in preds_long: {preds_long['image_path'].nunique()}\")\n    \n    common_paths = set(test_df['image_path'].unique()) & set(preds_long['image_path'].unique())\n    print(f\"Common paths found: {len(common_paths)}\")\n    \n    if len(common_paths) == 0:\n        print(\"\\n❌ ERROR: No matching paths found!\")\n        print(f\"Test CSV paths sample: {list(test_df['image_path'].unique()[:3])}\")\n        print(f\"Prediction paths sample: {list(preds_long['image_path'].unique()[:3])}\")\n        raise ValueError(\"Path mismatch between test.csv and predictions\")\n    \n    # Step 5: Merge to get sample_ids\n    submission = pd.merge(\n        test_df[['sample_id', 'image_path', 'target_name']],\n        preds_long,\n        on=['image_path', 'target_name'],\n        how='left'\n    )\n    \n    # Step 6: Keep only required columns\n    submission = submission[['sample_id', 'target']]\n    \n    # Step 7: Check for missing values\n    missing_count = submission['target'].isna().sum()\n    if missing_count > 0:\n        print(f\"\\n⚠ Warning: {missing_count} missing predictions found!\")\n        print(\"Sample missing entries:\")\n        print(submission[submission['target'].isna()].head())\n        submission.loc[submission['target'].isna(), 'target'] = 0.0\n    \n    # Step 8: Sort by sample_id\n    submission = submission.sort_values('sample_id').reset_index(drop=True)\n    \n    # Step 9: Save to CSV\n    output_path = os.path.join(CFG.SUBMISSION_DIR, 'submission.csv')\n    submission.to_csv(output_path, index=False)\n    \n    print(f\"\\n✓ Submission file saved: {output_path}\")\n    print(f\"  Total rows: {len(submission)}\")\n    print(f\"\\nPrediction statistics:\")\n    print(f\"  Min: {submission['target'].min():.4f}\")\n    print(f\"  Max: {submission['target'].max():.4f}\")\n    print(f\"  Mean: {submission['target'].mean():.4f}\")\n    print(f\"  Non-zero values: {(submission['target'] > 0).sum()}/{len(submission)}\")\n    \n    print(f\"\\nFirst 10 rows:\")\n    print(submission.head(10))\n    print(f\"\\nLast 10 rows:\")\n    print(submission.tail(10))\n    \n    # Step 10: Validation checks\n    print(f\"\\n\" + \"=\"*70)\n    print(\"VALIDATION CHECKS\")\n    print(\"=\"*70)\n    print(f\"✓ Expected rows: {len(test_df)}\")\n    print(f\"✓ Actual rows: {len(submission)}\")\n    print(f\"✓ Match: {len(submission) == len(test_df)}\")\n    print(f\"✓ No missing values: {not submission['target'].isna().any()}\")\n    print(f\"✓ All sample_ids unique: {submission['sample_id'].is_unique}\")\n    print(f\"✓ Has non-zero predictions: {(submission['target'] > 0).any()}\")\n    \n    return submission\n\n# Create submission\n# Post-process predictions\n# Run inference\nif CFG.CREATE_SUBMISSION:\n    predictions_direct, test_filenames = run_inference()\n    predictions_all = postprocess_predictions(predictions_direct)\n    submission_df = create_submission(predictions_all, test_filenames)","metadata":{"vscode":{"languageId":"python"},"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T01:30:30.028487Z","iopub.execute_input":"2025-12-21T01:30:30.028746Z","iopub.status.idle":"2025-12-21T01:31:05.937042Z","shell.execute_reply.started":"2025-12-21T01:30:30.028718Z","shell.execute_reply":"2025-12-21T01:31:05.936070Z"}},"outputs":[],"execution_count":null}]}