{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":691204,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":524045,"modelId":538058},{"sourceId":694054,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":526320,"modelId":540370},{"sourceId":694561,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":526712,"modelId":540755},{"sourceId":694566,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":526717,"modelId":540760},{"sourceId":694781,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":526879,"modelId":540929},{"sourceId":694830,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":526916,"modelId":540966},{"sourceId":695249,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":527277,"modelId":541324},{"sourceId":695268,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":527292,"modelId":541339},{"sourceId":696121,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":527976,"modelId":542033},{"sourceId":696150,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":527997,"modelId":542053},{"sourceId":696382,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":528178,"modelId":542225}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\nimport timm\nimport os\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport warnings\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nfrom PIL import Image\nfrom torch.amp import autocast\nimport gc\n\n# Filter warnings\nwarnings.filterwarnings('ignore')\n\n# ====================================================\n# 1. CONFIGURATION\n# ====================================================\nclass Config:\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    num_workers = 0 # Set to 2 for faster loading if compatible, else 0\n    batch_size = 2\n\n    test_csv = \"/kaggle/input/csiro-biomass/test.csv\"\n    test_img_dir = \"/kaggle/input/csiro-biomass/test\"\n\n    # Your separate seed directories\n    seed_dirs = [\n        # \"/kaggle/input/latest-1337/pytorch/default/1\",\n        # \"/kaggle/input/latest-42/pytorch/default/1\",\n        \"/kaggle/input/latest-108/pytorch/default/1\",\n        \"/kaggle/input/ver5/pytorch/default/1\"\n    ]\n\n    model_name = \"vit_base_patch14_reg4_dinov2.lvd142m\"\n    num_targets = 5\n    # img_size = 224\n\ncfg = Config()\n\n# ====================================================\n# 2. DATASET & TRANSFORMS\n# ====================================================\nclass TestDataset(Dataset):\n    def __init__(self, img_ids, img_dir, tfms):\n        self.img_ids = img_ids\n        self.img_dir = img_dir\n        self.tfms = tfms\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        \n        # Robust path handling\n        if \"/\" in img_id:\n            path = f\"{self.img_dir}/{img_id}\"\n        else:\n            path = f\"{self.img_dir}/{img_id}.jpg\"\n\n        # Fallback for png\n        if not os.path.exists(path):\n            path = path.replace(\".jpg\", \".png\")\n\n        try:\n            img = Image.open(path).convert(\"RGB\")\n        except:\n            # Return black image if corrupted/missing\n            img = Image.new(\"RGB\", (cfg.img_size, cfg.img_size))\n\n        # 4x Test Time Augmentation (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            self.tfms(img.rotate(180)) \n        ]\n\n        return torch.stack(imgs), img_id\n\n# ====================================================\n# 3. MODEL\n# ====================================================\nclass BioEstimatorD(nn.Module):\n    def __init__(self, name):\n        super().__init__()\n        self.backbone = timm.create_model(\n            name,\n            pretrained=False,\n            num_classes=0,\n            global_pool=\"\"\n        )\n        dim = self.backbone.num_features\n        self.head = nn.Sequential(\n            nn.Linear(dim, 512),\n            nn.ReLU(inplace=True),\n            nn.Dropout(0.1),\n            nn.Linear(512, 256),\n            nn.ReLU(inplace=True),\n            nn.Linear(256, 5)\n        )\n\n    def forward(self, x):\n        # Handle 5D input [Batch, TTA, C, H, W]\n        if x.ndim == 5:\n            b, tta, c, h, w = x.shape\n            x = x.view(-1, c, h, w)\n            tta_mode = True\n        else:\n            tta_mode = False\n            b = x.size(0)\n\n        feats = self.backbone.forward_features(x)\n        cls = feats[:, 0]\n        out = self.head(cls)\n\n        if tta_mode:\n            out = out.view(b, tta, -1)\n\n        return out\n\nclass BiomassModel(nn.Module):\n    def __init__(self,MODEL_NAME,NUM_TARGETS=5):\n        super().__init__()\n\n        self.backbone = timm.create_model(\n            MODEL_NAME,\n            pretrained=False,   # CRITICAL: offline-safe\n            num_classes=0,\n            global_pool=\"\"\n        )\n\n        embed_dim = self.backbone.embed_dim\n\n        self.regressor = nn.Sequential(\n            nn.Linear(embed_dim, 512),\n            nn.ReLU(),\n            nn.Linear(512, 256),\n            nn.Dropout(0.1),\n            nn.ReLU(),\n            nn.Linear(256, NUM_TARGETS)\n        )\n\n    def forward(self, x):\n        b, tta, c, h, w = x.shape\n        x = x.view(-1, c, h, w)\n\n        tokens = self.backbone.forward_features(x)\n        cls = tokens[:, 0]               # CLS token\n        preds = self.regressor(cls)\n\n        return preds.view(b, tta, 5)\n# ====================================================\n# 4. INFERENCE LOOP \n# ====================================================\ndef run_inference():\n    test_df = pd.read_csv(cfg.test_csv)\n    img_col = \"sample_id\" if \"sample_id\" in test_df.columns else test_df.columns[0]\n\n    # Extract clean IDs\n    img_ids = test_df[img_col].apply(lambda x: x.split(\"__\")[0]).unique()\n    print(f\"Total Images to Predict: {len(img_ids)}\")\n\n    # Setup Transforms\n    tmp_model = timm.create_model(cfg.model_name, pretrained=False)\n    data_config = timm.data.resolve_model_data_config(tmp_model)\n    tfms = timm.data.create_transform(**data_config, is_training=False)\n    del tmp_model\n    gc.collect()\n\n    # Loader\n    dataset = TestDataset(img_ids, cfg.test_img_dir, tfms)\n    loader = DataLoader(\n        dataset,\n        batch_size=cfg.batch_size,\n        shuffle=False,\n        num_workers=cfg.num_workers\n    )\n\n    all_seed_preds = []\n    seed_weights = np.array([1,1,1,1])\n    # --- LOOP SEEDS ---\n    for idx,seed_path in enumerate(cfg.seed_dirs):\n        print(f\"\\n>>> Processing Seed Directory: {seed_path}\")\n        \n        fold_ckpts = sorted([f for f in os.listdir(seed_path) if f.endswith(\".pth\")])\n        if not fold_ckpts:\n            print(f\"WARNING: No checkpoints found in {seed_path}\")\n            continue\n            \n        print(f\"Found {len(fold_ckpts)} folds.\")\n        fold_preds_accum = []\n\n        # --- LOOP FOLDS ---\n        for ckpt in fold_ckpts:\n            print(f\"   -> Loading {ckpt}...\")\n            if idx <3:\n                model = BioEstimatorD(cfg.model_name).to(cfg.device)\n            else:\n                model = BiomassModel(cfg.model_name).to(cfg.device)\n            state = torch.load(os.path.join(seed_path, ckpt), map_location=cfg.device)\n            model.load_state_dict(state)\n            model.eval()\n\n            preds_this_fold = []\n\n            with torch.no_grad():\n                for imgs, _ in tqdm(loader, leave=False):\n                    imgs = imgs.to(cfg.device)\n\n                    with autocast(\"cuda\"):\n                        out = model(imgs) # Output shape: [B, TTA, 5]\n\n                    out = torch.clamp(out, max=12) \n                    out = torch.nan_to_num(out, nan=0.0)\n                    out = out.mean(dim=1) # Shape: [B, 5]\n                    # print(out[0])\n                    # --- CRITICAL FIX END ---\n\n                    preds_this_fold.append(out.cpu().numpy())\n\n            # Stack batches -> [N_Images, 5]\n            fold_preds_accum.append(np.vstack(preds_this_fold))\n            \n            del model, state\n            torch.cuda.empty_cache()\n            gc.collect()\n\n        # Average all folds in this seed (Arithmetic Mean)\n        seed_avg = np.mean(fold_preds_accum, axis=0)\n        all_seed_preds.append(seed_avg)\n\n    # Average across all seeds (Arithmetic Mean)\n    final_preds = np.average(all_seed_preds, axis=0,weights=seed_weights)\n    final_preds = np.expm1(final_preds)\n    \n    return img_ids, final_preds\n\n# Run it\nimg_ids, final_preds = run_inference()\n\n# ====================================================\n# 5. SUBMISSION GENERATION (OPTIMIZED)\n# ====================================================\nprint(\"\\nGenerating Submission CSV...\")\n\ntarget_map = {\n    \"Dry_Green_g\": 0,\n    \"Dry_Dead_g\": 1,\n    \"Dry_Clover_g\": 2,\n    \"GDM_g\": 3,\n    \"Dry_Total_g\": 4\n}\n\n# Fast Lookup Dictionary\npred_dict = {img_id: preds for img_id, preds in zip(img_ids, final_preds)}\n\ntest_df = pd.read_csv(cfg.test_csv)\nrows = []\n\nfor _, r in tqdm(test_df.iterrows(), total=len(test_df)):\n    parts = r[\"sample_id\"].split(\"__\")\n    img_clean = parts[0]\n    tgt_name = parts[1]\n    \n    # Instant lookup O(1)\n    if img_clean in pred_dict:\n        val = pred_dict[img_clean][target_map[tgt_name]]\n    else:\n        val = 0.0\n        \n    rows.append({\n        \"sample_id\": r[\"sample_id\"],\n        \"target\": float(val)\n    })\n\nsubmission = pd.DataFrame(rows)\nsubmission.to_csv(\"submission.csv\", index=False)\n\nprint(\"✅ Submission saved successfully.\")\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T11:51:29.403497Z","iopub.execute_input":"2025-12-25T11:51:29.403681Z","iopub.status.idle":"2025-12-25T11:52:05.680349Z","shell.execute_reply.started":"2025-12-25T11:51:29.403661Z","shell.execute_reply":"2025-12-25T11:52:05.679321Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}