{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":691204,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":524045,"modelId":538058},{"sourceId":692402,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":525000,"modelId":539041},{"sourceId":692741,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":525255,"modelId":539300},{"sourceId":693373,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":525785,"modelId":539847},{"sourceId":693411,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":525820,"modelId":539881}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport gc\nimport warnings\nwarnings.filterwarnings('ignore')\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nimport timm\nfrom torch.amp import autocast\n\n# --- CONFIG ---\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nBATCH_SIZE = 1  # Reduced batch size to fit both models/TTA in memory if needed\nNUM_TARGETS = 5\n\n# PATHS (Adjust these to match your Kaggle Inputs)\nTEST_CSV_PATH = \"/kaggle/input/csiro-biomass/test.csv\"\nTEST_IMG_DIR = \"/kaggle/input/csiro-biomass\"\n\n# DINO CONFIG\nDINO_MODEL_NAME = \"vit_base_patch14_dinov2.lvd142m\"\nDINO_WEIGHTS_DIR = \"/kaggle/input/ver-09/pytorch/default/1\" # <--- UPDATE THIS\n\n# CONVNEXT CONFIG\nCONV_MODEL_NAME = \"convnext_base.fb_in22k_ft_in1k_384\" # or 'nano' depending on what you trained\nCONV_WEIGHTS_DIR = \"/kaggle/input/ver09-convnext/pytorch/default/1\" # <--- UPDATE THIS\n\n# ENSEMBLE WEIGHTS\nW_DINO = 0.0\nW_CONV = 0.45","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T12:30:40.167056Z","iopub.execute_input":"2025-12-21T12:30:40.167218Z","iopub.status.idle":"2025-12-21T12:30:52.111463Z","shell.execute_reply.started":"2025-12-21T12:30:40.167202Z","shell.execute_reply":"2025-12-21T12:30:52.110864Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------------------------------------\n# 1. SHARED UTILS (GeM, Residual, Dataset)\n# ---------------------------------------------------------\nclass GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super(GeM, self).__init__()\n        self.p = nn.Parameter(torch.ones(1)*p)\n        self.eps = eps\n    def forward(self, x):\n        return nn.functional.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, dim, dropout=0.1):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.LayerNorm(dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim, dim),\n            nn.Dropout(dropout)\n        )\n        self.norm = nn.LayerNorm(dim)\n    def forward(self, x):\n        return self.norm(x + self.block(x))\n\nclass TestDataset(Dataset):\n    def __init__(self, img_ids, img_dir, transforms):\n        self.img_ids = img_ids\n        self.img_dir = img_dir\n        self.tfms = transforms\n\n    def __len__(self):\n        return len(self.img_ids)\n\n    def __getitem__(self, idx):\n        img_id = self.img_ids[idx]\n        # Handle cases where ID might already include path\n        if \"/\" in str(img_id):\n             path = f\"{self.img_dir}/{img_id}\"\n        else:\n             path = f\"{self.img_dir}/{img_id}.jpg\" # Check extension!\n\n        # Try/Except for robust loading\n        if os.path.exists(path):\n            try:\n                img = Image.open(path).convert(\"RGB\")\n            except:\n                img = Image.new(\"RGB\", (518, 518))\n        else:\n             # Fallback search if extension is wrong\n             path_png = path.replace(\".jpg\", \".png\")\n             if os.path.exists(path_png):\n                 img = Image.open(path_png).convert(\"RGB\")\n             else:\n                 img = Image.new(\"RGB\", (518, 518))\n\n        # 4x TTA\n        imgs = [\n            self.tfms(img),\n            self.tfms(img.transpose(Image.FLIP_LEFT_RIGHT)),\n            self.tfms(img.transpose(Image.FLIP_TOP_BOTTOM)),\n            self.tfms(img.rotate(90)),\n        ]\n        return torch.stack(imgs), img_id\n\n# ---------------------------------------------------------\n# 2. MODEL DEFINITIONS\n# ---------------------------------------------------------\n\n# --- A. DINOv2 Architecture ---\n# --- UPDATE CONFIG ---\n# We switch to the 'reg4' version to match your checkpoint weights\nDINO_MODEL_NAME = \"vit_base_patch14_reg4_dinov2.lvd142m\" \n\n# --- CORRECTED DINO MODEL CLASS ---\nclass BiomassModelDINO(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n        # 1. Load the REGISTER version of DINOv2\n        self.backbone = timm.create_model(\n            DINO_MODEL_NAME, \n            pretrained=False, \n            num_classes=0, \n            global_pool=\"\",\n            img_size=518  # CRITICAL: Force 518 to match your checkpoint's pos_embed\n        )\n        \n        embed_dim = self.backbone.num_features \n        self.gem = GeM()\n\n        # Stream 1: Texture Projector\n        self.gem_proj = nn.Sequential(\n            nn.Linear(embed_dim, 512),\n            nn.LayerNorm(512),\n            nn.GELU()\n        )\n\n        # Stream 2: Context Projector\n        self.cls_proj = nn.Sequential(\n            nn.Linear(embed_dim, 512),\n            nn.LayerNorm(512),\n            nn.GELU()\n        )\n\n        # Fusion Head\n        self.regressor = nn.Sequential(\n            nn.Linear(1024, 512),\n            nn.GELU(),\n            ResidualBlock(512, dropout=0.1), \n            nn.Linear(512, 256),\n            nn.GELU(),\n            nn.Linear(256, NUM_TARGETS)\n        )\n\n    def forward(self, x):\n        # 1. Handle TTA Shape\n        b, tta, c, h, w = x.shape\n        x = x.view(-1, c, h, w)\n        \n        # 2. Extract Features\n        x_feat = self.backbone.forward_features(x)\n        \n        # 3. Robust Token Separation (Handle Registers)\n        # The 'reg4' model has 5 prefix tokens (1 CLS + 4 Registers)\n        n_prefix = self.backbone.num_prefix_tokens \n        \n        cls_token = x_feat[:, 0]\n        # Skip CLS (idx 0) AND Registers (idx 1-4) -> Start from n_prefix\n        patch_tokens = x_feat[:, n_prefix:] \n        \n        # 4. GeM Pooling\n        B_sq, N, C_emb = patch_tokens.shape\n        H_grid = W_grid = int(N**0.5) # Should be 37 for 518px input\n        \n        gem_feat = patch_tokens.permute(0, 2, 1).reshape(B_sq, C_emb, H_grid, W_grid)\n        gem_feat = self.gem(gem_feat).flatten(1) \n        \n        # 5. Dual Projection & Fusion\n        combined = torch.cat([self.gem_proj(gem_feat), self.cls_proj(cls_token)], dim=1)\n        \n        preds = self.regressor(combined)\n        return preds.view(b, tta, NUM_TARGETS)\n\n# --- B. ConvNeXt Architecture ---\nclass BiomassModelConvNeXt(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        # 1. Backbone: ConvNeXt (Match your training config)\n        self.backbone = timm.create_model(\n            CONV_MODEL_NAME,    # e.g. \"convnext_base.fb_in22k_ft_in1k_384\"\n            pretrained=False,   # Offline safe\n            num_classes=0,\n            global_pool=\"\" \n        )\n        \n        self.embed_dim = self.backbone.num_features \n        \n        # 2. Pooling Mechanisms (Must match training!)\n        self.gem = GeM() \n        self.avg_pool = nn.AdaptiveAvgPool2d(1) \n\n        # 3. Projections\n        # Branch 1: GeM Features\n        self.gem_proj = nn.Sequential(\n            nn.Linear(self.embed_dim, 512),\n            nn.LayerNorm(512),\n            nn.GELU()\n        )\n\n        # Branch 2: Context Features (Avg Pool)\n        # CRITICAL: This must be named 'ctx_proj' to match your weights\n        self.ctx_proj = nn.Sequential(\n            nn.Linear(self.embed_dim, 512),\n            nn.LayerNorm(512),\n            nn.GELU()\n        )\n\n        # 4. Regressor Head\n        self.regressor = nn.Sequential(\n            nn.Linear(1024, 512),\n            nn.GELU(),\n            ResidualBlock(512, dropout=0.1), # Lower dropout for inference\n            nn.Linear(512, 256),\n            nn.GELU(),\n            nn.Linear(256, NUM_TARGETS)\n        )\n\n    def forward(self, x):\n        # 1. Handle TTA Shape: [Batch, TTA, C, H, W] -> [Batch*TTA, C, H, W]\n        b, tta, c, h, w = x.shape\n        x = x.view(-1, c, h, w)\n\n        # 2. Extract Features\n        # Returns: [B*TTA, Channels, Height, Width]\n        x_feat = self.backbone.forward_features(x)\n        \n        # 3. Branch 1: Texture/Salient Focus (GeM)\n        gem_feat = self.gem(x_feat).flatten(1) \n        gem_emb = self.gem_proj(gem_feat) \n        \n        # 4. Branch 2: Context Focus (Global Average)\n        ctx_feat = self.avg_pool(x_feat).flatten(1) \n        ctx_emb = self.ctx_proj(ctx_feat) \n        \n        # 5. Concatenate & Regress\n        combined = torch.cat([gem_emb, ctx_emb], dim=1) \n        \n        preds = self.regressor(combined)\n\n        # 6. Reshape back for TTA averaging\n        return preds.view(b, tta, NUM_TARGETS)\n\n# ---------------------------------------------------------\n# 3. PREPARE DATA\n# ---------------------------------------------------------\n# Use DINO transforms as base (usually works fine for ConvNeXt too if resolution matches)\ntmp_model = timm.create_model(DINO_MODEL_NAME, pretrained=False)\ndata_config = timm.data.resolve_model_data_config(tmp_model)\nbase_tfms = timm.data.create_transform(**data_config, is_training=False)\ndel tmp_model\ngc.collect()\n\ntest_df = pd.read_csv(TEST_CSV_PATH)\n# Extract clean image IDs (removing the target suffix if present)\ncol = \"sample_id\" if \"sample_id\" in test_df.columns else test_df.columns[0]\nimg_ids = test_df[col].apply(lambda x: str(x).split(\"__\")[0]).unique()\n# If image_path column exists, use that to help Dataset find files\nif \"image_path\" in test_df.columns:\n    # Quick map: ID -> Path\n    path_map = dict(zip(test_df[col].apply(lambda x: str(x).split(\"__\")[0]), test_df[\"image_path\"]))\n    # Pass this map or just rely on IDs if paths are predictable\n\nprint(f\"Found {len(img_ids)} unique images to process.\")\n\ndataset = TestDataset(img_ids, TEST_IMG_DIR, base_tfms)\nloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n\n# ---------------------------------------------------------\n# 4. INFERENCE LOOP\n# ---------------------------------------------------------\ndef get_preds(ModelClass, weights_dir, model_name_log):\n    print(f\"\\n--- Running Inference: {model_name_log} ---\")\n    preds_accum = np.zeros((len(img_ids), NUM_TARGETS), dtype=np.float32)\n    \n    weight_files = sorted([f for f in os.listdir(weights_dir) if f.endswith(\".pth\")])\n    if not weight_files:\n        print(f\"WARNING: No weights found in {weights_dir}!\")\n        return preds_accum # Return zeros if failed\n\n    for i, wf in enumerate(weight_files):\n        print(f\"Loading Fold {i+1}: {wf}\")\n        model = ModelClass().to(DEVICE)\n        state = torch.load(os.path.join(weights_dir, wf), map_location=DEVICE)\n        model.load_state_dict(state, strict=True)\n        model.eval()\n\n        fold_preds = []\n        with torch.no_grad():\n            for imgs, _ in tqdm(loader, leave=False):\n                imgs = imgs.to(DEVICE)\n                with autocast('cuda'):\n                    out = model(imgs) # [B, TTA, 5]\n                \n                # Inverse Log1p -> Expm1\n                out = torch.expm1(out)\n                out = torch.nan_to_num(out, 0.0, 0.0, 0.0)\n                \n                # Average TTA\n                out = out.mean(dim=1) \n                fold_preds.append(out.float().cpu().numpy())\n\n        preds_accum += np.vstack(fold_preds)\n        \n        del model, state\n        torch.cuda.empty_cache()\n        gc.collect()\n        \n    # Average across folds\n    return np.maximum(preds_accum / len(weight_files), 0)\n\n# --- EXECUTE ---\n# preds_dino = get_preds(BiomassModelDINO, DINO_WEIGHTS_DIR, \"DINOv2\")\npreds_conv = get_preds(BiomassModelConvNeXt, CONV_WEIGHTS_DIR, \"ConvNeXt\")\n\n# --- ENSEMBLE ---\nprint(f\"\\nEnsembling: {W_DINO} * DINO + {W_CONV} * ConvNeXt\")\nfinal_preds = preds_conv\n\n# ---------------------------------------------------------\n# 5. SUBMISSION GENERATION\n# ---------------------------------------------------------\ntarget_map = {\n    \"Dry_Green_g\": 0, \"Dry_Dead_g\": 1, \"Dry_Clover_g\": 2, \"GDM_g\": 3, \"Dry_Total_g\": 4\n}\npred_dict = {img_ids[i]: final_preds[i] for i in range(len(img_ids))}\n\nrows = []\nfor _, r in test_df.iterrows():\n    # Parse ID: 'ID123__Dry_Green_g' -> 'ID123', 'Dry_Green_g'\n    parts = r[\"sample_id\"].split(\"__\")\n    img_id_clean = parts[0]\n    tgt_name = parts[1]\n    \n    rows.append({\n        \"sample_id\": r[\"sample_id\"],\n        \"target\": float(pred_dict.get(img_id_clean, [0]*5)[target_map[tgt_name]])\n    })\n\nsubmission = pd.DataFrame(rows)\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(\"Saved submission.csv\")\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T12:30:52.114840Z","iopub.execute_input":"2025-12-21T12:30:52.115073Z","iopub.status.idle":"2025-12-21T12:31:16.476344Z","shell.execute_reply.started":"2025-12-21T12:30:52.115055Z","shell.execute_reply":"2025-12-21T12:31:16.475712Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}