{"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":14091363,"sourceType":"datasetVersion","datasetId":8972760},{"sourceId":293031629,"sourceType":"kernelVersion"},{"sourceId":268942,"sourceType":"modelInstanceVersion","modelInstanceId":230141,"modelId":251887},{"sourceId":711299,"sourceType":"modelInstanceVersion","modelInstanceId":540408,"modelId":553604},{"sourceId":719660,"sourceType":"modelInstanceVersion","modelInstanceId":547330,"modelId":560139}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":441.06159,"end_time":"2026-01-05T05:15:36.197688","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-01-05T05:08:15.136098","version":"2.6.0"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"5cea4344fed44fffaac71b6fe02ac50b":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","background":null,"description_width":"","font_size":null,"text_color":null}},"5fd6570613f943ecb07631ffcecddacb":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"6b204ec19a394a818e17994aae3263e6":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","background":null,"description_width":"","font_size":null,"text_color":null}},"864a5d2d04374e688a854f72ad0d04d9":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"8f988be66966496085d6e6b8909bf0da":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HTMLView","description":"","description_allow_html":false,"layout":"IPY_MODEL_864a5d2d04374e688a854f72ad0d04d9","placeholder":"​","style":"IPY_MODEL_5cea4344fed44fffaac71b6fe02ac50b","tabbable":null,"tooltip":null,"value":" 1/1 [00:01<00:00,  1.98s/it]"}},"9a1a60c1824e4af4a02be7dd08d54449":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"af818733f9ec4e22a3500a5692c607b2":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_e8a7fcf1682642f9a95f077413387b58","IPY_MODEL_d548382ab4904dd9b1c2742bec3463f0","IPY_MODEL_8f988be66966496085d6e6b8909bf0da"],"layout":"IPY_MODEL_5fd6570613f943ecb07631ffcecddacb","tabbable":null,"tooltip":null}},"b695855014a24a7184902662cc419d1e":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"d372dd276f534b6ba203146ce2f78865":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"d548382ab4904dd9b1c2742bec3463f0":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"ProgressView","bar_style":"success","description":"","description_allow_html":false,"layout":"IPY_MODEL_9a1a60c1824e4af4a02be7dd08d54449","max":1,"min":0,"orientation":"horizontal","style":"IPY_MODEL_d372dd276f534b6ba203146ce2f78865","tabbable":null,"tooltip":null,"value":1}},"e8a7fcf1682642f9a95f077413387b58":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HTMLView","description":"","description_allow_html":false,"layout":"IPY_MODEL_b695855014a24a7184902662cc419d1e","placeholder":"​","style":"IPY_MODEL_6b204ec19a394a818e17994aae3263e6","tabbable":null,"tooltip":null,"value":"100%"}}},"version_major":2,"version_minor":0}}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# CSIRO Biomass Prediction - FIXED: Kernel-Safe with Proper Resource Management\n# Prevents multiprocessing errors and OOM issues\n\nimport os\nimport math\nimport warnings\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms as T\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.ensemble import IsolationForest\nfrom sklearn.metrics import r2_score\nimport lightgbm as lgb\nimport optuna\nimport timm\nimport gc\n\nwarnings.filterwarnings('ignore')\n\n# ⭐ FIX 1: Force single-threaded for stability\ntorch.set_num_threads(2)\nos.environ['OMP_NUM_THREADS'] = '2'\nos.environ['MKL_NUM_THREADS'] = '2'\n\ntorch.cuda.empty_cache()\ngc.collect()\n\nprint(\"=\"*70)\nprint(\"🚀 CSIRO BIOMASS - KERNEL-SAFE VERSION\")\nprint(\"=\"*70)\n\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\nclass CFG:\n    DATA_PATH = '/kaggle/input/csiro-biomass/'\n    TEST_DATA_PATH = '/kaggle/input/csiro-biomass/test.csv'\n    OUTPUT_DIR = './outputs'\n    DINOV2_PATH = '/kaggle/input/dinov2-large/dinov2-large'\n    \n    MODELS = {'dinov2_large': 'vit_large_patch14_dinov2.lvd142m'}\n    \n    EPOCHS = 12\n    BATCH_SIZE = 4\n    GRAD_ACC = 8\n    NUM_FOLDS = 3\n    SEED = 42\n    \n    IMG_SCHEDULE = [(0, 336), (7, 392)]\n    USE_GRADIENT_CHECKPOINTING = True\n    \n    LR_BACKBONE = 1e-5\n    LR_HEAD = 1e-3\n    WARMUP_EPOCHS = 2\n    \n    TARGETS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    TARGET_MAX = {\n        'Dry_Green_g': 157.9836,\n        'Dry_Clover_g': 71.7865,\n        'Dry_Dead_g': 83.8407,\n        'GDM_g': 157.9836,\n        'Dry_Total_g': 185.70\n    }\n    \n    LAMBDA_PHYSICS = 0.15\n    OPTUNA_TRIALS = 15\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    # ⭐ FIX 2: Reduce workers to prevent resource exhaustion\n    NUM_WORKERS = 0 if not torch.cuda.is_available() else 2  # ← REDUCED\n    USE_AMP = torch.cuda.is_available()\n    PATIENCE = 3\n    \n    # ⭐ FIX 3: Enable persistent workers only if workers > 0\n    PERSISTENT_WORKERS = False  # ← SAFER\n\nCFG.TARGET_MAX_ARRAY = np.array([CFG.TARGET_MAX[t] for t in CFG.TARGETS])\n\nos.makedirs(CFG.OUTPUT_DIR, exist_ok=True)\nprint(f\"Device: {CFG.DEVICE}\")\nprint(f\"Workers: {CFG.NUM_WORKERS} (persistent: {CFG.PERSISTENT_WORKERS})\")\nprint(\"=\"*70)\n\n# ============================================================================\n# MASS BALANCE (UNCHANGED)\n# ============================================================================\ndef enforce_mass_balance(df_wide, fixed_clover=True):\n    ordered_cols = ['Dry_Green_g', 'Dry_Clover_g', 'Dry_Dead_g', 'GDM_g', 'Dry_Total_g']\n    Y = df_wide[ordered_cols].values.T\n    \n    if fixed_clover:\n        clover_fixed = Y[1, :].copy()\n        Y[3, :] = Y[0, :] + clover_fixed\n        Y[4, :] = Y[3, :] + Y[2, :]\n        Y_reconciled = Y\n    else:\n        C = np.array([\n            [1, 1, 0, -1, 0],\n            [0, 0, 1, 1, -1]\n        ])\n        \n        CT = C.T\n        try:\n            inv_CCt = np.linalg.inv(C @ CT)\n            P = np.eye(5) - CT @ inv_CCt @ C\n        except np.linalg.LinAlgError:\n            print(\"⚠️  Singular matrix, skipping projection\")\n            return df_wide\n        \n        Y_reconciled = P @ Y\n    \n    Y_reconciled = Y_reconciled.T\n    Y_reconciled = np.maximum(0, Y_reconciled)\n    \n    df_out = df_wide.copy()\n    df_out[ordered_cols] = Y_reconciled\n    \n    return df_out\n\n# ============================================================================\n# DATA LOADING (UNCHANGED)\n# ============================================================================\ndef load_and_pivot_data(csv_path):\n    print(f\"Loading data from {csv_path}...\")\n    df_long = pd.read_csv(csv_path)\n    \n    print(f\"Loaded {len(df_long)} rows\")\n    df_long['image_id'] = df_long['sample_id'].str.split('__').str[0]\n    \n    df_wide = df_long.pivot_table(\n        index=['image_id', 'image_path', 'Sampling_Date', 'State', 'Species', \n               'Pre_GSHH_NDVI', 'Height_Ave_cm'],\n        columns='target_name',\n        values='target'\n    ).reset_index()\n    \n    df_wide.columns.name = None\n    df_wide['image_path'] = df_wide['image_path'].apply(\n        lambda x: os.path.join(CFG.DATA_PATH, x)\n    )\n    \n    print(f\"Pivoted to {len(df_wide)} images\")\n    \n    return df_wide\n\ndef load_test_data(csv_path):\n    print(f\"Loading test data from {csv_path}...\")\n    df_long = pd.read_csv(csv_path)\n    \n    print(f\"Loaded {len(df_long)} test rows\")\n    df_long['image_id'] = df_long['sample_id'].str.split('__').str[0]\n    \n    df_wide = df_long.drop_duplicates(subset=['image_id']).reset_index(drop=True)\n    df_wide['image_path'] = df_wide['image_path'].apply(\n        lambda x: os.path.join(CFG.DATA_PATH, x)\n    )\n    \n    print(f\"Test set: {len(df_wide)} images\")\n    \n    return df_wide, df_long[['sample_id', 'target_name']]\n\n# ============================================================================\n# DATASET (UNCHANGED)\n# ============================================================================\nclass BiomassDataset(Dataset):\n    def __init__(self, df, img_size=336, mode='train'):\n        self.df = df.reset_index(drop=True)\n        self.img_size = img_size\n        self.mode = mode\n        \n        if mode == 'train':\n            self.transform = T.Compose([\n                T.Resize((img_size, img_size)),\n                T.RandomResizedCrop(img_size, scale=(0.85, 1.0)),\n                T.RandomHorizontalFlip(p=0.5),\n                T.RandomVerticalFlip(p=0.5),\n                T.RandomRotation(10),\n                T.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.15, hue=0.05),\n                T.RandomGrayscale(p=0.05),\n                T.ToTensor(),\n                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n            ])\n        else:\n            self.transform = T.Compose([\n                T.Resize((img_size, img_size)),\n                T.ToTensor(),\n                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n            ])\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img = Image.open(row['image_path']).convert('RGB')\n        \n        w, h = img.size\n        left = img.crop((0, 0, w//2, h))\n        right = img.crop((w//2, 0, w, h))\n        \n        left = self.transform(left)\n        right = self.transform(right)\n        \n        if self.mode == 'test':\n            return left, right\n        \n        targets = torch.tensor([\n            row['Dry_Green_g'], row['Dry_Dead_g'], row['Dry_Clover_g'],\n            row['GDM_g'], row['Dry_Total_g']\n        ], dtype=torch.float32)\n        \n        targets = targets / torch.tensor(CFG.TARGET_MAX_ARRAY, dtype=torch.float32)\n        \n        return left, right, targets\n\n# ============================================================================\n# HELPER FUNCTIONS (UNCHANGED)\n# ============================================================================\ndef compute_vegetation_indices(images):\n    r, g, b = images[:, 0], images[:, 1], images[:, 2]\n    \n    exg = 2*g - r - b\n    exgr = exg - (1.4*r - g)\n    vari = (g - r) / (g + r + 1e-6)\n    ndi = (g - r) / (g + r + 1e-6)\n    \n    features = torch.stack([\n        exg.mean(dim=(1,2)),\n        exgr.mean(dim=(1,2)),\n        vari.mean(dim=(1,2)),\n        ndi.mean(dim=(1,2)),\n        g.mean(dim=(1,2)),\n        r.std(dim=(1,2))\n    ], dim=1)\n    \n    return features\n\nclass CrossAttentionFusion(nn.Module):\n    def __init__(self, dim=1024, num_heads=4):\n        super().__init__()\n        self.cross_attn_l2r = nn.MultiheadAttention(dim, num_heads, batch_first=True, dropout=0.1)\n        self.cross_attn_r2l = nn.MultiheadAttention(dim, num_heads, batch_first=True, dropout=0.1)\n        self.norm1 = nn.LayerNorm(dim)\n        self.norm2 = nn.LayerNorm(dim)\n        self.mlp = nn.Sequential(\n            nn.Linear(dim, dim * 2),\n            nn.GELU(),\n            nn.Dropout(0.15),\n            nn.Linear(dim * 2, dim),\n            nn.Dropout(0.1)\n        )\n        \n    def forward(self, left_feat, right_feat):\n        left_feat = left_feat.unsqueeze(1)\n        right_feat = right_feat.unsqueeze(1)\n        \n        l_enhanced, _ = self.cross_attn_l2r(left_feat, right_feat, right_feat)\n        r_enhanced, _ = self.cross_attn_r2l(right_feat, left_feat, left_feat)\n        \n        fused = (l_enhanced + r_enhanced) / 2\n        fused = self.norm1(fused)\n        fused = fused + self.mlp(self.norm2(fused))\n        \n        return fused.squeeze(1)\n\nclass UncertaintyLoss(nn.Module):\n    def __init__(self, n_tasks=5):\n        super().__init__()\n        self.log_vars = nn.Parameter(torch.zeros(n_tasks))\n    \n    def forward(self, predictions, targets):\n        losses = []\n        for i in range(predictions.size(1)):\n            task_loss = F.smooth_l1_loss(predictions[:, i], targets[:, i])\n            weighted = torch.exp(-self.log_vars[i]) * task_loss + self.log_vars[i]\n            losses.append(weighted)\n        return torch.stack(losses).mean()\n\ndef physics_constraint_loss(predictions):\n    green, dead, clover, gdm, total = predictions.unbind(1)\n    constraint1 = torch.abs(green + clover - gdm)\n    constraint2 = torch.abs(gdm + dead - total)\n    return (constraint1 + constraint2).mean()\n\n# ============================================================================\n# MODEL (UNCHANGED)\n# ============================================================================\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name='vit_large_patch14_dinov2.lvd142m', pretrained=True):\n        super().__init__()\n        \n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=pretrained,\n            num_classes=0,\n            global_pool='',\n            dynamic_img_size=True,\n            img_size=392\n        )\n        \n        if CFG.USE_GRADIENT_CHECKPOINTING and hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(enable=True)\n        \n        dummy = torch.randn(1, 3, 336, 336)\n        with torch.no_grad():\n            feat = self.backbone(dummy)\n            if feat.dim() == 3:\n                self.feat_dim = feat.size(-1)\n            else:\n                self.feat_dim = feat.size(1)\n        \n        print(f\"📏 Feature dim: {self.feat_dim}\")\n        \n        self.fusion = CrossAttentionFusion(dim=self.feat_dim, num_heads=4)\n        self.veg_indices_proj = nn.Sequential(\n            nn.Linear(6, 128),\n            nn.LayerNorm(128),\n            nn.GELU(),\n            nn.Dropout(0.1),\n            nn.Linear(128, 64)\n        )\n        \n        final_dim = self.feat_dim + 64\n        \n        self.heads = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(final_dim, 512),\n                nn.LayerNorm(512),\n                nn.GELU(),\n                nn.Dropout(0.25),\n                nn.Linear(512, 128),\n                nn.LayerNorm(128),\n                nn.GELU(),\n                nn.Dropout(0.15),\n                nn.Linear(128, 32),\n                nn.GELU(),\n                nn.Linear(32, 1)\n            ) for _ in range(5)\n        ])\n        \n        self._init_heads()\n    \n    def _init_heads(self):\n        for head in self.heads:\n            for m in head.modules():\n                if isinstance(m, nn.Linear):\n                    nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                    if m.bias is not None:\n                        nn.init.constant_(m.bias, 0)\n                elif isinstance(m, nn.LayerNorm):\n                    nn.init.constant_(m.weight, 1)\n                    nn.init.constant_(m.bias, 0)\n    \n    def extract_features(self, x):\n        feat = self.backbone(x)\n        \n        if feat.dim() == 3:\n            cls_token = feat[:, 0, :]\n            mean_pool = feat[:, 1:, :].mean(dim=1)\n            feat = (cls_token + mean_pool) / 2\n        elif feat.dim() == 4:\n            feat = F.adaptive_avg_pool2d(feat, 1).flatten(1)\n        \n        return feat\n    \n    def forward(self, left_img, right_img):\n        left_feat = self.extract_features(left_img)\n        right_feat = self.extract_features(right_img)\n        fused = self.fusion(left_feat, right_feat)\n        \n        veg_left = compute_vegetation_indices(left_img)\n        veg_right = compute_vegetation_indices(right_img)\n        veg_feat = self.veg_indices_proj((veg_left + veg_right) / 2)\n        \n        combined = torch.cat([fused, veg_feat], dim=1)\n        outputs = torch.cat([head(combined) for head in self.heads], dim=1)\n        \n        outputs = F.softplus(outputs)\n        outputs = torch.clamp(outputs, 0, 2.0)\n        \n        return outputs\n\n# ============================================================================\n# ⭐ FIX 4: SAFE DATALOADER CREATION WITH CLEANUP\n# ============================================================================\ndef create_dataloader(dataset, batch_size, shuffle=False, drop_last=False):\n    \"\"\"Create DataLoader with proper resource management\"\"\"\n    loader_kwargs = {\n        'batch_size': batch_size,\n        'shuffle': shuffle,\n        'num_workers': CFG.NUM_WORKERS,\n        'pin_memory': True if CFG.DEVICE.type == 'cuda' else False,\n        'drop_last': drop_last,\n    }\n    \n    # Only use persistent workers if we have workers\n    if CFG.NUM_WORKERS > 0 and CFG.PERSISTENT_WORKERS:\n        loader_kwargs['persistent_workers'] = True\n    \n    return DataLoader(dataset, **loader_kwargs)\n\ndef cleanup_dataloader(loader):\n    \"\"\"Safely cleanup DataLoader resources\"\"\"\n    if loader is not None and hasattr(loader, '_iterator'):\n        try:\n            if loader._iterator is not None:\n                loader._iterator._shutdown_workers()\n        except:\n            pass\n    \n    del loader\n    gc.collect()\n\n# ============================================================================\n# TRAINING (MODIFIED FOR SAFETY)\n# ============================================================================\ndef train_epoch(model, loader, optimizer, criterion_uncertainty, epoch, scaler=None):\n    model.train()\n    total_loss = 0\n    total_physics = 0\n    \n    pbar = tqdm(loader, desc=f'Epoch {epoch+1}')\n    \n    for batch_idx, (left, right, targets) in enumerate(pbar):\n        left = left.to(CFG.DEVICE, non_blocking=True)\n        right = right.to(CFG.DEVICE, non_blocking=True)\n        targets = targets.to(CFG.DEVICE, non_blocking=True)\n        \n        with torch.amp.autocast('cuda', enabled=CFG.USE_AMP):\n            predictions = model(left, right)\n            task_loss = criterion_uncertainty(predictions, targets)\n            physics_loss = physics_constraint_loss(predictions)\n            loss = task_loss + CFG.LAMBDA_PHYSICS * physics_loss\n            loss = loss / CFG.GRAD_ACC\n        \n        if scaler is not None:\n            scaler.scale(loss).backward()\n            \n            if (batch_idx + 1) % CFG.GRAD_ACC == 0:\n                scaler.unscale_(optimizer)\n                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad()\n        else:\n            loss.backward()\n            \n            if (batch_idx + 1) % CFG.GRAD_ACC == 0:\n                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n                optimizer.step()\n                optimizer.zero_grad()\n        \n        total_loss += task_loss.item()\n        total_physics += physics_loss.item()\n        \n        pbar.set_postfix({'loss': f'{task_loss.item():.4f}', 'physics': f'{physics_loss.item():.4f}'})\n    \n    return total_loss / len(loader), total_physics / len(loader)\n\n@torch.inference_mode()\ndef validate(model, loader):\n    model.eval()\n    preds_list = []\n    targets_list = []\n    \n    norm_tensor = torch.tensor(CFG.TARGET_MAX_ARRAY, dtype=torch.float32, device=CFG.DEVICE)\n    \n    for left, right, targets in tqdm(loader, desc='Validation', leave=False):\n        left = left.to(CFG.DEVICE, non_blocking=True)\n        right = right.to(CFG.DEVICE, non_blocking=True)\n        targets = targets.to(CFG.DEVICE, non_blocking=True)\n        \n        with torch.amp.autocast('cuda', enabled=CFG.USE_AMP):\n            predictions = model(left, right)\n        \n        predictions = predictions * norm_tensor\n        targets = targets * norm_tensor\n        \n        preds_list.append(predictions.cpu().numpy())\n        targets_list.append(targets.cpu().numpy())\n    \n    preds = np.vstack(preds_list)\n    targets = np.vstack(targets_list)\n    \n    r2_scores = [r2_score(targets[:, i], preds[:, i]) for i in range(5)]\n    avg_r2 = np.mean(r2_scores)\n    \n    return avg_r2, r2_scores, preds, targets\n\n@torch.inference_mode()\ndef predict_test(model, loader):\n    model.eval()\n    preds_list = []\n    \n    norm_tensor = torch.tensor(CFG.TARGET_MAX_ARRAY, dtype=torch.float32, device=CFG.DEVICE)\n    \n    for left, right in tqdm(loader, desc='Test Prediction'):\n        left = left.to(CFG.DEVICE, non_blocking=True)\n        right = right.to(CFG.DEVICE, non_blocking=True)\n        \n        with torch.amp.autocast('cuda', enabled=CFG.USE_AMP):\n            predictions = model(left, right)\n        \n        predictions = predictions * norm_tensor\n        \n        preds_list.append(predictions.cpu().numpy())\n    \n    return np.vstack(preds_list)\n\n# ============================================================================\n# META-MODEL (UNCHANGED)\n# ============================================================================\ndef train_meta_model(base_predictions, targets):\n    X_meta = np.hstack(base_predictions)\n    \n    print(\"\\n🔧 Tuning meta-model with Optuna...\")\n    \n    def objective(trial):\n        params = {\n            'n_estimators': trial.suggest_int('n_estimators', 150, 600),\n            'learning_rate': trial.suggest_float('learning_rate', 0.02, 0.08, log=True),\n            'max_depth': trial.suggest_int('max_depth', 3, 7),\n            'num_leaves': trial.suggest_int('num_leaves', 30, 100),\n            'subsample': trial.suggest_float('subsample', 0.7, 0.95),\n            'colsample_bytree': trial.suggest_float('colsample_bytree', 0.7, 0.95),\n            'reg_alpha': trial.suggest_float('reg_alpha', 0.0, 2.0),\n            'reg_lambda': trial.suggest_float('reg_lambda', 0.5, 4.0),\n            'min_child_samples': trial.suggest_int('min_child_samples', 10, 50)\n        }\n        \n        scores = []\n        for i in range(5):\n            model = lgb.LGBMRegressor(**params, random_state=CFG.SEED, verbose=-1, n_jobs=-1)\n            model.fit(X_meta, targets[:, i])\n            \n            from sklearn.model_selection import cross_val_score\n            cv_scores = cross_val_score(model, X_meta, targets[:, i], cv=3, scoring='r2', n_jobs=-1)\n            scores.append(cv_scores.mean())\n        \n        return np.mean(scores)\n    \n    study = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=CFG.SEED))\n    study.optimize(objective, n_trials=CFG.OPTUNA_TRIALS, show_progress_bar=True)\n    \n    print(f\"✅ Best meta-model R²: {study.best_value:.4f}\")\n    \n    best_params = study.best_params\n    meta_models = []\n    \n    for i in range(5):\n        model = lgb.LGBMRegressor(**best_params, random_state=CFG.SEED, verbose=-1, n_jobs=-1)\n        model.fit(X_meta, targets[:, i])\n        meta_models.append(model)\n    \n    return meta_models\n\n# ============================================================================\n# ⭐ FIX 5: MAIN PIPELINE WITH PROPER CLEANUP\n# ============================================================================\ndef main():\n    import time\n    start_time = time.time()\n    \n    print(\"\\n📊 Loading data...\")\n    train_df = load_and_pivot_data(os.path.join(CFG.DATA_PATH, 'train.csv'))\n    \n    print(\"\\n🔍 Removing outliers...\")\n    iso_forest = IsolationForest(contamination=0.05, random_state=CFG.SEED)\n    outliers = iso_forest.fit_predict(train_df[CFG.TARGETS].values)\n    train_df = train_df[outliers == 1].reset_index(drop=True)\n    print(f\"Final dataset: {len(train_df)} images\")\n    \n    bins = pd.qcut(train_df['Dry_Total_g'], q=10, labels=False, duplicates='drop')\n    skf = StratifiedGroupKFold(n_splits=CFG.NUM_FOLDS, shuffle=True, random_state=CFG.SEED)\n    \n    all_oof_preds = []\n    all_oof_targets = []\n    fold_models = []\n    \n    for fold, (train_idx, val_idx) in enumerate(skf.split(train_df, bins, groups=train_df['image_id'])):\n        print(f\"\\n{'='*70}\")\n        print(f\"🔥 FOLD {fold+1}/{CFG.NUM_FOLDS}\")\n        print(f\"{'='*70}\")\n        \n        # ⭐ AGGRESSIVE CLEANUP\n        torch.cuda.empty_cache()\n        gc.collect()\n        \n        train_fold = train_df.iloc[train_idx]\n        val_fold = train_df.iloc[val_idx]\n        \n        best_r2 = -999\n        patience_counter = 0\n        \n        # Variables to track for cleanup\n        train_loader = None\n        val_loader = None\n        \n        try:\n            for stage_idx, (start_epoch, img_size) in enumerate(CFG.IMG_SCHEDULE):\n                if stage_idx == 0:\n                    end_epoch = CFG.IMG_SCHEDULE[1][0]\n                else:\n                    end_epoch = CFG.EPOCHS\n                \n                print(f\"\\n📐 Resolution: {img_size}×{img_size} (Epochs {start_epoch+1}-{end_epoch})\")\n                \n                train_ds = BiomassDataset(train_fold, img_size=img_size, mode='train')\n                val_ds = BiomassDataset(val_fold, img_size=img_size, mode='val')\n                \n                # ⭐ SAFE DATALOADER CREATION\n                train_loader = create_dataloader(train_ds, CFG.BATCH_SIZE, shuffle=True, drop_last=True)\n                val_loader = create_dataloader(val_ds, CFG.BATCH_SIZE*2, shuffle=False, drop_last=False)\n                \n                if stage_idx == 0:\n                    print(f\"📦 Loading DINOv2...\")\n                    \n                    model = BiomassModel(model_name=CFG.MODELS['dinov2_large'], pretrained=False)\n                    \n                    try:\n                        safetensors_path = os.path.join(CFG.DINOV2_PATH, 'model.safetensors')\n                        pytorch_path = os.path.join(CFG.DINOV2_PATH, 'pytorch_model.bin')\n                        \n                        if os.path.exists(safetensors_path):\n                            from safetensors.torch import load_file\n                            state_dict = load_file(safetensors_path)\n                            print(\"✅ Loaded from safetensors\")\n                        elif os.path.exists(pytorch_path):\n                            state_dict = torch.load(pytorch_path, map_location='cpu', weights_only=True)\n                            print(\"✅ Loaded from pytorch_model.bin\")\n                        else:\n                            raise FileNotFoundError(f\"No weights found in {CFG.DINOV2_PATH}\")\n                        \n                        missing, unexpected = model.backbone.load_state_dict(state_dict, strict=False)\n                        print(f\"✅ Loaded (missing: {len(missing)}, unexpected: {len(unexpected)})\")\n                    \n                    except Exception as e:\n                        print(f\"⚠️  Using RANDOM init: {e}\")\n                    \n                    model = model.to(CFG.DEVICE)\n                    \n                    optimizer = torch.optim.AdamW([\n                        {'params': model.backbone.parameters(), 'lr': CFG.LR_BACKBONE},\n                        {'params': [p for n, p in model.named_parameters() if 'backbone' not in n], \n                         'lr': CFG.LR_HEAD}\n                    ], weight_decay=0.01)\n                    \n                    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n                        optimizer, T_0=5, T_mult=2, eta_min=1e-6\n                    )\n                    \n                    criterion_uncertainty = UncertaintyLoss(n_tasks=5).to(CFG.DEVICE)\n                    scaler = torch.amp.GradScaler('cuda') if CFG.USE_AMP else None\n                \n                for epoch in range(start_epoch, end_epoch):\n                    train_loss, physics_loss = train_epoch(\n                        model, train_loader, optimizer, criterion_uncertainty, epoch, scaler\n                    )\n                    \n                    val_r2, r2_scores, val_preds, val_targets = validate(model, val_loader)\n                    scheduler.step()\n                    \n                    print(f\"Epoch {epoch+1}: Train={train_loss:.4f}, Val R²={val_r2:.4f}\")\n                    print(f\"Per-target: {[f'{s:.3f}' for s in r2_scores]}\")\n                    \n                    if val_r2 > best_r2:\n                        best_r2 = val_r2\n                        patience_counter = 0\n                        torch.save(model.state_dict(), f'{CFG.OUTPUT_DIR}/best_fold{fold}.pth')\n                        print(f\"💾 Saved (R²={val_r2:.4f})\")\n                    else:\n                        patience_counter += 1\n                        if patience_counter >= CFG.PATIENCE and epoch > CFG.WARMUP_EPOCHS:\n                            print(f\"⚠️  Early stop at epoch {epoch+1}\")\n                            break\n                \n                # ⭐ CLEANUP DATALOADERS AFTER EACH STAGE\n                cleanup_dataloader(train_loader)\n                cleanup_dataloader(val_loader)\n                train_loader = None\n                val_loader = None\n                \n                torch.cuda.empty_cache()\n                gc.collect()\n            \n            # Final validation\n            model.load_state_dict(torch.load(f'{CFG.OUTPUT_DIR}/best_fold{fold}.pth'))\n            val_ds = BiomassDataset(val_fold, img_size=CFG.IMG_SCHEDULE[-1][1], mode='val')\n            val_loader = create_dataloader(val_ds, CFG.BATCH_SIZE*2, shuffle=False)\n            \n            _, _, oof_preds, oof_targets = validate(model, val_loader)\n            \n            all_oof_preds.append(oof_preds)\n            all_oof_targets.append(oof_targets)\n            fold_models.append(model)\n            \n        finally:\n            # ⭐ ENSURE CLEANUP EVEN IF ERROR\n            cleanup_dataloader(train_loader)\n            cleanup_dataloader(val_loader)\n            torch.cuda.empty_cache()\n            gc.collect()\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"🎯 Training Meta-Model\")\n    print(\"=\"*70)\n    \n    oof_preds_combined = np.vstack(all_oof_preds)\n    oof_targets_combined = np.vstack(all_oof_targets)\n    \n    meta_models = train_meta_model([oof_preds_combined], oof_targets_combined)\n    \n    meta_preds = np.column_stack([m.predict(oof_preds_combined) for m in meta_models])\n    final_r2 = r2_score(oof_targets_combined, meta_preds)\n    \n    total_time = (time.time() - start_time) / 3600\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"✅ TRAINING COMPLETE\")\n    print(\"=\"*70)\n    print(f\"🏆 Final R²: {final_r2:.4f}\")\n    print(f\"⏱️  Time: {total_time:.2f} hours\")\n    \n    return fold_models, meta_models, oof_preds_combined, oof_targets_combined\n\n# ============================================================================\n# ⭐ FIX 6: SAFE SUBMISSION GENERATION\n# ============================================================================\ndef generate_submission(fold_models, meta_models):\n    print(\"\\n\" + \"=\"*70)\n    print(\"🎯 GENERATING SUBMISSION\")\n    print(\"=\"*70)\n    \n    test_df, sample_mapping = load_test_data(CFG.TEST_DATA_PATH)\n    \n    final_img_size = CFG.IMG_SCHEDULE[-1][1]\n    test_ds = BiomassDataset(test_df, img_size=final_img_size, mode='test')\n    \n    all_fold_preds = []\n    \n    for fold_idx, model in enumerate(fold_models):\n        print(f\"\\n📊 Fold {fold_idx+1} prediction...\")\n        \n        # ⭐ RECREATE LOADER FOR EACH FOLD (prevents resource leakage)\n        test_loader = create_dataloader(test_ds, CFG.BATCH_SIZE*2, shuffle=False)\n        \n        try:\n            model.eval()\n            fold_preds = predict_test(model, test_loader)\n            all_fold_preds.append(fold_preds)\n        finally:\n            cleanup_dataloader(test_loader)\n            torch.cuda.empty_cache()\n            gc.collect()\n    \n    base_test_preds = np.mean(all_fold_preds, axis=0)\n    print(f\"✅ Base predictions: {base_test_preds.shape}\")\n    \n    print(\"\\n🔮 Applying meta-model...\")\n    final_preds = np.column_stack([m.predict(base_test_preds) for m in meta_models])\n    \n    print(\"\\n⚖️  Enforcing mass balance...\")\n    preds_df = pd.DataFrame(final_preds, columns=CFG.TARGETS)\n    preds_df['image_id'] = test_df['image_id'].values\n    \n    preds_balanced = enforce_mass_balance(preds_df, fixed_clover=True)\n    \n    print(\"✅ Mass balance applied!\")\n    \n    print(\"\\n📝 Creating submission.csv...\")\n    submission_rows = []\n    for idx, row in test_df.iterrows():\n        image_id = row['image_id']\n        for target_name in CFG.TARGETS:\n            sample_id = f\"{image_id}__{target_name}\"\n            prediction = preds_balanced.iloc[idx][target_name]\n            submission_rows.append({'sample_id': sample_id, 'target': prediction})\n    \n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv('submission.csv', index=False)\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"✅ SUBMISSION COMPLETE\")\n    print(\"=\"*70)\n    print(f\"📄 File: submission.csv\")\n    print(f\"📊 Rows: {len(submission_df)}\")\n    print(\"\\nFirst 10:\")\n    print(submission_df.head(10))\n    \n    print(\"\\n🔬 Mass Balance Verification:\")\n    for idx in range(min(5, len(preds_balanced))):\n        green = preds_balanced.iloc[idx]['Dry_Green_g']\n        clover = preds_balanced.iloc[idx]['Dry_Clover_g']\n        dead = preds_balanced.iloc[idx]['Dry_Dead_g']\n        gdm = preds_balanced.iloc[idx]['GDM_g']\n        total = preds_balanced.iloc[idx]['Dry_Total_g']\n        \n        err1 = abs((green + clover) - gdm)\n        err2 = abs((gdm + dead) - total)\n        \n        print(f\"Sample {idx}: GDM err={err1:.4f}, Total err={err2:.4f}\")\n    \n    return submission_df\n\nif __name__ == '__main__':\n    fold_models, meta_models, oof_preds, oof_targets = main()\n    submission = generate_submission(fold_models, meta_models)\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"🎉 PIPELINE COMPLETE!\")\n    print(\"=\"*70)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}