{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":663314,"sourceType":"modelInstanceVersion","modelInstanceId":471723,"modelId":487624}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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\nfor 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":"2026-01-03T06:32:30.283881Z","iopub.execute_input":"2026-01-03T06:32:30.284267Z","iopub.status.idle":"2026-01-03T06:32:31.105092Z","shell.execute_reply.started":"2026-01-03T06:32:30.284226Z","shell.execute_reply":"2026-01-03T06:32:31.104208Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport math\nimport gc\nimport random\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import List, Tuple, Dict, Optional\n\nimport cv2\nimport timm\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor\n\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.preprocessing import PowerTransformer, LabelEncoder, StandardScaler\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nwarnings.filterwarnings('ignore')\n\n# Check GPU availability\nprint(f\"CUDA Available: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:31.106302Z","iopub.execute_input":"2026-01-03T06:32:31.106687Z","iopub.status.idle":"2026-01-03T06:32:39.337192Z","shell.execute_reply.started":"2026-01-03T06:32:31.106659Z","shell.execute_reply":"2026-01-03T06:32:39.336329Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Config:\n    # Paths (Kaggle-specific)\n    DATA_DIR: str = '/kaggle/input/csiro-biomass'\n    TRAIN_CSV: str = '/kaggle/input/csiro-biomass/train.csv'\n    TEST_CSV: str = '/kaggle/input/csiro-biomass/test.csv'\n    \n    # ADD THIS: Model checkpoint path for inference\n    MODEL_DIR: str = '/kaggle/input/csiro/pytorch/default/12'\n    CKPT_PATTERN: str = '/kaggle/input/csiro/pytorch/default/12/fold_{fold}/checkpoints/best_wr2.pt'\n    \n    \n    DINO_CANDIDATES: Tuple[str, ...] = (\n        \"vit_base_patch14_dinov2\",\n        \"vit_base_patch14_reg4_dinov2\",\n        \"vit_small_patch14_dinov2\",\n    )\n    \n    # Training\n    SEED: int = 42\n    N_FOLDS: int = 5\n    EPOCHS: int = 50\n    BATCH_SIZE: int = 8  # Reduced for DINOv2 (larger model)\n    ACCUMULATE_GRAD_BATCHES: int = 4  # Effective batch size = 32\n    \n    # Image - DINOv2 native resolution\n    IMG_SIZE: int = 518  # DINOv2 native size (patch14 * 37)\n    \n    # Model architecture params (from competitor)\n    DROPOUT: float = 0.1\n    HIDDEN_RATIO: float = 0.35\n    \n    # Tiling configuration\n    SMALL_GRID: Tuple[int, int] = (4, 4)\n    BIG_GRID: Tuple[int, int] = (2, 2)\n    \n    # Transformer configuration\n    T2T_DEPTH: int = 2\n    CROSS_LAYERS: int = 2\n    CROSS_HEADS: int = 6\n    \n    # Pyramid mixer configuration\n    PYRAMID_DIMS: Tuple[int, int, int] = (384, 512, 640)\n    MOBILEVIT_HEADS: int = 4\n    MOBILEVIT_DEPTH: int = 2\n    SRA_HEADS: int = 8\n    SRA_RATIO: int = 2\n    MAMBA_DEPTH: int = 3\n    MAMBA_KERNEL: int = 5\n    \n    # Auxiliary head\n    AUX_HEAD: bool = True\n    AUX_LOSS_WEIGHT: float = 0.4\n    \n    # Optimizer\n    LR: float = 5e-5  # Lower LR for pretrained ViT\n    WEIGHT_DECAY: float = 1e-5\n    \n    # Target columns (order matters for physical constraints)\n    TARGET_COLS: Tuple[str, ...] = (\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"Dry_Total_g\", \"GDM_g\")\n    TARGET_WEIGHTS: np.ndarray = None  # Will be set in __post_init__\n    \n    # Feature columns (YOUR improvement - tabular features)\n    FEATURE_COLS: Tuple[str, ...] = ('Pre_GSHH_NDVI', 'Height_Ave_cm')\n    CAT_COLS: Tuple[str, ...] = ('State', 'Species')\n    \n    # TTA\n    USE_TTA: bool = True\n    TTA_STEPS: int = 4\n    \n    # DataLoader\n    NUM_WORKERS: int = 2\n    MIXED_PRECISION: bool = True\n    \n    # Post-processing (from competitor)\n    SMOOTH_FACTOR: float = 0.905\n    DRY_CLOVER_MIN: float = 1.25\n    DRY_DEAD_MIN: float = 1.00\n    \n    def __post_init__(self):\n        self.TARGET_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.5, 0.2])\n    \n    @property\n    def device(self):\n        return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nCFG = Config()\nCFG.__post_init__()\n\nprint(f\"\\n✅ Configuration loaded\")\nprint(f\"Model: DINOv2 with Tiled Encoding + Tabular Features\")\nprint(f\"Image Size: {CFG.IMG_SIZE}\")\nprint(f\"Batch Size: {CFG.BATCH_SIZE} x {CFG.ACCUMULATE_GRAD_BATCHES} = {CFG.BATCH_SIZE * CFG.ACCUMULATE_GRAD_BATCHES}\")\nprint(f\"Learning Rate: {CFG.LR}\")\nprint(f\"\\n🎯 Combined Approach:\")\nprint(f\"  ✓ DINOv2 backbone (competitor)\")\nprint(f\"  ✓ Left/Right image split (competitor)\")\nprint(f\"  ✓ Physical constraint prediction (competitor)\")\nprint(f\"  ✓ Tabular features: NDVI, Height, etc. (YOUR improvement)\")\nprint(f\"  ✓ PowerTransformer normalization (YOUR improvement)\")\nprint(f\"  ✓ Stratified K-Fold (YOUR improvement)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.339668Z","iopub.execute_input":"2026-01-03T06:32:39.340223Z","iopub.status.idle":"2026-01-03T06:32:39.351606Z","shell.execute_reply.started":"2026-01-03T06:32:39.340196Z","shell.execute_reply":"2026-01-03T06:32:39.350906Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed: int = 42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    pl.seed_everything(seed, workers=True)\n\nseed_everything(CFG.SEED)\nprint(f\"✅ Seed set to {CFG.SEED}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.352649Z","iopub.execute_input":"2026-01-03T06:32:39.353100Z","iopub.status.idle":"2026-01-03T06:32:39.376837Z","shell.execute_reply.started":"2026-01-03T06:32:39.353040Z","shell.execute_reply":"2026-01-03T06:32:39.375924Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray, weights=None):\n    \"\"\"Calculate weighted R² score - y_true, y_pred: shape (N, 5)\"\"\"\n    if weights is None:\n        weights = CFG.TARGET_WEIGHTS\n    \n    r2_scores = []\n    for i in range(5):\n        y_t = y_true[:, i]\n        y_p = y_pred[:, i]\n        ss_res = np.sum((y_t - y_p) ** 2)\n        ss_tot = np.sum((y_t - np.mean(y_t)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n    \n    r2_scores = np.array(r2_scores)\n    weighted_r2 = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted_r2, r2_scores","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.377859Z","iopub.execute_input":"2026-01-03T06:32:39.378171Z","iopub.status.idle":"2026-01-03T06:32:39.391806Z","shell.execute_reply.started":"2026-01-03T06:32:39.378145Z","shell.execute_reply":"2026-01-03T06:32:39.391132Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"                #Data Preparation (YOUR improvement - tabular features)\ndef prepare_data():\n    \"\"\"Load and prepare training data with YOUR feature engineering\"\"\"\n    \n    # Load data\n    train_df = pd.read_csv(CFG.TRAIN_CSV)\n    \n    # Pivot to wide format\n    train_wide = pd.pivot_table(\n        train_df,\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    # YOUR Feature engineering\n    train_wide['NDVI_Height_interaction'] = train_wide['Pre_GSHH_NDVI'] * train_wide['Height_Ave_cm']\n    train_wide['NDVI_squared'] = train_wide['Pre_GSHH_NDVI'] ** 2\n    train_wide['Height_log'] = np.log1p(train_wide['Height_Ave_cm'])\n    \n    # Extract date features\n    train_wide['Sampling_Date'] = pd.to_datetime(train_wide['Sampling_Date'])\n    train_wide['Month'] = train_wide['Sampling_Date'].dt.month\n    train_wide['Season'] = train_wide['Month'].map({\n        12: 0, 1: 0, 2: 0,  # Summer\n        3: 1, 4: 1, 5: 1,   # Autumn\n        6: 2, 7: 2, 8: 2,   # Winter\n        9: 3, 10: 3, 11: 3  # Spring\n    })\n    \n    state_encoder = LabelEncoder()\n    species_encoder = LabelEncoder()\n    \n    train_wide['State_encoded'] = state_encoder.fit_transform(train_wide['State'])\n    train_wide['Species_encoded'] = species_encoder.fit_transform(train_wide['Species'])\n    \n    target_scaler = PowerTransformer(method='yeo-johnson')\n    target_cols_list = list(CFG.TARGET_COLS)\n    train_wide[target_cols_list] = target_scaler.fit_transform(train_wide[target_cols_list].values)\n    \n    feature_cols = ['Pre_GSHH_NDVI', 'Height_Ave_cm', 'NDVI_Height_interaction', \n                   'NDVI_squared', 'Height_log', 'Month']\n    feature_scaler = StandardScaler()\n    train_wide[feature_cols] = feature_scaler.fit_transform(train_wide[feature_cols])\n    \n    numeric_cols = feature_cols + ['State_encoded', 'Species_encoded', 'Season'] + target_cols_list\n    for col in numeric_cols:\n        train_wide[col] = train_wide[col].astype(np.float32)\n    \n    train_wide['Dry_Total_bin'] = pd.qcut(train_wide['Dry_Total_g'], q=10, labels=False, duplicates='drop')\n    \n    return train_wide, target_scaler, feature_scaler, state_encoder, species_encoder\n\ntrain_df, target_scaler, feature_scaler, state_encoder, species_encoder = prepare_data()\nprint(f\"✅ Training samples: {len(train_df)}\")\nprint(f\"\\n📊 Feature columns: {train_df.columns.tolist()}\")\nprint(f\"\\n📈 Target stats (normalized):\")\nprint(train_df[list(CFG.TARGET_COLS)].describe())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.392636Z","iopub.execute_input":"2026-01-03T06:32:39.392852Z","iopub.status.idle":"2026-01-03T06:32:39.478904Z","shell.execute_reply.started":"2026-01-03T06:32:39.392831Z","shell.execute_reply":"2026-01-03T06:32:39.478111Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"                 #Dataset with Left/Right Split + Tabular Features\nclass BiomassDatasetDINO(Dataset):\n    \"\"\"\n    Combined dataset:\n    - Left/Right image split (competitor's approach)\n    - Tabular features (YOUR improvement)\n    \"\"\"\n    def __init__(self, df, transform=None, is_test=False):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.is_test = is_test\n        \n        self.feature_cols = ['Pre_GSHH_NDVI', 'Height_Ave_cm', 'NDVI_Height_interaction',\n                            'NDVI_squared', 'Height_log', 'State_encoded', 'Species_encoded', \n                            'Month', 'Season']\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        \n        img_path = os.path.join(CFG.DATA_DIR, row['image_path'])\n        img = cv2.imread(img_path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        h, w, _ = img.shape\n        mid = w // 2\n        left_img = img[:, :mid]\n        right_img = img[:, mid:]\n        \n        # Apply transforms\n        if self.transform:\n            left_tensor = self.transform(image=left_img)[\"image\"]\n            right_tensor = self.transform(image=right_img)[\"image\"]\n        else:\n            left_tensor = torch.from_numpy(left_img).permute(2, 0, 1).float() / 255.0\n            right_tensor = torch.from_numpy(right_img).permute(2, 0, 1).float() / 255.0\n        features = torch.tensor(row[self.feature_cols].values.astype(np.float32), dtype=torch.float32)\n        \n        if self.is_test:\n            return left_tensor, right_tensor, features, row['image_path']\n        else:\n            # Targets\n            targets = torch.tensor(row[list(CFG.TARGET_COLS)].values.astype(np.float32), dtype=torch.float32)\n            return left_tensor, right_tensor, features, targets\n\n\ndef get_train_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.5),\n        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2(),\n    ])\n\ndef get_valid_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2(),\n    ])\n\ndef get_tta_transforms():\n    \"\"\"Test Time Augmentation transforms (competitor's approach)\"\"\"\n    base = [\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2(),\n    ]\n    \n    return [\n        # Original\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n            *base,\n        ]),\n        # Horizontal flip\n        A.Compose([\n            A.HorizontalFlip(p=1.0),\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n            *base,\n        ]),\n        # Vertical flip\n        A.Compose([\n            A.VerticalFlip(p=1.0),\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n            *base,\n        ]),\n        # 90 degree rotation\n        A.Compose([\n            A.RandomRotate90(p=1.0),\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE, interpolation=cv2.INTER_AREA),\n            *base,\n        ]),\n    ]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.480201Z","iopub.execute_input":"2026-01-03T06:32:39.480873Z","iopub.status.idle":"2026-01-03T06:32:39.498196Z","shell.execute_reply.started":"2026-01-03T06:32:39.480844Z","shell.execute_reply":"2026-01-03T06:32:39.497365Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Model Building Blocks\nclass FeedForward(nn.Module):\n    def __init__(self, dim, mlp_ratio=4.0, dropout=0.0):\n        super().__init__()\n        hid = int(dim * mlp_ratio)\n        self.net = nn.Sequential(\n            nn.Linear(dim, hid),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hid, dim),\n            nn.Dropout(dropout),\n        )\n\n    def forward(self, x):\n        return self.net(x)\n\n\nclass AttentionBlock(nn.Module):\n    def __init__(self, dim, heads=8, dropout=0.0, mlp_ratio=4.0):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)\n        self.norm2 = nn.LayerNorm(dim)\n        self.ff = FeedForward(dim, mlp_ratio=mlp_ratio, dropout=dropout)\n\n    def forward(self, x):\n        h = self.norm1(x)\n        attn_out, _ = self.attn(h, h, h, need_weights=False)\n        x = x + attn_out\n        x = x + self.ff(self.norm2(x))\n        return x\n\n\nclass MobileViTBlock(nn.Module):\n    def __init__(self, dim, heads=4, depth=2, patch=(2, 2), dropout=0.0):\n        super().__init__()\n        self.local = nn.Sequential(\n            nn.Conv2d(dim, dim, 3, padding=1, groups=dim),\n            nn.Conv2d(dim, dim, 1),\n            nn.GELU(),\n        )\n        self.patch = patch\n        self.transformer = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(depth)]\n        )\n        self.fuse = nn.Conv2d(dim * 2, dim, kernel_size=1)\n\n    def forward(self, x: torch.Tensor):\n        local_feat = self.local(x)\n        B, C, H, W = local_feat.shape\n        ph, pw = self.patch\n        new_h = math.ceil(H / ph) * ph\n        new_w = math.ceil(W / pw) * pw\n        if new_h != H or new_w != W:\n            local_feat = F.interpolate(local_feat, size=(new_h, new_w), mode=\"bilinear\", align_corners=False)\n            H, W = new_h, new_w\n\n        tokens = local_feat.unfold(2, ph, ph).unfold(3, pw, pw)\n        tokens = tokens.contiguous().view(B, C, -1, ph, pw)\n        tokens = tokens.permute(0, 2, 3, 4, 1).reshape(B, -1, C)\n\n        for blk in self.transformer:\n            tokens = blk(tokens)\n\n        feat = tokens.view(B, -1, ph * pw, C).permute(0, 3, 1, 2)\n        nh = H // ph\n        nw = W // pw\n        feat = feat.view(B, C, nh, nw, ph, pw).permute(0, 1, 2, 4, 3, 5)\n        feat = feat.reshape(B, C, H, W)\n\n        if feat.shape[-2:] != x.shape[-2:]:\n            feat = F.interpolate(feat, size=x.shape[-2:], mode=\"bilinear\", align_corners=False)\n\n        out = self.fuse(torch.cat([x, feat], dim=1))\n        return out\n\n\nclass SpatialReductionAttention(nn.Module):\n    def __init__(self, dim, heads=8, sr_ratio=2, dropout=0.0):\n        super().__init__()\n        self.heads = heads\n        self.scale = (dim // heads) ** -0.5\n        self.q = nn.Linear(dim, dim)\n        self.kv = nn.Linear(dim, dim * 2)\n        self.sr_ratio = sr_ratio\n        if sr_ratio > 1:\n            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)\n            self.norm = nn.LayerNorm(dim)\n        else:\n            self.sr = None\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x, hw: Tuple[int, int]):\n        B, N, C = x.shape\n        q = self.q(x).reshape(B, N, self.heads, C // self.heads).permute(0, 2, 1, 3)\n\n        if self.sr is not None:\n            H, W = hw\n            feat = x.transpose(1, 2).reshape(B, C, H, W)\n            feat = self.sr(feat)\n            feat = feat.reshape(B, C, -1).transpose(1, 2)\n            feat = self.norm(feat)\n        else:\n            feat = x\n\n        kv = self.kv(feat)\n        k, v = kv.chunk(2, dim=-1)\n        k = k.reshape(B, -1, self.heads, C // self.heads).permute(0, 2, 3, 1)\n        v = v.reshape(B, -1, self.heads, C // self.heads).permute(0, 2, 1, 3)\n\n        attn = torch.matmul(q, k) * self.scale\n        attn = attn.softmax(dim=-1)\n        attn = self.drop(attn)\n        out = torch.matmul(attn, v).permute(0, 2, 1, 3).reshape(B, N, C)\n        out = self.proj(out)\n        return out\n\n\nclass PVTBlock(nn.Module):\n    def __init__(self, dim, heads=8, sr_ratio=2, dropout=0.0, mlp_ratio=4.0):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.sra = SpatialReductionAttention(dim, heads=heads, sr_ratio=sr_ratio, dropout=dropout)\n        self.norm2 = nn.LayerNorm(dim)\n        self.ff = FeedForward(dim, mlp_ratio=mlp_ratio, dropout=dropout)\n\n    def forward(self, x, hw: Tuple[int, int]):\n        x = x + self.sra(self.norm1(x), hw)\n        x = x + self.ff(self.norm2(x))\n        return x\n\n\nclass LocalMambaBlock(nn.Module):\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\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        shortcut = x\n        x = self.norm(x)\n        g = torch.sigmoid(self.gate(x))\n        x = (x * g).transpose(1, 2)\n        x = self.dwconv(x).transpose(1, 2)\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\nclass T2TRetokenizer(nn.Module):\n    def __init__(self, dim, depth=2, heads=4, dropout=0.0):\n        super().__init__()\n        self.blocks = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(depth)]\n        )\n\n    def forward(self, tokens: torch.Tensor, grid_hw: Tuple[int, int]):\n        B, T, C = tokens.shape\n        H, W = grid_hw\n        feat_map = tokens.transpose(1, 2).reshape(B, C, H, W)\n        seq = feat_map.flatten(2).transpose(1, 2)\n        for blk in self.blocks:\n            seq = blk(seq)\n        seq_map = seq.transpose(1, 2).reshape(B, C, H, W)\n        pooled = F.adaptive_avg_pool2d(seq_map, (2, 2))\n        retokens = pooled.flatten(2).transpose(1, 2)\n        return retokens, seq_map\n\n\nclass CrossScaleFusion(nn.Module):\n    def __init__(self, dim, heads=6, dropout=0.0, layers=2):\n        super().__init__()\n        self.layers_s = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(layers)]\n        )\n        self.layers_b = nn.ModuleList(\n            [AttentionBlock(dim, heads=heads, dropout=dropout, mlp_ratio=2.0) for _ in range(layers)]\n        )\n        self.cross_s = nn.ModuleList(\n            [nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) for _ in range(layers)]\n        )\n        self.cross_b = nn.ModuleList(\n            [nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) for _ in range(layers)]\n        )\n        self.norm_s = nn.LayerNorm(dim)\n        self.norm_b = nn.LayerNorm(dim)\n\n    def forward(self, tok_s: torch.Tensor, tok_b: torch.Tensor):\n        B, Ts, C = tok_s.shape\n        cls_s = tok_s.new_zeros(B, 1, C)\n        cls_b = tok_b.new_zeros(B, 1, C)\n        tok_s = torch.cat([cls_s, tok_s], dim=1)\n        tok_b = torch.cat([cls_b, tok_b], dim=1)\n\n        for ls, lb, cs, cb in zip(self.layers_s, self.layers_b, self.cross_s, self.cross_b):\n            tok_s = ls(tok_s)\n            tok_b = lb(tok_b)\n            q_s = self.norm_s(tok_s[:, :1])\n            q_b = self.norm_b(tok_b[:, :1])\n            cls_s_upd, _ = cs(q_s, torch.cat([tok_b, q_b], dim=1), torch.cat([tok_b, q_b], dim=1), need_weights=False)\n            cls_b_upd, _ = cb(q_b, torch.cat([tok_s, q_s], dim=1), torch.cat([tok_s, q_s], dim=1), need_weights=False)\n            tok_s = torch.cat([tok_s[:, :1] + cls_s_upd, tok_s[:, 1:]], dim=1)\n            tok_b = torch.cat([tok_b[:, :1] + cls_b_upd, tok_b[:, 1:]], dim=1)\n\n        tokens = torch.cat([tok_s[:, :1], tok_b[:, :1], tok_s[:, 1:], tok_b[:, 1:]], dim=1)\n        return tokens\n\n\nclass TileEncoder(nn.Module):\n    def __init__(self, backbone: nn.Module, input_res: int):\n        super().__init__()\n        self.backbone = backbone\n        self.input_res = input_res\n\n    def forward(self, x: torch.Tensor, grid: Tuple[int, int]):\n        B, C, H, W = x.shape\n        r, c = grid\n        hs = torch.linspace(0, H, steps=r + 1, device=x.device).round().long()\n        ws = torch.linspace(0, W, steps=c + 1, device=x.device).round().long()\n        tiles = []\n        for i in range(r):\n            for j in range(c):\n                rs, re = hs[i].item(), hs[i + 1].item()\n                cs, ce = ws[j].item(), ws[j + 1].item()\n                xt = x[:, :, rs:re, cs:ce]\n                if xt.shape[-2:] != (self.input_res, self.input_res):\n                    xt = F.interpolate(xt, size=(self.input_res, self.input_res), mode=\"bilinear\", align_corners=False)\n                tiles.append(xt)\n        tiles = torch.stack(tiles, dim=1)\n        flat = tiles.view(-1, C, self.input_res, self.input_res)\n        feats = self.backbone(flat)\n        feats = feats.view(B, -1, feats.shape[-1])\n        return feats\n\n\nclass PyramidMixer(nn.Module):\n    def __init__(\n        self,\n        dim_in: int,\n        dims: Tuple[int, int, int],\n        mobilevit_heads: int = 4,\n        mobilevit_depth: int = 2,\n        sra_heads: int = 6,\n        sra_ratio: int = 2,\n        mamba_depth: int = 3,\n        mamba_kernel: int = 5,\n        dropout: float = 0.0,\n    ):\n        super().__init__()\n        c1, c2, c3 = dims\n        self.proj1 = nn.Linear(dim_in, c1)\n        self.mobilevit = MobileViTBlock(c1, heads=mobilevit_heads, depth=mobilevit_depth, dropout=dropout)\n        self.proj2 = nn.Linear(c1, c2)\n        self.pvt = PVTBlock(c2, heads=sra_heads, sr_ratio=sra_ratio, dropout=dropout, mlp_ratio=3.0)\n        self.mamba_local = LocalMambaBlock(c2, kernel_size=mamba_kernel, dropout=dropout)\n        self.proj3 = nn.Linear(c2, c3)\n        self.mamba_global = nn.ModuleList(\n            [LocalMambaBlock(c3, kernel_size=mamba_kernel, dropout=dropout) for _ in range(mamba_depth)]\n        )\n        self.final_attn = AttentionBlock(c3, heads=min(8, c3 // 64 + 1), dropout=dropout, mlp_ratio=2.0)\n\n    def _tokens_to_map(self, tokens: torch.Tensor, target_hw: Tuple[int, int]):\n        B, N, C = tokens.shape\n        H, W = target_hw\n        need = H * W\n        if N < need:\n            pad = tokens.new_zeros(B, need - N, C)\n            tokens = torch.cat([tokens, pad], dim=1)\n        tokens = tokens[:, :need, :]\n        feat_map = tokens.transpose(1, 2).reshape(B, C, H, W)\n        return feat_map\n\n    @staticmethod\n    def _fit_hw(n_tokens: int) -> Tuple[int, int]:\n        h = int(math.sqrt(n_tokens))\n        w = h\n        while h * w < n_tokens:\n            w += 1\n            if h * w < n_tokens:\n                h += 1\n        return h, w\n\n    def forward(self, tokens: torch.Tensor):\n        B, N, C = tokens.shape\n        map_hw = (3, 4)\n        \n        t1 = self.proj1(tokens)\n        m1 = self._tokens_to_map(t1, map_hw)\n        m1 = self.mobilevit(m1)\n        t1_out = m1.flatten(2).transpose(1, 2)[:, :N]\n\n        t2 = self.proj2(t1_out)\n        new_len = max(4, N // 2)\n        t2 = t2[:, :new_len] + F.adaptive_avg_pool1d(t2.transpose(1, 2), new_len).transpose(1, 2)\n        hw2 = self._fit_hw(t2.size(1))\n        if t2.size(1) < hw2[0] * hw2[1]:\n            pad = t2.new_zeros(B, hw2[0] * hw2[1] - t2.size(1), t2.size(2))\n            t2 = torch.cat([t2, pad], dim=1)\n        t2 = self.pvt(t2, hw2)\n        t2 = self.mamba_local(t2)\n\n        t3 = self.proj3(t2)\n        pooled = torch.stack([t3.mean(dim=1), t3.max(dim=1).values], dim=1)\n        t3 = pooled\n        for blk in self.mamba_global:\n            t3 = blk(t3)\n        t3 = self.final_attn(t3)\n        global_feat = t3.mean(dim=1)\n        return global_feat, {\"stage1_map\": m1.detach(), \"stage2_tokens\": t2.detach(), \"stage3_tokens\": t3.detach()}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.499242Z","iopub.execute_input":"2026-01-03T06:32:39.499469Z","iopub.status.idle":"2026-01-03T06:32:39.551658Z","shell.execute_reply.started":"2026-01-03T06:32:39.499448Z","shell.execute_reply":"2026-01-03T06:32:39.551001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Main Model (Combined Architecture)\nclass CrossPVT_T2T_MambaDINO(nn.Module):\n    \"\"\"Competitor's original model architecture (no tabular features)\"\"\"\n    \n    def __init__(self, dropout: float = 0.1, hidden_ratio: float = 0.35):\n        super().__init__()\n        self.backbone, self.feat_dim, self.backbone_name, self.input_res = self._build_dino_backbone()\n        self.tile_encoder = TileEncoder(self.backbone, self.input_res)\n        self.t2t = T2TRetokenizer(self.feat_dim, depth=CFG.T2T_DEPTH, heads=CFG.CROSS_HEADS, dropout=dropout)\n        self.cross = CrossScaleFusion(self.feat_dim, heads=CFG.CROSS_HEADS, dropout=dropout, layers=CFG.CROSS_LAYERS)\n        self.pyramid = PyramidMixer(\n            dim_in=self.feat_dim,\n            dims=CFG.PYRAMID_DIMS,\n            mobilevit_heads=CFG.MOBILEVIT_HEADS,\n            mobilevit_depth=CFG.MOBILEVIT_DEPTH,\n            sra_heads=CFG.SRA_HEADS,\n            sra_ratio=CFG.SRA_RATIO,\n            mamba_depth=CFG.MAMBA_DEPTH,\n            mamba_kernel=CFG.MAMBA_KERNEL,\n            dropout=dropout,\n        )\n\n        combined = CFG.PYRAMID_DIMS[-1] * 2\n        self.combined_dim = combined\n        hidden = max(32, int(combined * hidden_ratio))\n\n        def head():\n            return nn.Sequential(\n                nn.Linear(combined, hidden),\n                nn.GELU(),\n                nn.Dropout(dropout),\n                nn.Linear(hidden, 1),\n            )\n\n        self.head_green = head()\n        self.head_clover = head()\n        self.head_dead = head()\n        self.aux_head = (\n            nn.Sequential(nn.LayerNorm(CFG.PYRAMID_DIMS[1]), nn.Linear(CFG.PYRAMID_DIMS[1], 5))\n            if CFG.AUX_HEAD else None\n        )\n        self.softplus = nn.Softplus(beta=1.0)\n        self.cross_gate_left = nn.Linear(CFG.PYRAMID_DIMS[-1], CFG.PYRAMID_DIMS[-1])\n        self.cross_gate_right = nn.Linear(CFG.PYRAMID_DIMS[-1], CFG.PYRAMID_DIMS[-1])\n\n    def _build_dino_backbone(self):\n        last_err = None\n        for name in CFG.DINO_CANDIDATES:\n            for gp in [\"token\", \"avg\", \"__default__\"]:\n                try:\n                    if gp == \"__default__\":\n                        m = timm.create_model(name, pretrained=False, num_classes=0)\n                    else:\n                        m = timm.create_model(name, pretrained=False, num_classes=0, global_pool=gp)\n                    feat = m.num_features\n                    input_res = self._infer_input_res(m)\n                    if hasattr(m, \"set_grad_checkpointing\"):\n                        m.set_grad_checkpointing(True)\n                    return m, feat, name, int(input_res)\n                except Exception as e:\n                    last_err = e\n                    continue\n        raise RuntimeError(f\"Cannot create DINO backbone. Last error: {last_err}\")\n\n    @staticmethod\n    def _infer_input_res(m) -> int:\n        if hasattr(m, \"patch_embed\") and hasattr(m.patch_embed, \"img_size\"):\n            isz = m.patch_embed.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        if hasattr(m, \"img_size\"):\n            isz = m.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        dc = getattr(m, \"default_cfg\", {}) or {}\n        ins = dc.get(\"input_size\", None)\n        if ins:\n            if isinstance(ins, (tuple, list)) and len(ins) >= 2:\n                return int(ins[1])\n            return int(ins if isinstance(ins, (int, float)) else 224)\n        return 518\n\n    def _half_forward(self, x_half: torch.Tensor):\n        tiles_small = self.tile_encoder(x_half, CFG.SMALL_GRID)\n        tiles_big = self.tile_encoder(x_half, CFG.BIG_GRID)\n        t2, stage1_map = self.t2t(tiles_small, CFG.SMALL_GRID)\n        fused = self.cross(t2, tiles_big)\n        feat, feat_maps = self.pyramid(fused)\n        feat_maps[\"stage1_map\"] = stage1_map\n        return feat, feat_maps\n\n    def forward(self, x_left, x_right, return_features=False):\n        \"\"\"Forward pass - compatible with competitor's checkpoint\"\"\"\n        feat_l, feats_l = self._half_forward(x_left)\n        feat_r, feats_r = self._half_forward(x_right)\n        \n        g_l = torch.sigmoid(self.cross_gate_left(feat_r))\n        g_r = torch.sigmoid(self.cross_gate_right(feat_l))\n        feat_l = feat_l * g_l\n        feat_r = feat_r * g_r\n        f = torch.cat([feat_l, feat_r], dim=1)\n        \n        green = self.softplus(self.head_green(f))\n        clover = self.softplus(self.head_clover(f))\n        dead = self.softplus(self.head_dead(f))\n        gdm = green + clover\n        total = gdm + dead\n        \n        # Return in same order as TARGET_COLS: [Green, Dead, Clover, Total, GDM]\n        output = torch.cat([green, dead, clover, total, gdm], dim=1)\n        return output\n\nclass CombinedBiomassModel(pl.LightningModule):\n    \"\"\"\n    Combined model that uses:\n    - Competitor's DINOv2 tiled architecture with left/right split\n    - YOUR tabular features\n    - Physical constraint prediction (predict 3, compute 2)\n    \"\"\"\n    def __init__(\n        self, \n        num_tabular_features: int = 9,\n        dropout: float = CFG.DROPOUT, \n        hidden_ratio: float = CFG.HIDDEN_RATIO,\n        lr: float = CFG.LR,\n        weight_decay: float = CFG.WEIGHT_DECAY\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        \n        # Build DINOv2 backbone\n        self.backbone, self.feat_dim, self.backbone_name, self.input_res = self._build_dino_backbone()\n        print(f\"✅ DINOv2 backbone: {self.backbone_name}, feat_dim: {self.feat_dim}, input_res: {self.input_res}\")\n        \n        # Tile encoder for multi-scale processing\n        self.tile_encoder = TileEncoder(self.backbone, self.input_res)\n        \n        # T2T Retokenizer\n        self.t2t = T2TRetokenizer(self.feat_dim, depth=CFG.T2T_DEPTH, heads=CFG.CROSS_HEADS, dropout=dropout)\n        \n        # Cross-scale fusion\n        self.cross = CrossScaleFusion(self.feat_dim, heads=CFG.CROSS_HEADS, dropout=dropout, layers=CFG.CROSS_LAYERS)\n        \n        # Pyramid mixer\n        self.pyramid = PyramidMixer(\n            dim_in=self.feat_dim,\n            dims=CFG.PYRAMID_DIMS,\n            mobilevit_heads=CFG.MOBILEVIT_HEADS,\n            mobilevit_depth=CFG.MOBILEVIT_DEPTH,\n            sra_heads=CFG.SRA_HEADS,\n            sra_ratio=CFG.SRA_RATIO,\n            mamba_depth=CFG.MAMBA_DEPTH,\n            mamba_kernel=CFG.MAMBA_KERNEL,\n            dropout=dropout,\n        )\n        \n        # Cross gates for left/right fusion\n        self.cross_gate_left = nn.Linear(CFG.PYRAMID_DIMS[-1], CFG.PYRAMID_DIMS[-1])\n        self.cross_gate_right = nn.Linear(CFG.PYRAMID_DIMS[-1], CFG.PYRAMID_DIMS[-1])\n        \n        # YOUR IMPROVEMENT: Tabular feature encoder\n        self.tabular_encoder = nn.Sequential(\n            nn.Linear(num_tabular_features, 64),\n            nn.BatchNorm1d(64),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(64, 128),\n            nn.BatchNorm1d(128),\n            nn.GELU(),\n            nn.Dropout(dropout),\n        )\n        \n        # Combined dimension: image features (left + right) + tabular features\n        img_combined_dim = CFG.PYRAMID_DIMS[-1] * 2\n        combined_dim = img_combined_dim + 128  # + tabular features\n        self.combined_dim = combined_dim\n        \n        # YOUR IMPROVEMENT: Attention for feature fusion\n        self.fusion_attention = nn.Sequential(\n            nn.Linear(combined_dim, combined_dim // 4),\n            nn.GELU(),\n            nn.Linear(combined_dim // 4, combined_dim),\n            nn.Sigmoid()\n        )\n        \n        # Prediction heads (predict 3, compute 2 - COMPETITOR's approach)\n        hidden = max(32, int(combined_dim * hidden_ratio))\n        \n        def make_head():\n            return nn.Sequential(\n                nn.Linear(combined_dim, hidden),\n                nn.GELU(),\n                nn.Dropout(dropout),\n                nn.Linear(hidden, 1),\n            )\n        \n        self.head_green = make_head()\n        self.head_clover = make_head()\n        self.head_dead = make_head()\n        \n        # Softplus for non-negative outputs (COMPETITOR's approach)\n        self.softplus = nn.Softplus(beta=1.0)\n        \n        # Auxiliary head (COMPETITOR's approach)\n        if CFG.AUX_HEAD:\n            self.aux_head = nn.Sequential(\n                nn.LayerNorm(CFG.PYRAMID_DIMS[1]),\n                nn.Linear(CFG.PYRAMID_DIMS[1], 5)\n            )\n        else:\n            self.aux_head = None\n        \n        # Loss function\n        self.criterion = nn.SmoothL1Loss()\n        \n        # Validation outputs\n        self.val_outputs = []\n        \n        print(f\"✅ Model initialized with combined_dim: {combined_dim}\")\n    \n    def _build_dino_backbone(self):\n        \"\"\"Build DINOv2 backbone with fallback options\"\"\"\n        last_err = None\n        for name in CFG.DINO_CANDIDATES:\n            for gp in [\"token\", \"avg\", \"__default__\"]:\n                try:\n                    if gp == \"__default__\":\n                        m = timm.create_model(name, pretrained=False, num_classes=0)  # Changed to False\n                    else:\n                        m = timm.create_model(name, pretrained=False, num_classes=0, global_pool=gp)  # Changed to False\n                    \n                    feat = m.num_features\n                    input_res = self._infer_input_res(m)\n                    \n                    if hasattr(m, \"set_grad_checkpointing\"):\n                        m.set_grad_checkpointing(True)\n                    \n                    return m, feat, name, int(input_res)\n                except Exception as e:\n                    last_err = e\n                    continue\n        raise RuntimeError(f\"Cannot create DINO backbone. Last error: {last_err}\")\n    \n    @staticmethod\n    def _infer_input_res(m) -> int:\n        if hasattr(m, \"patch_embed\") and hasattr(m.patch_embed, \"img_size\"):\n            isz = m.patch_embed.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        if hasattr(m, \"img_size\"):\n            isz = m.img_size\n            return int(isz if isinstance(isz, (int, float)) else isz[0])\n        dc = getattr(m, \"default_cfg\", {}) or {}\n        ins = dc.get(\"input_size\", None)\n        if ins:\n            if isinstance(ins, (tuple, list)) and len(ins) >= 2:\n                return int(ins[1])\n            return int(ins if isinstance(ins, (int, float)) else 224)\n        return 518\n    \n    def _half_forward(self, x_half: torch.Tensor):\n        \"\"\"Process one half of the image through tiled encoding\"\"\"\n        tiles_small = self.tile_encoder(x_half, CFG.SMALL_GRID)\n        tiles_big = self.tile_encoder(x_half, CFG.BIG_GRID)\n        t2, stage1_map = self.t2t(tiles_small, CFG.SMALL_GRID)\n        fused = self.cross(t2, tiles_big)\n        feat, feat_maps = self.pyramid(fused)\n        feat_maps[\"stage1_map\"] = stage1_map\n        return feat, feat_maps\n    \n    def forward(self, x_left: torch.Tensor, x_right: torch.Tensor, tabular_features: torch.Tensor):\n        \"\"\"\n        Forward pass with:\n        - Left/right image halves (competitor)\n        - Tabular features (YOUR improvement)\n        \"\"\"\n        # Process left and right image halves\n        feat_l, feats_l = self._half_forward(x_left)\n        feat_r, feats_r = self._half_forward(x_right)\n        \n        # Cross-gate fusion (competitor's approach)\n        g_l = torch.sigmoid(self.cross_gate_left(feat_r))\n        g_r = torch.sigmoid(self.cross_gate_right(feat_l))\n        feat_l = feat_l * g_l\n        feat_r = feat_r * g_r\n        \n        # Concatenate image features\n        img_combined = torch.cat([feat_l, feat_r], dim=1)\n        \n        # YOUR IMPROVEMENT: Encode tabular features\n        tab_feat = self.tabular_encoder(tabular_features)\n        \n        # Combine image and tabular features\n        combined = torch.cat([img_combined, tab_feat], dim=1)\n        \n        # YOUR IMPROVEMENT: Apply fusion attention\n        attention_weights = self.fusion_attention(combined)\n        combined = combined * attention_weights\n        \n        # COMPETITOR's approach: Predict 3 targets, compute 2\n        green = self.softplus(self.head_green(combined))\n        clover = self.softplus(self.head_clover(combined))\n        dead = self.softplus(self.head_dead(combined))\n        \n        # Physical constraints\n        gdm = green + clover\n        total = gdm + dead\n        \n        # Output: [Green, Dead, Clover, Total, GDM] (matches TARGET_COLS order)\n        output = torch.cat([green, dead, clover, total, gdm], dim=1)\n        \n        # Auxiliary output for auxiliary loss\n        aux_output = None\n        if self.aux_head is not None and self.training:\n            aux_tokens = torch.cat([feats_l[\"stage2_tokens\"], feats_r[\"stage2_tokens\"]], dim=1)\n            aux_output = self.softplus(self.aux_head(aux_tokens.mean(dim=1)))\n        \n        return output, aux_output\n    \n    def training_step(self, batch, batch_idx):\n        x_left, x_right, features, targets = batch\n        preds, aux_preds = self(x_left, x_right, features)\n        \n        # Main loss\n        loss = self.criterion(preds, targets)\n        \n        # Auxiliary loss (if enabled)\n        if aux_preds is not None:\n            aux_loss = self.criterion(aux_preds, targets)\n            loss = loss + CFG.AUX_LOSS_WEIGHT * aux_loss\n            self.log('aux_loss', aux_loss, on_step=False, on_epoch=True)\n        \n        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        x_left, x_right, features, targets = batch\n        preds, _ = self(x_left, x_right, features)\n        loss = self.criterion(preds, targets)\n        \n        self.val_outputs.append((preds.detach().cpu(), targets.detach().cpu()))\n        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True)\n        return loss\n    \n    def on_validation_epoch_end(self):\n        if len(self.val_outputs) == 0:\n            return\n        \n        preds, targets = zip(*self.val_outputs)\n        preds = torch.cat(preds).numpy()\n        targets = torch.cat(targets).numpy()\n        \n        weighted_r2, r2_scores = weighted_r2_score(targets, preds)\n        \n        self.log('val_weighted_r2', weighted_r2, prog_bar=True, on_epoch=True)\n        for i, name in enumerate(CFG.TARGET_COLS):\n            self.log(f'val_r2_{name}', r2_scores[i], on_epoch=True)\n        \n        self.val_outputs.clear()\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(\n            self.parameters(),\n            lr=self.hparams.lr,\n            weight_decay=self.hparams.weight_decay\n        )\n        \n        # YOUR improvement: Cosine annealing with warm restarts\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n            optimizer,\n            T_0=10,\n            T_mult=2,\n            eta_min=1e-6\n        )\n        \n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': {\n                'scheduler': scheduler,\n                'interval': 'epoch'\n            }\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.552685Z","iopub.execute_input":"2026-01-03T06:32:39.553326Z","iopub.status.idle":"2026-01-03T06:32:39.592646Z","shell.execute_reply.started":"2026-01-03T06:32:39.553294Z","shell.execute_reply":"2026-01-03T06:32:39.591982Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Data Module\nclass BiomassDataModule(pl.LightningDataModule):\n    def __init__(self, train_df, val_df, batch_size=CFG.BATCH_SIZE, num_workers=CFG.NUM_WORKERS):\n        super().__init__()\n        self.train_df = train_df\n        self.val_df = val_df\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n    \n    def setup(self, stage=None):\n        self.train_dataset = BiomassDatasetDINO(self.train_df, transform=get_train_transforms())\n        self.val_dataset = BiomassDatasetDINO(self.val_df, transform=get_valid_transforms())\n    \n    def train_dataloader(self):\n        return DataLoader(\n            self.train_dataset,\n            batch_size=self.batch_size,\n            shuffle=True,\n            num_workers=self.num_workers,\n            pin_memory=True,\n            persistent_workers=True if self.num_workers > 0 else False\n        )\n    \n    def val_dataloader(self):\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            num_workers=self.num_workers,\n            pin_memory=True,\n            persistent_workers=True if self.num_workers > 0 else False\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.593841Z","iopub.execute_input":"2026-01-03T06:32:39.594424Z","iopub.status.idle":"2026-01-03T06:32:39.614572Z","shell.execute_reply.started":"2026-01-03T06:32:39.594398Z","shell.execute_reply":"2026-01-03T06:32:39.613639Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Training Function\ndef train_kfold():\n    \"\"\"Train models using YOUR Stratified K-Fold cross-validation\"\"\"\n    \n    trained_models = []\n    oof_predictions = np.zeros((len(train_df), 5))\n    oof_targets = np.zeros((len(train_df), 5))\n    \n    # YOUR improvement: Stratified K-Fold based on target bins\n    skf = StratifiedKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\n    \n    for fold, (train_idx, val_idx) in enumerate(skf.split(train_df, train_df['Dry_Total_bin'])):\n        print(f\"\\n{'='*60}\")\n        print(f\"🚀 Training Fold {fold + 1}/{CFG.N_FOLDS}\")\n        print(f\"{'='*60}\")\n        \n        # Split data\n        train_fold = train_df.iloc[train_idx]\n        val_fold = train_df.iloc[val_idx]\n        \n        print(f\"Train size: {len(train_fold)}, Val size: {len(val_fold)}\")\n        \n        # Create data module\n        datamodule = BiomassDataModule(train_fold, val_fold)\n        \n        # Create model\n        model = CombinedBiomassModel()\n        \n        # Callbacks\n        checkpoint_callback = ModelCheckpoint(\n            dirpath=f'checkpoints/fold_{fold}',\n            filename='best_wr2',\n            monitor='val_weighted_r2',\n            mode='max',\n            save_top_k=1,\n            save_weights_only=False\n        )\n        \n        early_stop_callback = EarlyStopping(\n            monitor='val_weighted_r2',\n            patience=12,\n            mode='max',\n            verbose=True\n        )\n        \n        lr_monitor = LearningRateMonitor(logging_interval='epoch')\n        \n        # Trainer\n        trainer = pl.Trainer(\n            max_epochs=CFG.EPOCHS,\n            accelerator='gpu',\n            devices=1,\n            precision='16-mixed' if CFG.MIXED_PRECISION else 32,\n            callbacks=[checkpoint_callback, early_stop_callback, lr_monitor],\n            accumulate_grad_batches=CFG.ACCUMULATE_GRAD_BATCHES,\n            gradient_clip_val=1.0,\n            log_every_n_steps=5,\n            enable_progress_bar=True,\n            enable_model_summary=fold == 0,  # Only show summary for first fold\n            deterministic=False\n        )\n        \n        # Train\n        trainer.fit(model, datamodule=datamodule)\n        \n        # Load best model\n        best_model = CombinedBiomassModel.load_from_checkpoint(checkpoint_callback.best_model_path)\n        best_model.eval()\n        best_model.to('cuda')\n        \n        # Generate OOF predictions\n        val_dataset = BiomassDatasetDINO(val_fold, transform=get_valid_transforms())\n        val_loader = DataLoader(\n            val_dataset, \n            batch_size=CFG.BATCH_SIZE, \n            shuffle=False, \n            num_workers=CFG.NUM_WORKERS, \n            pin_memory=True,\n            persistent_workers=True if CFG.NUM_WORKERS > 0 else False\n        )\n        \n        val_preds = []\n        val_targets_list = []\n        \n        with torch.no_grad():\n            for x_left, x_right, features, targets in val_loader:\n                x_left = x_left.to(best_model.device)\n                x_right = x_right.to(best_model.device)\n                features = features.to(best_model.device)\n                \n                preds, _ = best_model(x_left, x_right, features)\n                val_preds.append(preds.cpu().numpy())\n                val_targets_list.append(targets.numpy())\n        \n        val_preds = np.concatenate(val_preds)\n        val_targets_arr = np.concatenate(val_targets_list)\n        \n        # Store OOF predictions\n        oof_predictions[val_idx] = val_preds\n        oof_targets[val_idx] = val_targets_arr\n        \n        # Calculate fold metrics\n        fold_r2, fold_r2_scores = weighted_r2_score(val_targets_arr, val_preds)\n        print(f\"\\n📊 Fold {fold + 1} Results:\")\n        print(f\"Weighted R²: {fold_r2:.4f}\")\n        for i, name in enumerate(CFG.TARGET_COLS):\n            print(f\"  {name}: {fold_r2_scores[i]:.4f}\")\n        \n        # Save model\n        trained_models.append(best_model)\n        \n        # Cleanup\n        del model, datamodule, trainer\n        torch.cuda.empty_cache()\n        gc.collect()\n    \n    # Calculate overall OOF score\n    oof_r2, oof_r2_scores = weighted_r2_score(oof_targets, oof_predictions)\n    print(f\"\\n{'='*60}\")\n    print(f\"📈 Overall OOF Results:\")\n    print(f\"{'='*60}\")\n    print(f\"Weighted R²: {oof_r2:.4f}\")\n    for i, name in enumerate(CFG.TARGET_COLS):\n        print(f\"  {name}: {oof_r2_scores[i]:.4f}\")\n    \n    return trained_models\n\n# In Cell 10, replace load_pretrained_models function:\n\ndef load_pretrained_models():\n    \"\"\"Load competitor's models from Kaggle input\"\"\"\n    models = []\n    \n    for fold in range(CFG.N_FOLDS):\n        # Try both naming patterns\n        ckpt_path = f'/kaggle/input/csiro/pytorch/default/12/fold_{fold}/checkpoints/best_wr2.pt'\n        alt_path = f'/kaggle/input/csiro/pytorch/default/12/fold{fold}/checkpoints/best_wr2.pt'\n        \n        if os.path.exists(ckpt_path):\n            path = ckpt_path\n        elif os.path.exists(alt_path):\n            path = alt_path\n        else:\n            print(f\"⚠️ Checkpoint not found for fold {fold}\")\n            continue\n        \n        print(f\"📦 Loading fold {fold} from: {path}\")\n        \n        # Load checkpoint\n        state = torch.load(path, map_location='cpu', weights_only=False)\n        \n        # Get config from checkpoint\n        cfg_dict = state.get(\"cfg\", {})\n        dropout = cfg_dict.get(\"dropout\", CFG.DROPOUT)\n        hidden_ratio = cfg_dict.get(\"hidden_ratio\", CFG.HIDDEN_RATIO)\n        \n        # Create COMPETITOR's model (not combined model)\n        model = CrossPVT_T2T_MambaDINO(dropout=dropout, hidden_ratio=hidden_ratio)\n        \n        # Get state dict\n        model_state = state.get('model_state', state)\n        \n        # Strip 'module.' prefix if present\n        if any(k.startswith('module.') for k in model_state.keys()):\n            model_state = {k.replace('module.', ''): v for k, v in model_state.items()}\n        \n        # Load weights\n        model.load_state_dict(model_state, strict=False)\n        model.eval()\n        model.to('cuda' if torch.cuda.is_available() else 'cpu')\n        \n        models.append(model)\n        print(f\"✅ Loaded fold {fold}\")\n    \n    return models\n\n# Check if we should train or load\nif os.path.exists(CFG.MODEL_DIR):\n    print(\"\\n\" + \"=\"*80)\n    print(\"📦 Loading Pre-trained Models from Kaggle Input\")\n    print(\"=\"*80)\n    trained_models = load_pretrained_models()\n    print(f\"\\n✅ Loaded {len(trained_models)} models\")\nelse:\n    print(\"\\n\" + \"=\"*80)\n    print(\"🚀 Starting Combined Model Training\")\n    print(\"=\"*80)\n    print(f\"Architecture: DINOv2 + Tabular Features\")\n    print(f\"Total Folds: {CFG.N_FOLDS}\")\n    print(f\"Max Epochs per Fold: {CFG.EPOCHS}\")\n    print(f\"Batch Size: {CFG.BATCH_SIZE} x {CFG.ACCUMULATE_GRAD_BATCHES} = {CFG.BATCH_SIZE * CFG.ACCUMULATE_GRAD_BATCHES}\")\n    print(\"=\"*80 + \"\\n\")\n    \n    trained_models = train_kfold()\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"✅ Models Ready!\")\nprint(f\"Total models: {len(trained_models)}\")\nprint(\"=\"*80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:39.615820Z","iopub.execute_input":"2026-01-03T06:32:39.616181Z","iopub.status.idle":"2026-01-03T06:32:54.286385Z","shell.execute_reply.started":"2026-01-03T06:32:39.616148Z","shell.execute_reply":"2026-01-03T06:32:54.285514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Prepare Test Data\ndef prepare_test_data():\n    \"\"\"Prepare test data with YOUR feature engineering\"\"\"\n    \n    test_df = pd.read_csv(CFG.TEST_CSV)\n    \n    # Get unique images\n    test_images = test_df.drop_duplicates('image_path')[['image_path']].reset_index(drop=True)\n    \n    # Use mean values from training data for tabular features\n    test_images['Pre_GSHH_NDVI'] = float(train_df['Pre_GSHH_NDVI'].mean())\n    test_images['Height_Ave_cm'] = float(train_df['Height_Ave_cm'].mean())\n    test_images['NDVI_Height_interaction'] = float(train_df['NDVI_Height_interaction'].mean())\n    test_images['NDVI_squared'] = float(train_df['NDVI_squared'].mean())\n    test_images['Height_log'] = float(train_df['Height_log'].mean())\n    test_images['State_encoded'] = float(train_df['State_encoded'].mode()[0])\n    test_images['Species_encoded'] = float(train_df['Species_encoded'].mode()[0])\n    test_images['Month'] = float(train_df['Month'].mean())\n    test_images['Season'] = float(train_df['Season'].mode()[0])\n    \n    # Ensure all columns are float32\n    feature_cols = ['Pre_GSHH_NDVI', 'Height_Ave_cm', 'NDVI_Height_interaction',\n                   'NDVI_squared', 'Height_log', 'State_encoded', 'Species_encoded', 'Month', 'Season']\n    for col in feature_cols:\n        test_images[col] = test_images[col].astype(np.float32)\n    \n    return test_images, test_df\n\ntest_images, test_df_original = prepare_test_data()\nprint(f\"✅ Test images: {len(test_images)}\")\nprint(f\"✅ Test samples: {len(test_df_original)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:54.287566Z","iopub.execute_input":"2026-01-03T06:32:54.288001Z","iopub.status.idle":"2026-01-03T06:32:54.304579Z","shell.execute_reply.started":"2026-01-03T06:32:54.287968Z","shell.execute_reply":"2026-01-03T06:32:54.303885Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Inference with TTA - Updated for competitor's model\ndef predict_with_tta(models, test_images, device='cuda'):\n    \"\"\"Generate predictions using ensemble and TTA\"\"\"\n    \n    # Move models to device\n    for model in models:\n        model.to(device)\n        model.eval()\n    \n    all_predictions = []\n    \n    # Get TTA transforms\n    tta_transforms = get_tta_transforms() if CFG.USE_TTA else [get_valid_transforms()]\n    \n    print(f\"\\n🔮 Generating predictions with {len(models)} models and {len(tta_transforms)} TTA steps...\")\n    \n    for tta_idx, transform in enumerate(tta_transforms):\n        print(f\"TTA step {tta_idx + 1}/{len(tta_transforms)}\")\n        \n        # Create dataset with current transform\n        test_dataset = BiomassDatasetDINO(test_images, transform=transform, is_test=True)\n        test_loader = DataLoader(\n            test_dataset,\n            batch_size=CFG.BATCH_SIZE,\n            shuffle=False,\n            num_workers=CFG.NUM_WORKERS,\n            pin_memory=True,\n            persistent_workers=True if CFG.NUM_WORKERS > 0 else False\n        )\n        \n        # Predict with each model\n        tta_predictions = []\n        \n        with torch.no_grad():\n            for model in models:\n                model_preds = []\n                \n                for x_left, x_right, features, _ in tqdm(test_loader, desc=f\"Model prediction\"):\n                    x_left = x_left.to(device)\n                    x_right = x_right.to(device)\n                    # Note: features NOT used - competitor's model doesn't have tabular input\n                    \n                    with torch.amp.autocast('cuda', enabled=CFG.MIXED_PRECISION):\n                        preds = model(x_left, x_right)  # No features parameter!\n                    \n                    model_preds.append(preds.cpu().numpy())\n                \n                model_preds = np.concatenate(model_preds)\n                tta_predictions.append(model_preds)\n        \n        # Average predictions across models for this TTA step\n        tta_avg = np.mean(tta_predictions, axis=0)\n        all_predictions.append(tta_avg)\n    \n    # Average across all TTA steps\n    final_predictions = np.mean(all_predictions, axis=0)\n    \n    # Note: NO inverse transform - competitor's model outputs raw values\n    \n    # COMPETITOR's post-processing\n    green = final_predictions[:, 0]\n    dead = final_predictions[:, 1]\n    clover = final_predictions[:, 2]\n    total = final_predictions[:, 3]\n    gdm = final_predictions[:, 4]\n    \n    # Apply minimum thresholds\n    clover = np.where(clover < CFG.DRY_CLOVER_MIN, 0, clover)\n    dead = np.where(dead < CFG.DRY_DEAD_MIN, 0, dead)\n    \n    # Ensure non-negative\n    green = np.maximum(green, 0)\n    dead = np.maximum(dead, 0)\n    clover = np.maximum(clover, 0)\n    total = np.maximum(total, 0)\n    gdm = np.maximum(gdm, 0)\n    \n    # Clean NaN/Inf\n    def clean(x):\n        return np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n    \n    final_predictions = np.column_stack([\n        clean(green), clean(dead), clean(clover), clean(total), clean(gdm)\n    ])\n    \n    return final_predictions\n\n# Generate predictions\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\npredictions = predict_with_tta(trained_models, test_images, device=device)\n\nprint(f\"\\n✅ Predictions shape: {predictions.shape}\")\nprint(f\"📊 Prediction stats:\")\nprint(pd.DataFrame(predictions, columns=list(CFG.TARGET_COLS)).describe())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:32:54.307786Z","iopub.execute_input":"2026-01-03T06:32:54.308187Z","iopub.status.idle":"2026-01-03T06:33:13.744243Z","shell.execute_reply.started":"2026-01-03T06:32:54.308160Z","shell.execute_reply":"2026-01-03T06:33:13.743381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Create Submission\ndef create_submission(predictions, test_images, test_df):\n    \"\"\"Create submission file in required format\"\"\"\n    \n    # Create prediction dataframe\n    pred_df = pd.DataFrame(predictions, columns=list(CFG.TARGET_COLS))\n    pred_df['image_path'] = test_images['image_path'].values\n    \n    # Melt to long format\n    pred_long = pred_df.melt(\n        id_vars=['image_path'],\n        value_vars=list(CFG.TARGET_COLS),\n        var_name='target_name',\n        value_name='target'\n    )\n    \n    # Create sample_id\n    pred_long['image_id'] = pred_long['image_path'].str.replace('test/', '').str.replace('.jpg', '')\n    pred_long['sample_id'] = pred_long['image_id'] + '__' + pred_long['target_name']\n    \n    # Select and sort columns\n    submission = pred_long[['sample_id', 'target']].copy()\n    \n    # Merge with test_df to ensure correct order\n    submission = test_df[['sample_id']].merge(submission, on='sample_id', how='left')\n    \n    # Save\n    submission.to_csv('submission.csv', index=False)\n    print(f\"\\n✅ Submission file created: submission.csv\")\n    print(f\"Shape: {submission.shape}\")\n    print(f\"\\nFirst few rows:\")\n    print(submission.head(10))\n    \n    return submission\n\nsubmission = create_submission(predictions, test_images, test_df_original)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T06:33:13.745704Z","iopub.execute_input":"2026-01-03T06:33:13.746006Z","iopub.status.idle":"2026-01-03T06:33:13.768600Z","shell.execute_reply.started":"2026-01-03T06:33:13.745974Z","shell.execute_reply":"2026-01-03T06:33:13.767857Z"}},"outputs":[],"execution_count":null}]}