{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":272104372,"sourceType":"kernelVersion"},{"sourceId":630489,"sourceType":"modelInstanceVersion","modelInstanceId":471723,"modelId":487624},{"sourceId":638209,"sourceType":"modelInstanceVersion","modelInstanceId":471723,"modelId":487624}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":108.02501,"end_time":"2025-11-23T16:42:36.83923","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-11-23T16:40:48.81422","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\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 albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\n\n# =============================== Configuration =========================================\nclass CFG:\n    # Data paths (Kaggle default)\n    BASE_PATH = \"/kaggle/input/csiro-biomass\"\n    TEST_CSV = os.path.join(BASE_PATH, \"test.csv\")\n    TEST_IMAGE_DIR = os.path.join(BASE_PATH, \"test\")\n\n    # -- Ensemble A (first code) --\n    MODEL_DIR_A = \"/kaggle/input/m/gothamjocker/csiro/pytorch/default/6\"   # Change to your dataset path\n    CKPTS_A = [\n        os.path.join(MODEL_DIR_A, \"tiled_film_best_model_fold1.pth\"),\n        os.path.join(MODEL_DIR_A, \"tiled_film_best_model_fold2.pth\"),\n        os.path.join(MODEL_DIR_A, \"tiled_film_best_model_fold3.pth\"),\n        os.path.join(MODEL_DIR_A, \"tiled_film_best_model_fold4.pth\"),\n        os.path.join(MODEL_DIR_A, \"tiled_film_best_model_fold5.pth\"),\n    ]\n\n    # -- Ensemble B (second code) --\n    MODEL_DIR_B = \"/kaggle/input/m/gothamjocker/csiro/pytorch/default/11\"  # Change to your dataset path\n    CKPTS_B = [\n        os.path.join(MODEL_DIR_B, \"tiled_film_best_model_fold1.pth\"),\n        os.path.join(MODEL_DIR_B, \"tiled_film_best_model_fold2.pth\"),\n        os.path.join(MODEL_DIR_B, \"tiled_film_best_model_fold3.pth\"),\n        os.path.join(MODEL_DIR_B, \"tiled_film_best_model_fold4.pth\"),\n        os.path.join(MODEL_DIR_B, \"tiled_film_best_model_fold5.pth\"),\n    ]\n\n    # Fusion weights\n    W_A = 0.93\n    W_B = 0.07\n\n    SUBMISSION_FILE = \"submission.csv\"\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    BATCH_SIZE = 1\n    NUM_WORKERS = 0\n\n    # Consistent with training\n    DROPOUT = 0.30\n    HIDDEN_RATIO = 0.25\n    GRID = (2, 2)  # Grid for tiled / tiled_film (auto-detected if weights aren't tiled type)\n\n    # Output column order (consistent with training)\n    ALL_TARGET_COLS = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n\n    # DINO candidates during training (priority from high to low)\n    DINO_CANDIDATES = [\n        \"vit_base_patch14_dinov3\",\n        \"vit_base_patch14_reg4_dinov3\",\n        \"vit_small_patch14_dinov3\",\n        \"vit_base_patch14_reg4_dinov2\",\n        \"vit_base_patch14_dinov2\",\n        \"vit_small_patch14_dinov2\",\n    ]\n\nprint(f\"Device: {CFG.DEVICE}\")\n\n# =============================== Dataset (left-right two streams, consistent with training) =================\nclass TestBiomassDataset(Dataset):\n    def __init__(self, df, transform, image_dir):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.image_dir = image_dir\n        self.paths = self.df[\"image_path\"].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        filename = os.path.basename(self.paths[idx])\n        full_path = os.path.join(self.image_dir, filename)\n        img = cv2.imread(full_path)\n        if img is None:\n            # Fallback: if reading fails, use black image\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # Split left-right (unified train/val strategy)\n        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n\n        left_t = self.transform(image=left)[\"image\"]\n        right_t = self.transform(image=right)[\"image\"]\n        return left_t, right_t\n\n# =============================== Build DINO backbone (inference end) ======================\ndef _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    name = getattr(m, \"default_cfg\", {}).get(\"architecture\", \"\") or str(type(m))\n    return 518 if (\"dinov2\" in name.lower()) else 224\n\ndef _build_dino_by_name(name: str):\n    # Inference doesn't need to download pretrained, directly pretrained=False (will be completely overwritten by weights)\n    m = timm.create_model(name, pretrained=False, num_classes=0)\n    feat = m.num_features\n    input_res = _infer_input_res(m)\n    return m, feat, input_res\n\n# =============================== Model (strictly aligned with training) ========================\nclass TwoStreamDINOBase(nn.Module):\n    def __init__(self, backbone_name: str, dropout: float = 0.3, hidden_ratio: float = 0.25):\n        super().__init__()\n        self.backbone, feat, input_res = _build_dino_by_name(backbone_name)\n        self.used_backbone_name = backbone_name\n        self.input_res = int(input_res)\n        self.feat_dim = feat\n        self.combined = feat * 2\n\n        hidden = max(8, int(self.combined * hidden_ratio))\n\n        def head():\n            return nn.Sequential(\n                nn.Linear(self.combined, hidden),\n                nn.ReLU(inplace=True),\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.softplus = nn.Softplus(beta=1.0)\n\n    def _merge_heads(self, f_l: torch.Tensor, f_r: torch.Tensor):\n        f = torch.cat([f_l, f_r], dim=1)\n        green_pos = self.softplus(self.head_green(f))\n        clover_pos = self.softplus(self.head_clover(f))\n        dead_pos = self.softplus(self.head_dead(f))\n        gdm = green_pos + clover_pos\n        total = gdm + dead_pos\n        return total, gdm, green_pos\n\nclass TwoStreamDINOPlain(TwoStreamDINOBase):\n    def forward(self, x_left: torch.Tensor, x_right: torch.Tensor):\n        f_l = self.backbone(x_left)\n        f_r = self.backbone(x_right)\n        return self._merge_heads(f_l, f_r)\n\ndef _make_edges(L: int, parts: int):\n    step = L // parts\n    edges = []\n    start = 0\n    for _ in range(parts - 1):\n        edges.append((start, start + step))\n        start += step\n    edges.append((start, L))\n    return edges\n\nclass TwoStreamDINOTiled(TwoStreamDINOBase):\n    def __init__(self, backbone_name: str, grid=(2, 2), **kwargs):\n        super().__init__(backbone_name, **kwargs)\n        self.grid = tuple(grid)\n\n    def _encode_tiles(self, x: torch.Tensor) -> torch.Tensor:\n        B, C, H, W = x.shape\n        r, c = self.grid\n        rows = _make_edges(H, r)\n        cols = _make_edges(W, c)\n        feats = []\n        for (rs, re) in rows:\n            for (cs, ce) in cols:\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                ft = self.backbone(xt)\n                feats.append(ft)\n        feats = torch.stack(feats, dim=0).permute(1, 0, 2)  # (B, T, F)\n        feat_stream = feats.mean(dim=1)  # (B, F)\n        return feat_stream\n\n    def forward(self, x_left: torch.Tensor, x_right: torch.Tensor):\n        f_l = self._encode_tiles(x_left)\n        f_r = self._encode_tiles(x_right)\n        return self._merge_heads(f_l, f_r)\n\nclass FiLM(nn.Module):\n    def __init__(self, in_dim: int):\n        super().__init__()\n        hid = max(32, in_dim // 2)\n        self.mlp = nn.Sequential(nn.Linear(in_dim, hid), nn.ReLU(inplace=True), nn.Linear(hid, in_dim * 2))\n\n    def forward(self, context: torch.Tensor):\n        gb = self.mlp(context)\n        gamma, beta = torch.chunk(gb, 2, dim=1)\n        return gamma, beta\n\nclass TwoStreamDINOTiledFiLM(TwoStreamDINOBase):\n    def __init__(self, backbone_name: str, grid=(2, 2), **kwargs):\n        super().__init__(backbone_name, **kwargs)\n        self.grid = tuple(grid)\n        self.film_left = FiLM(self.feat_dim)\n        self.film_right = FiLM(self.feat_dim)\n\n    def _tiles_backbone(self, x: torch.Tensor) -> torch.Tensor:\n        B, C, H, W = x.shape\n        r, c = self.grid\n        rows = _make_edges(H, r)\n        cols = _make_edges(W, c)\n        feats = []\n        for (rs, re) in rows:\n            for (cs, ce) in cols:\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                ft = self.backbone(xt)\n                feats.append(ft)\n        feats = torch.stack(feats, dim=0).permute(1, 0, 2)  # (B, T, F)\n        return feats\n\n    def _encode_stream(self, x: torch.Tensor, film: FiLM) -> torch.Tensor:\n        tiles = self._tiles_backbone(x)  # (B, T, F)\n        context = tiles.mean(dim=1)      # (B, F)\n        gamma, beta = film(context)      # (B, F)\n        tiles = tiles * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)\n        feat_stream = tiles.mean(dim=1)  # (B, F)\n        return feat_stream\n\n    def forward(self, x_left: torch.Tensor, x_right: torch.Tensor):\n        f_l = self._encode_stream(x_left, self.film_left)\n        f_r = self._encode_stream(x_right, self.film_right)\n        return self._merge_heads(f_l, f_r)\n\n# =============================== Weight cleaning and loading =================================\ndef _strip_module_prefix(sd: dict):\n    if not sd:\n        return sd\n    keys = list(sd.keys())\n    if all(k.startswith(\"module.\") for k in keys):\n        return {k[len(\"module.\") :]: v for k, v in sd.items()}\n    return sd\n\ndef _extract_student_substate(sd: dict) -> dict:\n    \"\"\"\n    Training saves MultiModalStudentTeacher's state_dict.\n    Here we strip:\n      - First remove 'module.' prefix\n      - If 'student.' prefix exists, keep only that subtree and remove 'student.' prefix\n      - Discard training-only multi-modal/distillation related layers (txt_enc.*, img_proj.*, txt_film_left.*, txt_film_right.*)\n    \"\"\"\n    sd = _strip_module_prefix(sd)\n    has_student = any(k.startswith(\"student.\") for k in sd.keys())\n    if has_student:\n        sd = {k[len(\"student.\") :]: v for k, v in sd.items() if k.startswith(\"student.\")}\n    drop_prefixes = (\"txt_enc.\", \"img_proj.\", \"txt_film_left.\", \"txt_film_right.\")\n    sd = {k: v for k, v in sd.items() if not k.startswith(drop_prefixes)}\n    return sd\n\ndef _has_film(sd_keys: set) -> bool:\n    return any(k.startswith(\"film_left.mlp.\") for k in sd_keys) or any(k.startswith(\"film_right.mlp.\") for k in sd_keys)\n\ndef _try_build_and_load(sd: dict, backbone_name: str, variant: str, grid=(2, 2)):\n    if variant == \"tiled_film\":\n        model = TwoStreamDINOTiledFiLM(backbone_name, grid=grid, dropout=CFG.DROPOUT, hidden_ratio=CFG.HIDDEN_RATIO)\n    elif variant == \"tiled\":\n        model = TwoStreamDINOTiled(backbone_name, grid=grid, dropout=CFG.DROPOUT, hidden_ratio=CFG.HIDDEN_RATIO)\n    else:\n        model = TwoStreamDINOPlain(backbone_name, dropout=CFG.DROPOUT, hidden_ratio=CFG.HIDDEN_RATIO)\n\n    # Strict matching: missing/unexpected must both be 0\n    result = model.load_state_dict(sd, strict=False)\n    missing = getattr(result, \"missing_keys\", [])\n    unexpected = getattr(result, \"unexpected_keys\", [])\n    if len(missing) == 0 and len(unexpected) == 0:\n        model.to(CFG.DEVICE)\n        model.eval()\n        return model\n    return None\n\ndef load_fold_model_auto(path: str, grid=(2, 2)):\n    if not os.path.exists(path):\n        raise FileNotFoundError(path)\n    # Compatible with PyTorch 2.6 weights_only change\n    try:\n        raw_sd = torch.load(path, map_location=CFG.DEVICE, weights_only=True)\n    except TypeError:\n        raw_sd = torch.load(path, map_location=CFG.DEVICE)\n    sd = _extract_student_substate(raw_sd)\n    keys = set(sd.keys())\n\n    # Determine if FiLM\n    is_film = _has_film(keys)\n    variant_order = [\"tiled_film\"] if is_film else [\"tiled\", \"plain\"]\n\n    # Enumerate backbone + variant until strict match\n    for variant in variant_order:\n        for backbone in CFG.DINO_CANDIDATES:\n            try:\n                m = _try_build_and_load(sd, backbone, variant, grid=grid)\n                if m is not None:\n                    return m, variant, backbone\n            except Exception:\n                continue\n\n    # If we reach here, failed to match\n    raise RuntimeError(\n        f\"Cannot find matching variant/backbone for {os.path.basename(path)}.\"\n        f\" Detected {'tiled_film' if is_film else 'non-film'} weights, please check training-inference consistency.\"\n    )\n\n# =============================== TTA transforms (dynamic based on model input_res) ===============\ndef get_tta_transforms(img_size: int):\n    base = [A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2()]\n    original = A.Compose([A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA), *base])\n    hflip = A.Compose([A.HorizontalFlip(p=1.0), A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA), *base])\n    vflip = A.Compose([A.VerticalFlip(p=1.0), A.Resize(img_size, img_size, interpolation=cv2.INTER_AREA), *base])\n    return [original, hflip, vflip]\n\n# =============================== Inference (unified output 5 targets) =======================\n@torch.no_grad()\ndef predict_one_view(models, loader):\n    out_list = []\n    amp_dtype = \"cuda\" if CFG.DEVICE.type == \"cuda\" else \"cpu\"\n\n    for (xl, xr) in tqdm(loader, desc=\"  Predicting View\", leave=False):\n        xl = xl.to(CFG.DEVICE, non_blocking=True)\n        xr = xr.to(CFG.DEVICE, non_blocking=True)\n\n        per_model_preds = []\n        with torch.amp.autocast(amp_dtype, enabled=(CFG.DEVICE.type == \"cuda\")):\n            for m in models:\n                total, gdm, green = m(xl, xr)\n                dead = total - gdm\n                clover = gdm - green\n                five = torch.cat([green, dead, clover, gdm, total], dim=1)\n                five = torch.clamp(five, min=0.0)  # Non-negative constraint\n                per_model_preds.append(five.float().cpu())\n\n        stacked = torch.mean(torch.stack(per_model_preds, dim=0), dim=0)\n        out_list.append(stacked.numpy())\n\n    return np.concatenate(out_list, axis=0)\n\ndef run_inference_for_ckpts(ckpt_list, test_unique, image_dir):\n    print(\"\\n================= Loading models (5 folds) =================\")\n    models = []\n\n    # First use fold 1 to determine input resolution\n    m1, v1, b1 = load_fold_model_auto(ckpt_list[0], grid=CFG.GRID)\n    models.append(m1)\n    backbone_res = int(getattr(m1, \"input_res\", 518))\n    print(f\"fold1 => variant={v1}, backbone={b1}, input_res={backbone_res}\")\n\n    for p in ckpt_list[1:]:\n        m, v, b = load_fold_model_auto(p, grid=CFG.GRID)\n        print(f\"{os.path.basename(p)} => variant={v}, backbone={b}, input_res={getattr(m, 'input_res', '?')}\")\n        models.append(m)\n\n    # Prepare TTA views\n    tta_trans = get_tta_transforms(backbone_res)\n    per_view_preds = []\n    for i, t in enumerate(tta_trans):\n        print(f\"\\n--- TTA view {i+1}/{len(tta_trans)} (resize={backbone_res}) ---\")\n        ds = TestBiomassDataset(test_unique, t, image_dir)\n        dl = DataLoader(ds, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n        view_5 = predict_one_view(models, dl)  # [N,5]\n        per_view_preds.append(view_5)\n\n    # TTA average\n    final_5 = np.mean(per_view_preds, axis=0)  # [N,5]: [green, dead, clover, gdm, total]\n    return final_5\n\ndef run_dual_ensembles_and_fuse():\n    print(\"\\n================= Loading test data =================\")\n    test_long = pd.read_csv(CFG.TEST_CSV)\n    test_unique = test_long.drop_duplicates(subset=[\"image_path\"]).reset_index(drop=True)\n    print(f\"Found {len(test_unique)} unique test images.\")\n\n    print(\"\\n================= Ensemble A inference =================\")\n    final_a = run_inference_for_ckpts(CFG.CKPTS_A, test_unique, CFG.TEST_IMAGE_DIR)\n\n    print(\"\\n================= Ensemble B inference =================\")\n    final_b = run_inference_for_ckpts(CFG.CKPTS_B, test_unique, CFG.TEST_IMAGE_DIR)\n\n    assert final_a.shape == final_b.shape, \"A/B results shape mismatch, cannot perform weighted fusion.\"\n    final = CFG.W_A * final_a + CFG.W_B * final_b\n    return final, test_long, test_unique\n\n# =============================== Generate submission ======================================\ndef create_submission(final_5, test_long, test_unique):\n    green = final_5[:, 0]\n    dead = final_5[:, 1]\n    clover = final_5[:, 2]\n    gdm = final_5[:, 3]\n    total = final_5[:, 4]\n\n    # Final non-negative clipping and NaN/Inf handling\n    def nnz(x):\n        x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n        return np.maximum(0, x)\n\n    green, dead, clover, gdm, total = map(nnz, [green, dead, clover, gdm, total])\n\n    wide = pd.DataFrame(\n        {\n            \"image_path\": test_unique[\"image_path\"],\n            \"Dry_Green_g\": green,\n            \"Dry_Dead_g\": dead,\n            \"Dry_Clover_g\": clover -1,\n            \"GDM_g\": green + clover -1  ,\n            \"Dry_Total_g\": green + clover -1  + dead ,\n        }\n    )\n\n    long_preds = wide.melt(\n        id_vars=[\"image_path\"],\n        value_vars=CFG.ALL_TARGET_COLS,\n        var_name=\"target_name\",\n        value_name=\"target\",\n    )\n\n    sub = pd.merge(\n        test_long[[\"sample_id\", \"image_path\", \"target_name\"]],\n        long_preds,\n        on=[\"image_path\", \"target_name\"],\n        how=\"left\",\n    )[[\"sample_id\", \"target\"]]\n\n    sub[\"target\"] = np.nan_to_num(sub[\"target\"], nan=0.0, posinf=0.0, neginf=0.0)\n    sub.to_csv(CFG.SUBMISSION_FILE, index=False)\n    print(f\"\\n🎉 Generated submission file: {CFG.SUBMISSION_FILE}\")\n    print(sub.head())\n    return sub\n\n# =============================== Entry point =========================================\nif __name__ == \"__main__\":\n    final_5, df_long, df_unique = run_dual_ensembles_and_fuse()\n    _ = create_submission(final_5, df_long, df_unique)\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2025-11-24T20:03:29.799413Z","iopub.execute_input":"2025-11-24T20:03:29.800022Z","iopub.status.idle":"2025-11-24T20:03:53.146311Z","shell.execute_reply.started":"2025-11-24T20:03:29.799993Z","shell.execute_reply":"2025-11-24T20:03:53.145561Z"},"papermill":{"duration":101.995539,"end_time":"2025-11-23T16:42:34.554752","exception":false,"start_time":"2025-11-23T16:40:52.559213","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}