{"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":293691576,"sourceType":"kernelVersion"},{"sourceId":268942,"sourceType":"modelInstanceVersion","modelInstanceId":230141,"modelId":251887},{"sourceId":719631,"sourceType":"modelInstanceVersion","modelInstanceId":546422,"modelId":559381},{"sourceId":732135,"sourceType":"modelInstanceVersion","modelInstanceId":551547,"modelId":564133}],"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&lt;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":"markdown","source":"# SigLip","metadata":{}},{"cell_type":"code","source":"# ============================================================\n# SIGLIP EMBEDDINGS + TABULAR ENSEMBLE\n# - TRAIN on 3x augmented embeddings (TTA modes 1/2/3 with same label)\n# - INFER by predicting on the 3 TTA embeddings and averaging predictions\n# - SigLIP mean/std + input size are taken from AutoImageProcessor\n# - FIX: Albumentations Gamma -> RandomGamma\n# ============================================================\n\nimport os\nimport gc\nimport random\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nfrom pathlib import Path\nfrom copy import deepcopy\nfrom dataclasses import dataclass\n\n# Albumentations for TTA\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# Sklearn & Models\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.decomposition import PCA\nfrom sklearn.cross_decomposition import PLSRegression\nfrom sklearn.mixture import GaussianMixture\nfrom sklearn.ensemble import GradientBoostingRegressor, HistGradientBoostingRegressor\nfrom lightgbm import LGBMRegressor\nfrom catboost import CatBoostRegressor\n\nfrom transformers import AutoModel, AutoImageProcessor, AutoTokenizer\n\nwarnings.filterwarnings(\"ignore\")\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\n\n# =========================================================================================\n# 1. CONFIGURATION & SEEDING\n# =========================================================================================\n@dataclass\nclass Config:\n    DATA_PATH: Path = Path(\"/kaggle/input/csiro-biomass/\")\n    SPLIT_PATH: Path = Path(\"/kaggle/input/csiro-datasplit/csiro_data_split.csv\")\n    SIGLIP_PATH: str = \"/kaggle/input/google-siglip-so400m-patch14-384/transformers/default/1\"\n\n    SEED: int = 42\n    DEVICE: str = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Tiling (for high-res)\n    PATCH_SIZE: int = 520\n    OVERLAP: int = 16\n\n    # >>> 3x augmentation for training (exactly 3 modes)\n    TRAIN_TTA_MODES = (1, 2, 3)   # HFlip, VFlip, photometric\n    INFER_TTA_MODES = (1, 2, 3)   # match training distribution; add 0 if you want\n\n    # Patch mini-batch for SigLIP forward (VRAM control)\n    PATCH_BATCH: int = 16   # lower this if VRAM is tight (e.g., 8 or 6)\n\n    TARGET_NAMES = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'Dry_Total_g', 'GDM_g']\n    TARGET_MAX = {\n        \"Dry_Clover_g\": 71.7865,\n        \"Dry_Dead_g\": 83.8407,\n        \"Dry_Green_g\": 157.9836,\n        \"Dry_Total_g\": 185.70,\n        \"GDM_g\": 157.9836,\n    }\n\ncfg = Config()\n\n\ndef seeding(SEED: int):\n    np.random.seed(SEED)\n    random.seed(SEED)\n    os.environ[\"PYTHONHASHSEED\"] = str(SEED)\n    torch.manual_seed(SEED)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(SEED)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\nseeding(cfg.SEED)\n\n\n# =========================================================================================\n# 2. DATA LOADING & PRE-PROCESSING\n# =========================================================================================\ndef pivot_table(df: pd.DataFrame) -> pd.DataFrame:\n    if 'target' in df.columns.tolist():\n        df_pt = pd.pivot_table(\n            df,\n            values='target',\n            index=['image_path', 'Sampling_Date', 'State', 'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm'],\n            columns='target_name',\n            aggfunc='mean'\n        ).reset_index()\n    else:\n        df = df.copy()\n        df['target'] = 0\n        df_pt = pd.pivot_table(\n            df,\n            values='target',\n            index='image_path',\n            columns='target_name',\n            aggfunc='mean'\n        ).reset_index()\n    return df_pt\n\n\ndef melt_table(df: pd.DataFrame) -> pd.DataFrame:\n    melted = df.melt(\n        id_vars='image_path',\n        value_vars=cfg.TARGET_NAMES,\n        var_name='target_name',\n        value_name='target'\n    )\n    melted['sample_id'] = (\n        melted['image_path']\n        .str.replace(r'^.*/', '', regex=True)\n        .str.replace('.jpg', '', regex=False)\n        + '__' + melted['target_name']\n    )\n    return melted[['sample_id', 'target']]\n\n\ndef post_process_biomass(df_preds: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"\n    Mass-balance post-process (your original behavior):\n      - Keep Dry_Clover_g fixed at 0\n      - Derive GDM = Green + Clover\n      - Derive Total = GDM + Dead\n      - Clip derived targets nonneg\n    \"\"\"\n    ordered_cols = [\"Dry_Green_g\", \"Dry_Clover_g\", \"Dry_Dead_g\", \"GDM_g\", \"Dry_Total_g\"]\n    df_out = df_preds.copy()\n    for c in ordered_cols:\n        if c not in df_out.columns:\n            df_out[c] = 0.0\n\n    df_out['Dry_Clover_g'] = 0.0\n    df_out['GDM_g'] = df_out['Dry_Green_g'] + df_out['Dry_Clover_g']\n    df_out['Dry_Total_g'] = df_out['GDM_g'] + df_out['Dry_Dead_g']\n    df_out['GDM_g'] = df_out['GDM_g'].clip(lower=0.0)\n    df_out['Dry_Total_g'] = df_out['Dry_Total_g'].clip(lower=0.0)\n    return df_out\n\n\nprint(\"Loading Data...\")\ntrain_df = pd.read_csv(cfg.SPLIT_PATH)\n\n# Remove any pre-existing emb* columns\ncols_to_keep = [c for c in train_df.columns if not c.startswith('emb')]\ntrain_df = train_df[cols_to_keep]\n\n# Ensure train paths match local environment\nif not str(train_df['image_path'].iloc[0]).startswith('/'):\n    train_df['image_path'] = train_df['image_path'].apply(\n        lambda p: str(cfg.DATA_PATH / 'train' / os.path.basename(p))\n    )\n\n# Load Test (wide)\ntest_df_raw = pd.read_csv(cfg.DATA_PATH / 'test.csv')\ntest_df_base = pivot_table(test_df_raw)\ntest_df_base['image_path'] = test_df_base['image_path'].apply(lambda p: str(cfg.DATA_PATH / p))\n\n\n# =========================================================================================\n# 3. FEATURE EXTRACTION: SIGLIP IMAGE EMBEDDINGS (WITH TTA MODES)\n# =========================================================================================\ndef split_image(image: np.ndarray, patch_size: int = 520, overlap: int = 16):\n    h, w, c = image.shape\n    stride = patch_size - overlap\n    patches = []\n    for y in range(0, h, stride):\n        for x in range(0, w, stride):\n            y2 = min(y + patch_size, h)\n            x2 = min(x + patch_size, w)\n            y1 = max(0, y2 - patch_size)\n            x1 = max(0, x2 - patch_size)\n            patches.append(image[y1:y2, x1:x2, :])\n    return patches\n\n\ndef _hw(img_size):\n    if isinstance(img_size, (tuple, list)) and len(img_size) == 2:\n        return int(img_size[0]), int(img_size[1])\n    return int(img_size), int(img_size)\n\n\ndef _processor_hw(proc, fallback=384):\n    size = getattr(proc, \"size\", None)\n    if isinstance(size, dict):\n        if \"height\" in size and \"width\" in size:\n            return int(size[\"height\"]), int(size[\"width\"])\n        if \"shortest_edge\" in size:\n            s = int(size[\"shortest_edge\"])\n            return s, s\n    if isinstance(size, (int, float)):\n        s = int(size)\n        return s, s\n    crop = getattr(proc, \"crop_size\", None)\n    if isinstance(crop, dict) and \"height\" in crop and \"width\" in crop:\n        return int(crop[\"height\"]), int(crop[\"width\"])\n    return int(fallback), int(fallback)\n\n\ndef get_tta_transforms(img_size=(384, 384), mode=0, mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5)):\n    \"\"\"\n    Modes:\n      0: base\n      1: HFlip\n      2: VFlip\n      3: photometric (brightness/contrast + gamma)\n\n    NOTE: FIXED A.Gamma -> A.RandomGamma for Kaggle albumentations versions.\n    \"\"\"\n    h, w = _hw(img_size)\n    transforms_list = [A.Resize(h, w)]\n\n    if mode == 1:\n        transforms_list.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        transforms_list.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        transforms_list.extend([\n            A.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, p=1.0),\n            A.RandomGamma(gamma_limit=(90, 110), p=1.0),\n        ])\n\n    transforms_list.extend([\n        A.Normalize(mean=mean, std=std, max_pixel_value=255.0),\n        ToTensorV2()\n    ])\n\n    return A.Compose(transforms_list)\n\n\n@torch.inference_mode()\ndef _embed_patches_batched(model, patches_rgb, tfm, device, patch_batch: int = 16):\n    \"\"\"\n    Apply Albumentations to each patch -> pixel_values -> SigLIP get_image_features\n    Returns (D,) numpy embedding.\n    \"\"\"\n    feat_sum = None\n    n = 0\n    buf = []\n\n    for p in patches_rgb:\n        out = tfm(image=p)\n        buf.append(out[\"image\"])  # torch CHW\n\n        if len(buf) >= patch_batch:\n            pv = torch.stack(buf, dim=0).to(device, non_blocking=True)\n            feats = model.get_image_features(pixel_values=pv).detach().float().cpu()\n            feat_sum = feats.sum(dim=0) if feat_sum is None else (feat_sum + feats.sum(dim=0))\n            n += feats.shape[0]\n            buf.clear()\n            del pv, feats\n            if device.type == \"cuda\":\n                torch.cuda.empty_cache()\n\n    if buf:\n        pv = torch.stack(buf, dim=0).to(device, non_blocking=True)\n        feats = model.get_image_features(pixel_values=pv).detach().float().cpu()\n        feat_sum = feats.sum(dim=0) if feat_sum is None else (feat_sum + feats.sum(dim=0))\n        n += feats.shape[0]\n        buf.clear()\n        del pv, feats\n        if device.type == \"cuda\":\n            torch.cuda.empty_cache()\n\n    if feat_sum is None or n == 0:\n        return None\n    return (feat_sum / float(n)).numpy()\n\n\ndef compute_embeddings_by_mode(model_path, df_base, tta_modes, patch_batch: int = 16):\n    \"\"\"\n    Compute embeddings for each image in df_base for each tta_mode, returning:\n      embeddings_by_mode[mode] = (N, D)\n\n    This reads each image once and then runs each mode.\n    \"\"\"\n    tta_modes = tuple(tta_modes)\n    print(f\"Computing embeddings: N={len(df_base)} | modes={tta_modes}\")\n    device = torch.device(cfg.DEVICE)\n\n    model = AutoModel.from_pretrained(model_path, local_files_only=True).eval().to(device)\n    processor = AutoImageProcessor.from_pretrained(model_path)\n\n    mean = getattr(processor, \"image_mean\", (0.5, 0.5, 0.5))\n    std  = getattr(processor, \"image_std\",  (0.5, 0.5, 0.5))\n    in_h, in_w = _processor_hw(processor, fallback=384)\n    siglip_size = (in_h, in_w)\n\n    tfms = {m: get_tta_transforms(img_size=siglip_size, mode=m, mean=mean, std=std) for m in tta_modes}\n\n    # store per-mode lists then stack\n    per_mode_lists = {m: [] for m in tta_modes}\n    D_fallback = 1152  # safe fallback if something fails\n\n    for i, row in tqdm(df_base.iterrows(), total=len(df_base)):\n        path = row[\"image_path\"]\n        try:\n            img = cv2.imread(path)\n            if img is None:\n                raise ValueError(\"Image not found\")\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n            patches = split_image(img, patch_size=cfg.PATCH_SIZE, overlap=cfg.OVERLAP)\n\n            for m in tta_modes:\n                emb = _embed_patches_batched(\n                    model=model,\n                    patches_rgb=patches,\n                    tfm=tfms[m],\n                    device=device,\n                    patch_batch=int(patch_batch),\n                )\n                if emb is None:\n                    raise ValueError(f\"Empty embedding (mode={m})\")\n                per_mode_lists[m].append(emb.astype(np.float32))\n\n        except Exception as e:\n            print(f\"Error processing {path}: {e}\")\n            z = np.zeros((D_fallback,), dtype=np.float32)\n            for m in tta_modes:\n                per_mode_lists[m].append(z)\n\n        if (i + 1) % 32 == 0 and device.type == \"cuda\":\n            torch.cuda.empty_cache()\n\n    # stack\n    embeddings_by_mode = {m: np.stack(per_mode_lists[m], axis=0) for m in tta_modes}\n\n    # cleanup\n    del model, processor, tfms\n    if device.type == \"cuda\":\n        torch.cuda.empty_cache()\n    gc.collect()\n\n    return embeddings_by_mode\n\n\ndef build_augmented_df_and_embeddings(df_base, embeddings_by_mode, tta_modes, add_targets=True):\n    \"\"\"\n    Given:\n      - df_base: (N, ...)\n      - embeddings_by_mode[mode]: (N, D)\n    Build:\n      - df_aug: (N*M, ...) with a tta_mode column\n      - X_aug:  (N*M, D) concatenated in mode-block order\n    \"\"\"\n    blocks = []\n    X_blocks = []\n    for m in tta_modes:\n        d = df_base.copy()\n        d[\"tta_mode\"] = int(m)\n        blocks.append(d)\n        X_blocks.append(embeddings_by_mode[m])\n    df_aug = pd.concat(blocks, ignore_index=True)\n    X_aug = np.concatenate(X_blocks, axis=0).astype(np.float32)\n    return df_aug, X_aug\n\n\n# --- Compute SigLIP embeddings per mode ---\ntrain_emb_by_mode = compute_embeddings_by_mode(\n    cfg.SIGLIP_PATH, train_df, tta_modes=cfg.TRAIN_TTA_MODES, patch_batch=cfg.PATCH_BATCH\n)\ntest_emb_by_mode = compute_embeddings_by_mode(\n    cfg.SIGLIP_PATH, test_df_base, tta_modes=cfg.INFER_TTA_MODES, patch_batch=cfg.PATCH_BATCH\n)\n\n# --- Augment TRAIN (3x) and TEST (3x) ---\ntrain_aug_df, train_aug_emb = build_augmented_df_and_embeddings(\n    train_df, train_emb_by_mode, tta_modes=cfg.TRAIN_TTA_MODES\n)\ntest_aug_df, test_aug_emb = build_augmented_df_and_embeddings(\n    test_df_base, test_emb_by_mode, tta_modes=cfg.INFER_TTA_MODES\n)\n\nemb_cols = [f\"emb{i}\" for i in range(train_aug_emb.shape[1])]\ntrain_feat_df = pd.concat([train_aug_df.reset_index(drop=True), pd.DataFrame(train_aug_emb, columns=emb_cols)], axis=1)\ntest_feat_df  = pd.concat([test_aug_df.reset_index(drop=True),  pd.DataFrame(test_aug_emb,  columns=emb_cols)], axis=1)\n\nprint(f\"Train (aug) shape: {train_feat_df.shape}  | modes={cfg.TRAIN_TTA_MODES}\")\nprint(f\"Test  (aug) shape: {test_feat_df.shape}   | modes={cfg.INFER_TTA_MODES}\")\n\n\n# =========================================================================================\n# 4. FEATURE EXTRACTION: SEMANTIC FEATURES (TEXT PROBING)\n# =========================================================================================\n@torch.inference_mode()\ndef generate_semantic_features(image_embeddings_np, model_path):\n    \"\"\"\n    Uses SigLIP text encoder to compute concept similarities from the image embeddings.\n    \"\"\"\n    print(\"Generating Semantic Features...\")\n    device = torch.device(cfg.DEVICE)\n\n    model = AutoModel.from_pretrained(model_path, local_files_only=True).eval().to(device)\n    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)\n\n    concept_groups = {\n        \"bare\": [\"bare soil\", \"dirt ground\", \"sparse vegetation\", \"exposed earth\"],\n        \"sparse\": [\"low density pasture\", \"thin grass\", \"short clipped grass\"],\n        \"medium\": [\"average pasture cover\", \"medium height grass\", \"grazed pasture\"],\n        \"dense\": [\"dense tall pasture\", \"thick grassy volume\", \"high biomass\", \"overgrown vegetation\"],\n        \"green\": [\"lush green vibrant pasture\", \"photosynthesizing leaves\", \"fresh growth\"],\n        \"dead\": [\"dry brown dead grass\", \"yellow straw\", \"senesced material\", \"standing hay\"],\n        \"clover\": [\"white clover\", \"trifolium repens\", \"broadleaf legume\", \"clover flowers\"],\n        \"grass\": [\"ryegrass\", \"blade-like leaves\", \"fescue\", \"grassy sward\"]\n    }\n\n    # normalize image embeddings\n    img_tensor = torch.tensor(image_embeddings_np, dtype=torch.float32, device=device)\n    img_tensor = img_tensor / (img_tensor.norm(p=2, dim=-1, keepdim=True) + 1e-12)\n\n    concept_vectors = {}\n    for name, prompts in concept_groups.items():\n        inputs = tokenizer(prompts, padding=True, truncation=True, return_tensors=\"pt\").to(device)\n        txt = model.get_text_features(**inputs)\n        txt = txt / (txt.norm(p=2, dim=-1, keepdim=True) + 1e-12)\n        concept_vectors[name] = txt.mean(dim=0, keepdim=True)\n        del inputs, txt\n        if device.type == \"cuda\":\n            torch.cuda.empty_cache()\n\n    scores = {}\n    for name, vec in concept_vectors.items():\n        scores[name] = torch.matmul(img_tensor, vec.T).detach().cpu().numpy().flatten()\n\n    df_scores = pd.DataFrame(scores)\n    df_scores['ratio_greenness'] = df_scores['green'] / (df_scores['green'] + df_scores['dead'] + 1e-6)\n    df_scores['ratio_clover']    = df_scores['clover'] / (df_scores['clover'] + df_scores['grass'] + 1e-6)\n    df_scores['ratio_cover']     = (df_scores['dense'] + df_scores['medium']) / (df_scores['bare'] + df_scores['sparse'] + 1e-6)\n\n    del model, tokenizer, img_tensor, concept_vectors\n    if device.type == \"cuda\":\n        torch.cuda.empty_cache()\n    gc.collect()\n\n    return df_scores.values.astype(np.float32)\n\n\n# semantic for augmented (train+test) so shapes match exactly\nall_emb = np.vstack([train_aug_emb, test_aug_emb]).astype(np.float32)\nall_semantic = generate_semantic_features(all_emb, cfg.SIGLIP_PATH)\nsem_train = all_semantic[:len(train_feat_df)]\nsem_test  = all_semantic[len(train_feat_df):]\n\n\n# =========================================================================================\n# 5. SUPERVISED EMBEDDING ENGINE\n# =========================================================================================\nclass SupervisedEmbeddingEngine:\n    def __init__(self, n_pca=0.80, n_pls=8, n_gmm=6, random_state=42):\n        self.scaler = StandardScaler()\n        self.pca = PCA(n_components=n_pca, random_state=random_state)\n        self.pls = PLSRegression(n_components=n_pls, scale=False)\n        self.gmm = GaussianMixture(n_components=n_gmm, covariance_type='diag', random_state=random_state)\n        self.pls_fitted_ = False\n\n    def fit(self, X, y=None, X_semantic=None):\n        X_scaled = self.scaler.fit_transform(X)\n        self.pca.fit(X_scaled)\n        self.gmm.fit(X_scaled)\n        if y is not None:\n            self.pls.fit(X_scaled, y)\n            self.pls_fitted_ = True\n        return self\n\n    def transform(self, X, X_semantic=None):\n        X_scaled = self.scaler.transform(X)\n        feats = [self.pca.transform(X_scaled)]\n        if self.pls_fitted_:\n            feats.append(self.pls.transform(X_scaled))\n        feats.append(self.gmm.predict_proba(X_scaled))\n        if X_semantic is not None:\n            sem_norm = (X_semantic - np.mean(X_semantic, axis=0)) / (np.std(X_semantic, axis=0) + 1e-6)\n            feats.append(sem_norm.astype(np.float32))\n        return np.hstack(feats).astype(np.float32)\n\n\n# =========================================================================================\n# 6. TRAINING & INFERENCE (5-FOLD CV)  +  TTA prediction averaging\n# =========================================================================================\ndef cross_validate_predict(model_cls, model_params, train_data, test_data, sem_tr, sem_te, feature_engine):\n    \"\"\"\n    train_data is augmented 3x (tta_mode column exists)\n    test_data is augmented 3x in the SAME mode-block order\n\n    We train on augmented samples with sample_weight=1/M to avoid overweighting duplicates.\n    We predict on augmented test samples and return (N_test_aug, T).\n    \"\"\"\n    target_max_arr = np.array([cfg.TARGET_MAX[t] for t in cfg.TARGET_NAMES], dtype=np.float32)\n\n    n_splits = int(train_data['fold'].nunique())\n\n    X_train_full = train_data[emb_cols].values.astype(np.float32)\n    X_test_raw   = test_data[emb_cols].values.astype(np.float32)\n    y_train_full = train_data[cfg.TARGET_NAMES].values.astype(np.float32)\n\n    # weight each TTA copy so each original image has total weight ~1\n    if \"tta_mode\" in train_data.columns:\n        M = int(train_data[\"tta_mode\"].nunique())\n        base_w = 1.0 / max(M, 1)\n        w_full = np.full((len(train_data),), base_w, dtype=np.float32)\n    else:\n        w_full = None\n\n    y_pred_test_accum = np.zeros((len(test_data), len(cfg.TARGET_NAMES)), dtype=np.float32)\n\n    for fold in range(n_splits):\n        print(f\"Processing Fold {fold}...\")\n        train_mask = (train_data['fold'].values != fold)\n\n        X_tr = X_train_full[train_mask]\n        y_tr = (y_train_full[train_mask] / target_max_arr).astype(np.float32)\n        sem_tr_fold = sem_tr[train_mask]\n\n        sw_tr = w_full[train_mask] if w_full is not None else None\n\n        engine = deepcopy(feature_engine)\n        engine.fit(X_tr, y=y_tr, X_semantic=sem_tr_fold)\n\n        x_tr_eng = engine.transform(X_tr, X_semantic=sem_tr_fold)\n        x_te_eng = engine.transform(X_test_raw, X_semantic=sem_te)\n\n        fold_test_pred = np.zeros((len(test_data), len(cfg.TARGET_NAMES)), dtype=np.float32)\n\n        for k, target_name in enumerate(cfg.TARGET_NAMES):\n            # keep your behavior: clover always 0\n            if target_name == 'Dry_Clover_g':\n                fold_test_pred[:, k] = 0.0\n                continue\n\n            model = model_cls(**model_params)\n\n            # pass sample_weight when supported\n            try:\n                model.fit(x_tr_eng, y_tr[:, k], sample_weight=sw_tr)\n            except TypeError:\n                model.fit(x_tr_eng, y_tr[:, k])\n\n            pred_raw = model.predict(x_te_eng).astype(np.float32)\n            fold_test_pred[:, k] = pred_raw * target_max_arr[k]\n\n        y_pred_test_accum += fold_test_pred\n\n    return y_pred_test_accum / float(n_splits)\n\n\n# Model Parameters (same as your baseline)\nparams_cat = {\n    'iterations': 2000, 'learning_rate': 0.045, 'depth': 4, 'l2_leaf_reg': 0.56,\n    'random_strength': 0.045, 'bagging_temperature': 0.98, 'verbose': 0, 'random_state': 42,\n    'allow_writing_files': False\n}\nparams_xgb = {  # GradientBoostingRegressor proxy\n    'n_estimators': 1450, 'learning_rate': 0.010, 'max_depth': 3, 'subsample': 0.60,\n    'random_state': 42\n}\nparams_lgbm = {\n    'n_estimators': 900, 'learning_rate': 0.014, 'num_leaves': 48, 'min_child_samples': 19,\n    'subsample': 0.745, 'colsample_bytree': 0.745, 'reg_alpha': 0.21, 'reg_lambda': 3.78,\n    'verbose': -1, 'random_state': 42\n}\nparams_hist = {\n    'max_iter': 330, 'learning_rate': 0.05, 'max_depth': None, 'l2_regularization': 0.44,\n    'random_state': 42\n}\n\nfeat_engine = SupervisedEmbeddingEngine(n_pca=0.80, n_pls=8, n_gmm=6)\n\nprint(\"Training & Inferring Models (train is 3x augmented; test preds averaged across TTA)...\")\n\nprint(\"Model: HistGradientBoosting\")\npred_hist_aug = cross_validate_predict(\n    HistGradientBoostingRegressor, params_hist,\n    train_feat_df, test_feat_df, sem_train, sem_test, feat_engine\n)\n\nprint(\"Model: GradientBoosting\")\npred_gb_aug = cross_validate_predict(\n    GradientBoostingRegressor, params_xgb,\n    train_feat_df, test_feat_df, sem_train, sem_test, feat_engine\n)\n\nprint(\"Model: CatBoost\")\npred_cat_aug = cross_validate_predict(\n    CatBoostRegressor, params_cat,\n    train_feat_df, test_feat_df, sem_train, sem_test, feat_engine\n)\n\nprint(\"Model: LightGBM\")\npred_lgbm_aug = cross_validate_predict(\n    LGBMRegressor, params_lgbm,\n    train_feat_df, test_feat_df, sem_train, sem_test, feat_engine\n)\n\n\n# =========================================================================================\n# 7. ENSEMBLING & SUBMISSION (average across models, then average across TTA modes)\n# =========================================================================================\nprint(\"Ensembling, collapsing TTA, and post-processing...\")\n\n# Average across model families (still on augmented test rows)\nfinal_pred_aug = (pred_hist_aug + pred_gb_aug + pred_cat_aug + pred_lgbm_aug) / 4.0\n\n# Collapse augmented test back to base test size by averaging across TTA modes\nM = len(cfg.INFER_TTA_MODES)\nN_test = len(test_df_base)\nT = len(cfg.TARGET_NAMES)\n\n# We built test_feat_df as blocks per mode in order, so reshape works:\nfinal_pred = final_pred_aug.reshape(M, N_test, T).mean(axis=0).astype(np.float32)\n\n# assign to base test df (not augmented)\ntest_out = test_df_base.copy()\ntest_out[cfg.TARGET_NAMES] = final_pred\n\n# Mass-balance post-process (your behavior)\ntest_processed = post_process_biomass(test_out)\n\n# Create submission\nsub_df = melt_table(test_processed)\noutput_path = \"submission_siglip.csv\"\nsub_df.to_csv(output_path, index=False)\n\nprint(f\"✓ Submission generated: {output_path}\")\nprint(sub_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T05:20:01.140245Z","iopub.execute_input":"2026-01-16T05:20:01.140898Z","iopub.status.idle":"2026-01-16T05:34:03.683646Z","shell.execute_reply.started":"2026-01-16T05:20:01.140864Z","shell.execute_reply":"2026-01-16T05:34:03.683026Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dinov3","metadata":{}},{"cell_type":"code","source":"%%writefile mydino_infer.py\n# Config\nimport os, gc, math, cv2, re, numpy as np, pandas as pd\nfrom tqdm import tqdm\nimport torch, torch.nn as nn, torch.optim as optim\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nimport time, random\nfrom timm.utils import ModelEmaV2\nfrom sklearn.model_selection import StratifiedGroupKFold\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = True\ntorch.backends.cudnn.benchmark = True\n# (PyTorch 2.x)\ntorch.set_float32_matmul_precision(\"high\")\n\nALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\nGRAD_CHECKPOINTING = False\nR2_WEIGHTS      = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\nos.makedirs('/kaggle/CSIRO/working/baseline-ensemble_v2/', exist_ok=True)\nBASE_PATH       = '/kaggle/input/csiro-biomass'\nTRAIN_CSV       = os.path.join(BASE_PATH, 'train.csv')\nTRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\nTEST_IMAGE_DIR = '/kaggle/input/csiro-biomass/test'\nTEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\nSUBMISSION_DIR  = './'\nMODEL_DIR = '/kaggle/input/csiro-ensemble/pytorch/default/2/'\n\n\nclass DINOV3_CFG:\n    CREATE_SUBMISSION = True\n    USE_TQDM        = False\n    PRETRAINED_DIR  = None\n    PRETRAINED      = True\n    NSEEDS           = 1\n    SEED            = 82947501\n    N_FOLDS         = 5\n\n    MODEL_NAME      =  'vit_huge_plus_patch16_dinov3.lvd1689m'\n    BACKBONE_PATH   = None \n  \n    IMG_SIZE        = 512\n\n    VAL_TTA_TIMES   = 3\n    TTA_STEPS       = 2\n    \n    BATCH_SIZE      = 4\n    GRAD_ACC        = 1\n    NUM_WORKERS     = 16\n    EPOCHS          = 100\n    FREEZE_EPOCHS   = 3\n    WARMUP_EPOCHS   = 3\n    LR_REST         = 5e-4\n    LR_BACKBONE     = 1e-5\n    WD              = 1e-2\n    EMA_DECAY       = 0.9\n    CLIP_GRAD_NORM  = 1.0\n    PATIENCE        = 7\n    TARGET_COLS     = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS    = ['Dry_Clover_g', 'Dry_Dead_g']\n\n    LOSS_WEIGHTS    = np.array([0.1, 0.1, 0.1, 0.0, 0.0])\n    DEVICE          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n    # Aux heads + aux losses (single switch for all aux tasks)\n    USE_AUX = False \n    LAMBDA_HEIGHT = 0.2   # height regression aux loss weight\n    LAMBDA_NDVI   = 70.0   # NDVI regression aux loss weight\n    LAMBDA_STATE  = 0.35   # state classification aux loss weight\n    END_AUX_DECAY_EPOCH = 7\n\n    # Head regularization\n    HEAD_DROPOUT = 0.1\n\n\n\n# Transforms (surgical update: accept int OR (H, W) tuple)\n\n\nSIGLIP2_MEAN = [0.5, 0.5, 0.5]\nSIGLIP2_STD  = [0.5, 0.5, 0.5]\n\nIMAGENET_MEAN = [0.485, 0.456, 0.406]\nIMAGENET_STD  = [0.229, 0.224, 0.225]\n\ndef _hw(img_size):\n    # img_size can be 512 or (448, 224)\n    if isinstance(img_size, (tuple, list)):\n        assert len(img_size) == 2, \"img_size tuple must be (H, W)\"\n        return int(img_size[0]), int(img_size[1])\n    return int(img_size), int(img_size)\n\ndef clean_image(img):\n    h, w = img.shape[:2]\n    img = img[0:int(h * 0.90), :]\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n    return img\n\n\n\ndef get_tta_transforms(img_size=512, mode=0, mean=IMAGENET_MEAN, std=IMAGENET_STD):\n    \"\"\"\n    Stereo-friendly TTA:\n      - Uses ReplayCompose(additional_targets={\"image_right\":\"image\"}) so geometry matches.\n      - If mode includes HorizontalFlip, your dataset's replay-based swap will kick in.\n\n    Modes (suggested):\n      0: identity\n      1: horizontal flip (panorama-consistent when dataset swaps halves)\n      2: vertical flip\n      3: hflip + vflip (== 180° rot; dataset will swap halves because hflip present)\n      4: rotate +10°\n      5: rotate -10°\n      6: rotate +90°\n      7: rotate -90°\n    \"\"\"\n    h, w = _hw(img_size)\n\n    t = [A.Resize(h, w)]\n\n    if mode == 1:\n        t.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        t.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        t.append([\n             A.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, p=1.0),\n             A.Gamma(gamma_limit=(90, 110), p=1.0), \n        ])\n    elif mode == 4:\n        t.append(A.Rotate(limit=(10, 10), p=1.0,\n                          interpolation=cv2.INTER_LINEAR,\n                          border_mode=cv2.BORDER_REFLECT_101))\n    elif mode == 5:\n        t.append(A.Rotate(limit=(-10, -10), p=1.0,\n                          interpolation=cv2.INTER_LINEAR,\n                          border_mode=cv2.BORDER_REFLECT_101))\n    elif mode == 6:\n        t.append(A.Rotate(limit=(90, 90), p=1.0,\n                          interpolation=cv2.INTER_LINEAR,\n                          border_mode=cv2.BORDER_REFLECT_101))\n    elif mode == 7:\n        t.append(A.Rotate(limit=(-90, -90), p=1.0,\n                          interpolation=cv2.INTER_LINEAR,\n                          border_mode=cv2.BORDER_REFLECT_101))\n\n    t += [\n        A.Normalize(mean=mean, std=std, max_pixel_value=255.0),\n        ToTensorV2(),\n    ]\n\n    return A.ReplayCompose(\n        t,\n        additional_targets={\"image_right\": \"image\"},\n        p=1.0\n    )\n\n\n\n\n# Layers\nclass Local2DTokenMixerBlock(nn.Module):\n    \"\"\"\n    2D local token mixing on patch grid:\n    LN -> gated -> depthwise Conv2D -> proj -> dropout -> residual\n    Expects x_grid: (B, H, W, D)\n    \"\"\"\n    def __init__(self, dim, kernel_size=5, dropout=0.1):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        self.gate = nn.Linear(dim, dim)\n        self.dwconv = nn.Conv2d(\n            dim, dim,\n            kernel_size=kernel_size,\n            padding=kernel_size // 2,\n            groups=dim,\n            bias=True\n        )\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x_grid):\n        shortcut = x_grid\n        x = self.norm(x_grid)\n        x = x * torch.sigmoid(self.gate(x))\n        x = x.permute(0, 3, 1, 2)   # (B, D, H, W)\n        x = self.dwconv(x)\n        x = x.permute(0, 2, 3, 1)   # (B, H, W, D)\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\n# Model\nclass BiomassModel(nn.Module):\n    \"\"\"\n    Unified BiomassModel that supports:\n      - token backbones (ViT/DINO/SigLIP/DeiT3, etc.) via backbone_mode=\"tokens\"\n      - feature-map backbones (ConvNeXt/MaxViT/MaxxViT, etc.) via backbone_mode=\"features\"\n\n    Shared idea:\n      left/right -> backbone -> fuse as a wider 2D grid -> 2x Local2DTokenMixerBlock -> pool -> heads\n      heads predict (green, clover, dead) with Softplus, then:\n        gdm   = green + clover\n        total = gdm + dead\n\n    cls_mode:\n      - \"none\"     : no cls tokens at all (recommended; simplest & avoids fake cls compute)\n      - \"learned\"  : 2 learnable tokens (one per half) concatenated before pooling\n      - \"backbone\" : (tokens mode only) use backbone's first token from each half as cls\n    \"\"\"\n\n    def __init__(\n        self,\n        model_name: str,\n        pretrained: bool = True,\n        backbone_path: str | None = None,\n        use_aux: bool = False,\n        num_states: int = 4,\n        head_dropout: float = 0.1,\n        # NEW:\n        backbone_mode: str = \"tokens\",   # \"tokens\" (ViT) or \"features\" (ConvNeXt/MaxViT)\n        fuse_dim: int | None = None,     # optional projection dim for feature backbones\n    ):\n        super().__init__()\n        self.model_name = model_name\n        self.backbone_path = backbone_path\n        self.use_aux = bool(use_aux)\n        self.num_states = int(num_states)\n\n        self.backbone_mode = str(backbone_mode)\n\n        if self.backbone_mode not in (\"tokens\", \"features\"):\n            raise ValueError(f\"backbone_mode must be 'tokens' or 'features', got: {self.backbone_mode}\")\n\n        # ----------------------------\n        # 1) Backbone\n        # ----------------------------\n        if self.backbone_mode == \"features\":\n            # ConvNeXt/MaxViT family: returns list of feature maps with features_only=True\n            self.backbone = timm.create_model(\n                self.model_name,\n                pretrained=pretrained,\n                features_only=True,\n            )\n            in_ch = int(self.backbone.feature_info.channels()[-1])\n            self.nf = int(fuse_dim or in_ch)\n            self.proj_feat = nn.Identity() if self.nf == in_ch else nn.Conv2d(in_ch, self.nf, kernel_size=1, bias=False)\n\n            self.num_prefix = 0  # not used in features mode\n            self.gh = self.gw = None  # not used in features mode\n\n        else:\n            # Token backbones (ViT-like)\n            self.backbone = timm.create_model(\n                self.model_name,\n                pretrained=pretrained,\n                num_classes=0,\n            )\n\n            if hasattr(self.backbone, \"set_grad_checkpointing\") and GRAD_CHECKPOINTING:\n                self.backbone.set_grad_checkpointing(True)\n                print(\"✓ Gradient Checkpointing enabled\")\n\n            self.nf = int(self.backbone.num_features)\n            self.num_prefix = int(getattr(self.backbone, \"num_prefix_tokens\", 1))\n\n            # cache grid if available (good for square inputs)\n            gh = gw = None\n            if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"grid_size\"):\n                gs = self.backbone.patch_embed.grid_size\n                if isinstance(gs, (tuple, list)) and len(gs) == 2:\n                    gh, gw = int(gs[0]), int(gs[1])\n            self.gh, self.gw = gh, gw\n            if (self.gh is None) or (self.gw is None):\n                print(\"Warning: backbone.patch_embed.grid_size not found; will infer grid from input or factor Np at runtime.\")\n\n            self.proj_feat = None  # not used\n\n        # 2) Optional local checkpoint/backbone load\n        if pretrained:\n            self._load_backbone()\n\n        # ----------------------------\n        # 3) 2D fusion on a channels-last grid (B,H,W,D)\n        # ----------------------------\n        self.fusion2d = nn.Sequential(\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n        )\n\n        # ----------------------------\n        # 4) Pooling + LN\n        # ----------------------------\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        self.feat_ln = nn.LayerNorm(self.nf)\n\n        # ----------------------------\n        # 5) Main heads (positive)\n        # ----------------------------\n        def make_pos_head():\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf // 2, 1),\n                nn.Softplus(),\n            )\n\n        self.head_green  = make_pos_head()\n        self.head_clover = make_pos_head()\n        self.head_dead   = make_pos_head()\n\n        # ----------------------------\n        # 6) Aux heads\n        # ----------------------------\n        def make_reg_head():\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf // 2, 1),\n            )\n\n        def make_cls_head(n_classes: int):\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf // 2, n_classes),\n            )\n\n        self.aux_height = make_reg_head() if self.use_aux else None\n        self.aux_ndvi   = make_reg_head() if self.use_aux else None\n        self.aux_state  = make_cls_head(self.num_states) if (self.use_aux and self.num_states >= 2) else None\n\n        self._init_heads()\n\n    # ----------------------------\n    # Init heads\n    # ----------------------------\n    def _init_heads(self):\n        def init_module(m: nn.Module):\n            if isinstance(m, nn.Linear):\n                nn.init.trunc_normal_(m.weight, std=0.02)\n                if m.bias is not None:\n                    nn.init.zeros_(m.bias)\n\n        for mod in [self.head_green, self.head_clover, self.head_dead,\n                    self.aux_height, self.aux_ndvi, self.aux_state]:\n            if mod is not None:\n                mod.apply(init_module)\n\n    # ----------------------------\n    # Loading backbone weights\n    # ----------------------------\n    def _load_backbone(self):\n        try:\n            if self.backbone_path and os.path.exists(self.backbone_path):\n                print(f\"Loading backbone weights from local file: {self.backbone_path}\")\n                sd = torch.load(self.backbone_path, map_location=\"cpu\")\n                if isinstance(sd, dict) and \"model\" in sd:\n                    sd = sd[\"model\"]\n                elif isinstance(sd, dict) and \"state_dict\" in sd:\n                    sd = sd[\"state_dict\"]\n            else:\n                # Let timm handle pretrained internally; only do extra here if you really need to.\n                # Keeping this for compatibility with your previous flow.\n                if self.backbone_mode == \"features\":\n                    sd = timm.create_model(self.model_name, pretrained=True, features_only=True).state_dict()\n                else:\n                    sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool=\"\").state_dict()\n\n            # If token backbone and pos_embed mismatches, safest is to drop it (avoid load crash).\n            if self.backbone_mode == \"tokens\":\n                if \"pos_embed\" in sd and hasattr(self.backbone, \"pos_embed\"):\n                    pe_ck = sd[\"pos_embed\"]\n                    pe_m  = self.backbone.pos_embed\n                    if hasattr(pe_ck, \"shape\") and hasattr(pe_m, \"shape\") and pe_ck.shape != pe_m.shape:\n                        print(f\"pos_embed mismatch {tuple(pe_ck.shape)} -> {tuple(pe_m.shape)}; dropping pos_embed from checkpoint\")\n                        sd.pop(\"pos_embed\", None)\n\n            self.backbone.load_state_dict(sd, strict=False)\n            print(\"Pretrained weights loaded.\")\n        except Exception as e:\n            print(f\"Warning: pretrained load failed: {e}\")\n\n    # ----------------------------\n    # Token utilities (token mode)\n    # ----------------------------\n    @staticmethod\n    def _infer_grid_hw(num_patches: int) -> tuple[int, int]:\n        gh = int(math.isqrt(num_patches))\n        while gh > 1 and (num_patches % gh) != 0:\n            gh -= 1\n        gw = num_patches // gh\n        if gh * gw != num_patches:\n            raise ValueError(f\"Cannot factor num_patches={num_patches} into (gh*gw).\")\n        return gh, gw\n\n    def _grid_from_input(self, x: torch.Tensor, num_patches: int) -> tuple[int, int]:\n        \"\"\"\n        Robust for rectangular halves (e.g. 448x224):\n        compute grid from input H/W and patch size if possible, else fallback.\n        \"\"\"\n        if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"patch_size\"):\n            ps = self.backbone.patch_embed.patch_size\n            ph, pw = (ps if isinstance(ps, (tuple, list)) else (ps, ps))\n            gh = int(x.shape[-2] // ph)\n            gw = int(x.shape[-1] // pw)\n            if gh * gw == num_patches:\n                return gh, gw\n        return self._infer_grid_hw(num_patches)\n\n    @staticmethod\n    def _concat_halves_as_full_grid(p_l: torch.Tensor, p_r: torch.Tensor, gh: int, gw: int) -> torch.Tensor:\n        \"\"\"\n        p_l, p_r: (B, Np, D) where Np = gh*gw (for a single half)\n        returns: (B, gh*(2*gw), D) representing fused left-right panorama grid\n        \"\"\"\n        B, Np, D = p_l.shape\n        assert p_r.shape == (B, Np, D)\n        assert Np == gh * gw\n\n        pl2 = p_l.view(B, gh, gw, D)\n        pr2 = p_r.view(B, gh, gw, D)\n        full2 = torch.cat([pl2, pr2], dim=2)          # (B, gh, 2*gw, D)\n        return full2.reshape(B, gh * (2 * gw), D)     # (B, gh*(2*gw), D)\n\n    # ----------------------------\n    # Backbone forward\n    # ----------------------------\n    def _forward_backbone(self, x: torch.Tensor):\n        \"\"\"\n        Returns either:\n          - tokens: (B, N, D)  if backbone_mode=\"tokens\"\n          - fmap:   (B, C, H, W) if backbone_mode=\"features\"\n        \"\"\"\n        if self.backbone_mode == \"features\":\n            feats = self.backbone(x)  # list of feature maps\n            if isinstance(feats, (list, tuple)):\n                return feats[-1]\n            return feats\n        else:\n            return self.backbone.forward_features(x)\n\n    # ----------------------------\n    # Forward\n    # ----------------------------\n    def forward(self, left: torch.Tensor, right: torch.Tensor):\n        x_l = self._forward_backbone(left)\n        x_r = self._forward_backbone(right)\n\n        # ---------------------------------------------------------------------\n        # FEATURE MAP MODE (ConvNeXt/MaxViT): x_* = (B,C,H,W)\n        # ---------------------------------------------------------------------\n        if x_l.ndim == 4:\n            f_l = self.proj_feat(x_l)  # (B,D,H,W)\n            f_r = self.proj_feat(x_r)  # (B,D,H,W)\n\n            # B,H,W,D for your fusion2d\n            p_l = f_l.permute(0, 2, 3, 1).contiguous()\n            p_r = f_r.permute(0, 2, 3, 1).contiguous()\n\n            with torch.amp.autocast(\"cuda\", enabled=False):\n                p_grid = torch.cat([p_l.float(), p_r.float()], dim=2)  # (B,H,2W,D)\n                p_grid = self.fusion2d(p_grid)\n                B, H, W2, D = p_grid.shape\n                p_fused = p_grid.view(B, H * W2, D)                    # (B,N,D)\n                toks = p_fused.float()\n                feat = self.pool(toks.transpose(1, 2)).squeeze(-1)     # (B,D)\n                feat = self.feat_ln(feat)\n\n        # ---------------------------------------------------------------------\n        # TOKEN MODE (ViT/DINO): x_* = (B,N,D)\n        # ---------------------------------------------------------------------\n        else:\n            # patches\n            p_l = x_l[:, self.num_prefix:, :]\n            p_r = x_r[:, self.num_prefix:, :]\n            Np = int(p_l.size(1))\n\n            # infer grid for a half\n            gh, gw = self.gh, self.gw\n            if (gh is None) or (gw is None) or (gh * gw != Np):\n                gh, gw = self._grid_from_input(left, Np)\n\n            p_seq = self._concat_halves_as_full_grid(p_l, p_r, gh, gw)  # (B, gh*(2*gw), D)\n            B, N, D = p_seq.shape\n            H, W = gh, 2 * gw\n\n            # cls tokens\n            cls_l = x_l[:, :1, :]\n            cls_r = x_r[:, :1, :]\n            cls_q = torch.cat([cls_l, cls_r], dim=1)               # (B,2,D)\n\n            with torch.amp.autocast(\"cuda\", enabled=False):\n                p_grid  = p_seq.float().view(B, H, W, D).contiguous()\n                p_grid  = self.fusion2d(p_grid)\n                p_fused = p_grid.view(B, H * W, D)                     # (B,N,D)\n\n                if cls_q is None:\n                    toks = p_fused.float()\n                else:\n                    toks = torch.cat([cls_q.float(), p_fused.float()], dim=1)\n\n                feat = self.pool(toks.transpose(1, 2)).squeeze(-1)\n                feat = self.feat_ln(feat)\n\n        # heads + mass balance\n        green  = self.head_green(feat)\n        clover = self.head_clover(feat)\n        dead   = self.head_dead(feat)\n\n        gdm   = green + clover\n        total = gdm + dead\n\n        if not self.use_aux:\n            return total, gdm, green, clover, dead\n\n        aux = {\n            \"height\": self.aux_height(feat) if self.aux_height is not None else None,\n            \"Pre_GSHH_NDVI\": self.aux_ndvi(feat) if self.aux_ndvi is not None else None,\n            \"state_logits\": self.aux_state(feat) if self.aux_state is not None else None,\n        }\n        return total, gdm, green, clover, dead, aux\n\n\ndef seed_everything(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\n# inference utils\n\nclass BiomassTestDataset(Dataset):\n    \"\"\"\n    Returns raw numpy halves so we can apply Albumentations TTA on the fly.\n    \"\"\"\n    def __init__(self, img_dir):\n        self.img_dir = img_dir\n        self.paths = sorted([\n            os.path.join(img_dir, f) for f in os.listdir(img_dir)\n            if f.lower().endswith((\".jpg\", \".jpeg\", \".png\"))\n        ])\n        self.filenames = [os.path.basename(p) for p in self.paths]\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        path = self.paths[idx]\n        img = cv2.imread(path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n\n        mid = img.shape[1] // 2\n        left = img[:, :mid].copy()\n        right = img[:, mid:].copy()\n        return left, right, self.filenames[idx]\n\n\ndef _collate_raw_np(batch):\n    # keep numpy arrays as lists (avoid default torch tensor collation + extra copies)\n    lefts, rights, fns = zip(*batch)\n    return list(lefts), list(rights), list(fns)\n\n\ndef _postprocess_3_to_5(preds_direct_3: np.ndarray) -> np.ndarray:\n    # preds_direct_3: (N,3) [total, gdm, green]\n    pred_total = preds_direct_3[:, 0]\n    pred_gdm   = preds_direct_3[:, 1]\n    pred_green = preds_direct_3[:, 2]\n\n    pred_clover = np.maximum(0.0, pred_gdm - pred_green)\n    pred_dead   = np.maximum(0.0, pred_total - pred_gdm)\n\n    # order: ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    return np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1).astype(np.float32)\n\n\ndef _make_submission_from_preds(test_csv_path, all_target_cols, preds_5, filenames, out_csv_path):\n    test_df = pd.read_csv(test_csv_path)\n\n    # match image_path prefix if test has \"test/xxx.jpg\"\n    test_path_example = test_df[\"image_path\"].iloc[0]\n    if \"/\" in test_path_example:\n        prefix = test_path_example.rsplit(\"/\", 1)[0] + \"/\"\n        corrected_filenames = [prefix + fn for fn in filenames]\n    else:\n        corrected_filenames = filenames\n\n    preds_wide = pd.DataFrame(preds_5, columns=all_target_cols)\n    preds_wide.insert(0, \"image_path\", corrected_filenames)\n\n    preds_long = preds_wide.melt(\n        id_vars=[\"image_path\"],\n        value_vars=all_target_cols,\n        var_name=\"target_name\",\n        value_name=\"target\",\n    )\n\n    submission = pd.merge(\n        test_df[[\"sample_id\", \"image_path\", \"target_name\"]],\n        preds_long,\n        on=[\"image_path\", \"target_name\"],\n        how=\"left\",\n    )[[\"sample_id\", \"target\"]]\n\n    if submission[\"target\"].isna().any():\n        submission.loc[submission[\"target\"].isna(), \"target\"] = 0.0\n\n    submission = submission.sort_values(\"sample_id\").reset_index(drop=True)\n    submission.to_csv(out_csv_path, index=False)\n    return submission\n\n\ndef _replay_hflip_applied(replay: dict) -> bool:\n    if not replay:\n        return False\n    for t in replay.get(\"transforms\", []):\n        name = t.get(\"__class_fullname__\", \"\") or t.get(\"__class__\", \"\")\n        if name.endswith(\"HorizontalFlip\") and t.get(\"applied\", False):\n            return True\n    return False\n\n\n@torch.inference_mode()\ndef _predict_fold_tta_batched(\n    model: nn.Module,\n    loader: DataLoader,\n    device: torch.device,\n    img_size,\n    tta_steps: int = 1,\n    use_amp: bool = True,\n    amp_dtype: torch.dtype = torch.bfloat16,\n    prefix: str = '',\n) -> np.ndarray:\n    \"\"\"\n    Batched TTA using YOUR get_tta_transforms(img_size, mode).\n\n    Stereo-friendly:\n      - applies ONE ReplayCompose to (left,right) together per sample\n      - if HorizontalFlip was applied, swaps halves (panorama-consistent)\n    \"\"\"\n    model.eval()\n\n    N = len(loader.dataset)\n    preds_sum = np.zeros((N, 3), dtype=np.float32)\n\n    tta_steps = int(tta_steps)\n    tta_steps = max(1, min(tta_steps, 8))  # allow your new modes 0..7\n\n    if 'siglip' in prefix.lower():\n        transform_mean = SIGLIP2_MEAN\n        transform_std  = SIGLIP2_STD\n    else:\n        transform_mean = IMAGENET_MEAN\n        transform_std  = IMAGENET_STD\n\n    tta_tfms = [\n        (m, get_tta_transforms(img_size, mode=m, mean=transform_mean, std=transform_std))\n        for m in range(tta_steps)\n    ]\n    nb = (device.type == \"cuda\")\n\n    for mode, tfm in tta_tfms:\n        offset = 0\n        for left_list, right_list, _ in loader:\n            bs = len(left_list)\n\n            lt_list, rt_list = [], []\n            for im_l, im_r in zip(left_list, right_list):\n                out = tfm(image=im_l, image_right=im_r)   # ✅ ONE replay for both halves\n                lt = out[\"image\"]\n                rt = out[\"image_right\"]\n\n                # ✅ swap halves if HFlip actually happened (mode 1 and mode 3, etc.)\n                if _replay_hflip_applied(out.get(\"replay\", {})):\n                    lt, rt = rt, lt\n\n                lt_list.append(lt)\n                rt_list.append(rt)\n\n            lt = torch.stack(lt_list, dim=0)\n            rt = torch.stack(rt_list, dim=0)\n\n            lt = lt.to(device, non_blocking=nb)\n            rt = rt.to(device, non_blocking=nb)\n\n            if nb and use_amp:\n                with torch.autocast(\"cuda\", dtype=amp_dtype):\n                    outs = model(lt, rt)\n            else:\n                outs = model(lt, rt)\n\n            total, gdm, green = outs[0], outs[1], outs[2]\n            total = total.view(-1)\n            gdm   = gdm.view(-1)\n            green = green.view(-1)\n\n            pred3 = torch.stack([total, gdm, green], dim=1).float().cpu().numpy()\n            preds_sum[offset:offset + bs] += pred3\n            offset += bs\n\n            del lt, rt, outs, total, gdm, green, pred3, lt_list, rt_list\n\n    preds_sum /= float(tta_steps)\n    return preds_sum.astype(np.float32)\n\n\n\ndef ensemble_infer_to_submission(\n    cfgs,                   # list of cfg objects/namespaces\n    model_fns,              # list of callables: (cfg, weight_path)->nn.Module (loads weights on CPU)\n    prefixes,               # list[str] subfolder under MODEL_DIR where folds are saved\n    weights,                # list[float] model-level weights\n    out_csv_path: str,\n\n    # speed knobs\n    infer_batch_size: int = 8,\n    num_workers: int = 2,\n    tta_steps_default: int = 1,\n\n    # device/amp\n    device=None,                 # if None: auto CUDA if available else CPU\n    use_amp: bool = True,\n    amp_dtype: torch.dtype = torch.bfloat16,\n    pin_memory: bool = True,\n):\n    \"\"\"\n    Reads test images once, runs (model x folds x TTA) efficiently, writes submission.\n    - test_dir + test_csv are taken from global constants\n    - ALL_TARGET_COLS assumed global\n    - device auto-detects CUDA\n    \"\"\"\n    assert len(cfgs) == len(model_fns) == len(prefixes) == len(weights)\n\n    # reference cfg (common paths live here)\n    test_dir = TEST_IMAGE_DIR\n    test_csv = TEST_CSV\n    all_target_cols = ALL_TARGET_COLS  # noqa: F821\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    # normalize weights\n    w = np.asarray(weights, dtype=np.float32)\n    w = w / (w.sum() + 1e-12)\n\n    dataset = BiomassTestDataset(test_dir)\n    filenames = dataset.filenames.copy()\n\n    loader = DataLoader(\n        dataset,\n        batch_size=int(infer_batch_size),\n        shuffle=False,\n        num_workers=int(num_workers),\n        pin_memory=bool(pin_memory) and (device.type == \"cuda\"),\n        persistent_workers=(int(num_workers) > 0),\n        prefetch_factor=4 if int(num_workers) > 0 else None,\n        collate_fn=_collate_raw_np,\n    )\n\n    ensemble_sum = np.zeros((len(dataset), 3), dtype=np.float32)\n\n    for mi, (cfg, build_loaded_model, prefix, mw) in enumerate(zip(cfgs, model_fns, prefixes, w)):\n        model_dir = os.path.join(MODEL_DIR, prefix) if prefix else MODEL_DIR\n\n        # tuple-safe img_size (int or (H,W))\n        img_size = getattr(cfg, \"IMG_SIZE\", getattr(cfg, \"img_size\", 512))\n        tta_steps = int(getattr(cfg, \"VAL_TTA_TIMES\", getattr(cfg, \"TTA_STEPS\", tta_steps_default)))\n\n        # folds to use: default NSEEDS*N_FOLDS\n        n_seeds = int(getattr(cfg, \"NSEEDS\", getattr(cfg, \"N_SEEDS\", 1)) or 1)\n        n_folds = int(getattr(cfg, \"N_FOLDS\", getattr(cfg, \"NSPLITS\", getattr(cfg, \"N_SPLITS\", 5))) or 5)\n        folds_to_use = list(range(n_seeds * n_folds))\n\n        model_preds_sum = np.zeros((len(dataset), 3), dtype=np.float32)\n        successful_folds = 0\n\n        print(\"\\n\" + \"=\" * 80)\n        print(\n            f\"[MODEL {mi}] dir={model_dir}\\n\"\n            f\"  name={getattr(cfg,'MODEL_NAME',None)}  weight={float(mw):.4f}\\n\"\n            f\"  img_size={img_size}  TTA={tta_steps}  folds={len(folds_to_use)}  bs={infer_batch_size}  nw={num_workers}\\n\"\n            f\"  device={device}\"\n        )\n        print(\"=\" * 80)\n\n        for fold in folds_to_use:\n            weight_path = os.path.join(model_dir, prefix, f\"best_model_fold{fold}.pth\")\n            if not os.path.exists(weight_path):\n                continue\n\n            model = build_loaded_model(cfg, weight_path)\n            model.to(device)\n            model.eval()\n\n            fold_preds3 = _predict_fold_tta_batched(\n                model=model,\n                loader=loader,\n                device=device,\n                img_size=img_size,\n                tta_steps=tta_steps,\n                use_amp=use_amp,\n                amp_dtype=amp_dtype,\n                prefix=prefix,\n            )\n\n            model_preds_sum += fold_preds3\n            successful_folds += 1\n\n            # cleanup per fold\n            del model, fold_preds3\n            if device.type == \"cuda\":\n                torch.cuda.empty_cache()\n            gc.collect()\n\n        if successful_folds == 0:\n            raise FileNotFoundError(f\"[MODEL {mi}] No weights found in {model_dir}\")\n\n        model_preds = model_preds_sum / float(successful_folds)\n        ensemble_sum += float(mw) * model_preds\n\n        print(f\"[MODEL {mi}] used {successful_folds} folds, added to ensemble.\")\n\n        del model_preds_sum, model_preds\n        gc.collect()\n\n    preds_5 = _postprocess_3_to_5(ensemble_sum)\n\n    sub = _make_submission_from_preds(\n        test_csv_path=test_csv,\n        all_target_cols=all_target_cols,\n        preds_5=preds_5,\n        filenames=filenames,\n        out_csv_path=out_csv_path,\n    )\n    print(f\"\\n✓ saved submission: {out_csv_path}\")\n    return sub\n\n\n# ----------------------------\n# example model_fn you provide per model\n# ----------------------------\ndef load_model_for_fold(cfg, weight_path):\n    \"\"\"\n    Must return a ready-to-run model with weights loaded (on CPU).\n    \"\"\"\n    model = BiomassModel(\n        cfg.MODEL_NAME,\n        pretrained=False,\n        backbone_path=None,\n        backbone_mode=getattr(cfg, \"BACKBONE_MODE\", \"tokens\"),\n        fuse_dim=getattr(cfg, \"FUSE_DIM\", None),\n        use_aux=False,\n    )\n\n    state = torch.load(weight_path, map_location=\"cpu\")\n    if isinstance(state, dict) and (\"model_state_dict\" in state or \"state_dict\" in state):\n        key = \"model_state_dict\" if \"model_state_dict\" in state else \"state_dict\"\n        sd = state[key]\n    else:\n        sd = state\n\n    model.load_state_dict(sd, strict=False)\n    return model\n\ndef load_model_for_foldv2(cfg, weight_path):\n    \"\"\"\n    Must return a ready-to-run model with weights loaded (on CPU).\n    \"\"\"\n    model = BiomassModelv2(\n        cfg.MODEL_NAME,\n        pretrained=False,\n        backbone_path=None,\n        backbone_mode=getattr(cfg, \"BACKBONE_MODE\", \"tokens\"),\n        fuse_dim=getattr(cfg, \"FUSE_DIM\", None),\n        use_aux=False,\n    )\n\n    state = torch.load(weight_path, map_location=\"cpu\")\n    if isinstance(state, dict) and (\"model_state_dict\" in state or \"state_dict\" in state):\n        key = \"model_state_dict\" if \"model_state_dict\" in state else \"state_dict\"\n        sd = state[key]\n    else:\n        sd = state\n\n    model.load_state_dict(sd, strict=False)\n    return model\n\n\n# --- FINAL ENSEMBLE INFERENCE + SUBMISSION ---\n\n# ensemble ALL models\nprefixes = [  \"dinov3\"]\ncv_weights = { 'dinov3':1.}\n\nweights  = [cv_weights[p] for p in prefixes]\nweights = np.array([cv_weights[p] for p in prefixes], dtype=float)\ns = weights.sum()\nif s == 0 or not np.isfinite(s):\n    weights = np.ones_like(weights) / len(weights)\nelse:\n    weights = weights / s  # sums to 1.0\n\nweights = weights.tolist()\n\n# cfgs aligned with prefixes\ncfgs      = [ DINOV3_CFG ]\nmodel_fns = [ load_model_for_fold  ] \n\nsub = ensemble_infer_to_submission(\n    cfgs=cfgs,\n    model_fns=model_fns,\n    prefixes=prefixes,   # folder prefixes\n    weights=weights,     # CV-proportional\n    out_csv_path=os.path.join(SUBMISSION_DIR, \"submission_mydino.csv\"),\n)\n\n\nprint(sub.head())","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python mydino_infer.py","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dinov3 distilled","metadata":{}},{"cell_type":"code","source":"%%writefile dino_distill_infer.py\n# =========================\n# CSIRO Biomass — TGGC-only inference (self-contained, Kaggle)\n# - no training\n# - no multi-model ensembling\n# - loads best_model_fold*.pth from one folder (student_tggc) and averages folds\n# =========================\n\nimport os, gc, math\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport timm\n\n\n# -------------------------\n# Config (edit paths if needed)\n# -------------------------\nclass CFG:\n    # competition data\n    BASE_PATH = \"/kaggle/input/csiro-biomass\"\n    TEST_IMAGE_DIR = f\"{BASE_PATH}/test\"\n    TEST_CSV = f\"{BASE_PATH}/test.csv\"\n\n    # weights dataset (add your weights as Kaggle Dataset input)\n    MODEL_DIR = \"/kaggle/input/dinov3-classroom/pytorch/default/5\"\n    PREFIX = \"student_tggc\"  # folder under MODEL_DIR with best_model_fold*.pth\n\n    # model\n    MODEL_NAME = \"vit_large_patch16_dinov3.lvd1689m\"\n    IMG_SIZE = 512\n    TTA_STEPS = 3\n    USE_AUX = True\n    HEAD_DROPOUT = 0.1\n\n    # loader\n    BATCH_SIZE = 8\n    NUM_WORKERS = 2\n\n    # folds\n    N_FOLDS = 5\n    NSEEDS = 1  # total folds used = N_FOLDS * NSEEDS\n\n    # output\n    SUBMISSION_PATH = \"submission_distill.csv\"\n\n\nALL_TARGET_COLS = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\nTAIL_TARGET_ORDER = ALL_TARGET_COLS[:]  # used by calibration below\n\n\n# -------------------------\n# OOF tail calibration params (from your notebook)\n# -------------------------\ndef make_oof_tail_calibration_params_optionA():\n    return {\n        \"Dry_Green_g\": dict(qL=2.52037, qR=55.0177, aL=0.654279, bL=0.520678, aR=1.153237, bR=-8.055886, widthL=5.25,  widthR=5.25,  alpha=0.350),\n        \"Dry_Dead_g\":  dict(qL=1.02866, qR=24.5511, aL=0.357310, bL=0.779024, aR=0.836502, bR= 5.284353, widthL=2.352, widthR=2.352, alpha=0.350),\n        \"Dry_Clover_g\":dict(qL=0.000149536, qR=23.3832,aL=1.750000, bL=0.062449, aR=1.093193, bR=-4.270560, widthL=2.338, widthR=2.338, alpha=0.350),\n        \"GDM_g\":       dict(qL=11.1627, qR=62.4147, aL=0.980426, bL=0.362193, aR=1.241384, bR=-10.000000, widthL=5.125, widthR=5.125, alpha=0.350),\n        \"Dry_Total_g\": dict(qL=19.5531, qR=75.7549, aL=1.154826, bL=-1.405406,aR=1.036987, bR= 0.191613, widthL=5.62,  widthR=5.62,  alpha=0.350),\n    }\n\n\n# -------------------------\n# Transforms (TTA)\n# -------------------------\nIMAGENET_MEAN = [0.485, 0.456, 0.406]\nIMAGENET_STD  = [0.229, 0.224, 0.225]\n\ndef _hw(img_size):\n    if isinstance(img_size, (tuple, list)):\n        return int(img_size[0]), int(img_size[1])\n    return int(img_size), int(img_size)\n\ndef clean_image(img):\n    h, w = img.shape[:2]\n    img = img[0:int(h * 0.90), :]\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n    return img\n\ndef get_tta_transforms(img_size=512, mode=0, mean=IMAGENET_MEAN, std=IMAGENET_STD):\n    h, w = _hw(img_size)\n    t = [A.Resize(h, w)]\n    if mode == 1:\n        t.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        t.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        t.append([\n             A.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, p=1.0),\n             A.Gamma(gamma_limit=(90, 110), p=1.0), \n        ])\n    t += [A.Normalize(mean=mean, std=std, max_pixel_value=255.0), ToTensorV2()]\n    return A.ReplayCompose(t, additional_targets={\"image_right\": \"image\"}, p=1.0)\n\n\n# -------------------------\n# Minimal layer used by TGGC student\n# -------------------------\nclass Local2DTokenMixerBlock(nn.Module):\n    def __init__(self, dim, kernel_size=5, dropout=0.1):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        self.gate = nn.Linear(dim, dim)\n        self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim, bias=True)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x_grid):\n        shortcut = x_grid\n        x = self.norm(x_grid)\n        x = x * torch.sigmoid(self.gate(x))\n        x = x.permute(0, 3, 1, 2)     # (B,D,H,W)\n        x = self.dwconv(x)\n        x = x.permute(0, 2, 3, 1)     # (B,H,W,D)\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\n# -------------------------\n# TGGC student model (images-only)\n# -------------------------\nclass BiomassStudentModelv2(nn.Module):\n    def __init__(\n        self,\n        model_name: str,\n        pretrained: bool = False,\n        head_dropout: float = 0.1,\n        use_siglip_attn_pool: bool = True,\n        gate_dc: bool = True,\n        gate_tau_dead_g: float | None = 1.25,\n        gate_tau_clover_g: float | None = 0.2,\n        gate_sharp: float = 0.02,\n        use_aux: bool = True,\n        aux_height_film: bool = True,\n        aux_height_film_strength: float = 0.1,\n        aux_height_film_clip: float = 0.30,\n        aux_height_film_dim: int = 16,\n        aux_height_film_w: tuple[float, float, float, float] = (1.0, 1.0, 0.8, 0.15),\n        aux_height_log1p: bool = True,\n    ):\n        super().__init__()\n        self.model_name = model_name\n        self.use_siglip_attn_pool = bool(use_siglip_attn_pool)\n\n        self.backbone = timm.create_model(self.model_name, pretrained=pretrained, num_classes=0)\n        self.nf = int(getattr(self.backbone, \"num_features\", None) or getattr(self.backbone, \"embed_dim\", None))\n        self.num_prefix = int(getattr(self.backbone, \"num_prefix_tokens\", 1))\n\n        gh = gw = None\n        if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"grid_size\"):\n            gs = self.backbone.patch_embed.grid_size\n            if isinstance(gs, (tuple, list)) and len(gs) == 2:\n                gh, gw = int(gs[0]), int(gs[1])\n        self.gh, self.gw = gh, gw\n\n        self._siglip_map_pool = (\n            self.use_siglip_attn_pool\n            and (getattr(self.backbone, \"global_pool\", None) == \"map\")\n            and hasattr(self.backbone, \"attn_pool\")\n            and (self.num_prefix == 0)\n        )\n\n        self.fusion2d = nn.Sequential(\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n            Local2DTokenMixerBlock(self.nf, kernel_size=5, dropout=0.1),\n        )\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        self.feat_ln = nn.LayerNorm(self.nf)\n\n        def make_pos_head():\n            return nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf // 2, 1),\n                nn.Softplus(),\n            )\n\n        # TGGC heads: total, gdm, green, clover_raw\n        self.head_1 = make_pos_head()\n        self.head_2 = make_pos_head()\n        self.head_3 = make_pos_head()\n        self.head_4 = make_pos_head()\n\n        # aux height + mild FiLM to TGGC heads\n        self.use_aux = bool(use_aux)\n        self.aux_height_log1p = bool(aux_height_log1p)\n        self.aux_height_film = bool(aux_height_film)\n        self.aux_height_film_strength = float(aux_height_film_strength)\n        self.aux_height_film_clip = float(aux_height_film_clip)\n\n        self.aux_height_head = None\n        self.height_film = None\n        if self.use_aux:\n            self.aux_height_head = nn.Sequential(\n                nn.Linear(self.nf, self.nf // 2),\n                nn.GELU(),\n                nn.Dropout(head_dropout),\n                nn.Linear(self.nf // 2, 1),\n                nn.Softplus(),\n            )\n            if self.aux_height_film:\n                self.height_film = nn.Sequential(\n                    nn.Linear(1, int(aux_height_film_dim)),\n                    nn.GELU(),\n                    nn.Linear(int(aux_height_film_dim), 4),  # total,gdm,green,clover\n                )\n                w = torch.tensor(aux_height_film_w, dtype=torch.float32).view(1, 4)\n                self.register_buffer(\"_height_film_w\", w, persistent=False)\n            else:\n                self.register_buffer(\"_height_film_w\", torch.ones(1, 4), persistent=False)\n        else:\n            self.register_buffer(\"_height_film_w\", torch.ones(1, 4), persistent=False)\n\n        self._init_heads()\n\n        self.gate_dc = bool(gate_dc)\n        self.gate_tau_dead_g = gate_tau_dead_g\n        self.gate_tau_clover_g = gate_tau_clover_g\n        self.gate_sharp = float(gate_sharp)\n\n    def _init_heads(self):\n        def init_linear(m):\n            if isinstance(m, nn.Linear):\n                nn.init.trunc_normal_(m.weight, std=0.02)\n                if m.bias is not None:\n                    nn.init.zeros_(m.bias)\n\n        for m in [self.head_1, self.head_2, self.head_3, self.head_4, self.aux_height_head, self.height_film]:\n            if m is not None:\n                m.apply(init_linear)\n\n    @staticmethod\n    def _infer_grid_hw(num_patches: int):\n        gh = int(math.isqrt(num_patches))\n        while gh > 1 and (num_patches % gh) != 0:\n            gh -= 1\n        gw = num_patches // gh\n        if gh * gw != num_patches:\n            raise ValueError(f\"Cannot factor num_patches={num_patches}.\")\n        return gh, gw\n\n    def _grid_from_input(self, x: torch.Tensor, num_patches: int):\n        if hasattr(self.backbone, \"patch_embed\") and hasattr(self.backbone.patch_embed, \"patch_size\"):\n            ps = self.backbone.patch_embed.patch_size\n            ph, pw = (ps if isinstance(ps, (tuple, list)) else (ps, ps))\n            gh = int(x.shape[-2] // ph)\n            gw = int(x.shape[-1] // pw)\n            if gh * gw == num_patches:\n                return gh, gw\n        return self._infer_grid_hw(num_patches)\n\n    @staticmethod\n    def _concat_halves_as_full_grid(p_l: torch.Tensor, p_r: torch.Tensor, gh: int, gw: int):\n        B, Np, D = p_l.shape\n        pl2 = p_l.view(B, gh, gw, D)\n        pr2 = p_r.view(B, gh, gw, D)\n        full2 = torch.cat([pl2, pr2], dim=2)\n        return full2.reshape(B, gh * (2 * gw), D)\n\n    def _siglip_attn_pool_cls(self, patch_tokens: torch.Tensor):\n        cls = self.backbone.attn_pool(patch_tokens)\n        if cls.ndim == 2:\n            cls = cls.unsqueeze(1)\n        return cls\n\n    def _forward_backbone(self, x: torch.Tensor):\n        x = self.backbone.forward_features(x)\n        if isinstance(x, (tuple, list)):\n            x = x[0]\n        if torch.is_tensor(x) and x.ndim == 2:\n            x = x.unsqueeze(1)\n        return x\n\n    @staticmethod\n    def _smooth_min(a: torch.Tensor, b: torch.Tensor):\n        return a - F.softplus(a - b)\n\n    def _soft_gate(self, x: torch.Tensor, thr: torch.Tensor, gamma: float = 2.0):\n        sharp = self.gate_sharp + 1e-6\n        return x * torch.sigmoid((x - thr) / sharp).pow(gamma)\n\n    def _apply_gate_dc(self, total, gdm, green, clover, dead, eps: float):\n        if (self.gate_tau_dead_g is None) or (self.gate_tau_clover_g is None):\n            return total, gdm, green, clover, dead\n\n        total_d = total.detach()\n\n        dead_ratio = dead / (total + eps)\n        dead_thr_r = (float(self.gate_tau_dead_g) / (total_d + eps))\n        dead_g = self._soft_gate(dead_ratio, dead_thr_r) * total\n\n        gdm_g = total - dead_g\n\n        gdm_d = gdm_g.detach()\n        clover_ratio = clover / (gdm_g + eps)\n        clover_thr_r = (float(self.gate_tau_clover_g) / (gdm_d + eps))\n        clover_g = self._soft_gate(clover_ratio, clover_thr_r) * gdm_g\n\n        green_g = gdm_g - clover_g\n        return total, gdm_g, green_g, clover_g, dead_g\n\n    def _apply_height_film_tggc(self, total_hat, gdm_hat, green_hat, clover_raw, height_pred):\n        if (self.height_film is None) or (height_pred is None):\n            return total_hat, gdm_hat, green_hat, clover_raw\n\n        h = height_pred.clamp_min(0.0)\n        if self.aux_height_log1p:\n            h = torch.log1p(h)\n        h = h.view(-1, 1)\n\n        d = self.height_film(h)  # (B,4)\n        scale = 1.0 + (self.aux_height_film_strength * self._height_film_w) * torch.tanh(d)\n        scale = torch.clamp(scale, 1.0 - self.aux_height_film_clip, 1.0 + self.aux_height_film_clip)\n\n        total_hat  = total_hat  * scale[:, 0:1]\n        gdm_hat    = gdm_hat    * scale[:, 1:2]\n        green_hat  = green_hat  * scale[:, 2:3]\n        clover_raw = clover_raw * scale[:, 3:4]\n        return total_hat, gdm_hat, green_hat, clover_raw\n\n    def forward(self, left: torch.Tensor, right: torch.Tensor):\n        x_l = self._forward_backbone(left)\n        x_r = self._forward_backbone(right)\n\n        if self._siglip_map_pool:\n            p_l = x_l\n            p_r = x_r\n            cls_l = self._siglip_attn_pool_cls(p_l)\n            cls_r = self._siglip_attn_pool_cls(p_r)\n        else:\n            p_l = x_l[:, self.num_prefix:, :]\n            p_r = x_r[:, self.num_prefix:, :]\n            cls_l = x_l[:, :self.num_prefix, :][:, :1, :]\n            cls_r = x_r[:, :self.num_prefix, :][:, :1, :]\n            if self.num_prefix == 0:\n                cls_l = p_l.mean(dim=1, keepdim=True)\n                cls_r = p_r.mean(dim=1, keepdim=True)\n\n        Np = int(p_l.size(1))\n        gh, gw = self.gh, self.gw\n        if (gh is None) or (gw is None) or (gh * gw != Np):\n            gh, gw = self._grid_from_input(left, Np)\n\n        p_seq = self._concat_halves_as_full_grid(p_l, p_r, gh, gw)\n        B, _, D = p_seq.shape\n        H, W = gh, 2 * gw\n\n        cls_q = torch.cat([cls_l, cls_r], dim=1)  # (B,2,D)\n\n        # keep mixer in fp32 (like your notebook)\n        with torch.amp.autocast(\"cuda\", enabled=False):\n            p_grid = p_seq.float().view(B, H, W, D).contiguous()\n            p_grid = self.fusion2d(p_grid)\n            p_fused = p_grid.view(B, H * W, D)\n\n            toks = torch.cat([cls_q.float(), p_fused.float()], dim=1)\n            feat = self.pool(toks.transpose(1, 2)).squeeze(-1)\n            feat = self.feat_ln(feat)\n\n        eps = 1e-6\n\n        height_pred = None\n        aux_dict = None\n        if self.use_aux and (self.aux_height_head is not None):\n            height_pred = self.aux_height_head(feat)\n            aux_dict = {\"height\": height_pred, \"Pre_GSHH_NDVI\": None, \"state_logits\": None}\n\n        total_hat  = self.head_1(feat)\n        gdm_hat    = self.head_2(feat)\n        green_hat  = self.head_3(feat)\n        clover_raw = self.head_4(feat)\n\n        if self.use_aux and self.aux_height_film and (height_pred is not None):\n            total_hat, gdm_hat, green_hat, clover_raw = self._apply_height_film_tggc(\n                total_hat, gdm_hat, green_hat, clover_raw, height_pred\n            )\n\n        gdm   = self._smooth_min(gdm_hat, total_hat)\n        green = self._smooth_min(green_hat, gdm)\n        max_clover = F.relu(gdm - green)\n        clover = self._smooth_min(clover_raw, max_clover)\n        dead   = total_hat - gdm\n        total  = total_hat\n\n        if self.gate_dc:\n            total, gdm, green, clover, dead = self._apply_gate_dc(total, gdm, green, clover, dead, eps)\n\n        if aux_dict is not None:\n            return total, gdm, green, clover, dead, aux_dict\n        return total, gdm, green, clover, dead\n\n\n# -------------------------\n# Submission utilities\n# -------------------------\ndef enforce_mass_balance(df_wide, fixed_clover=None):\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_rec = Y\n    else:\n        C = np.array([[1, 1, 0, -1, 0],\n                      [0, 0, 1,  1, -1]])\n        P = np.eye(5) - C.T @ np.linalg.inv(C @ C.T) @ C\n        Y_rec = P @ Y\n\n    Y_rec = np.maximum(0, Y_rec.T)\n    out = df_wide.copy()\n    out[ordered_cols] = Y_rec\n    return out\n\ndef _sigmoid_np(z):\n    z = np.clip(z, -60.0, 60.0)\n    return 1.0 / (1.0 + np.exp(-z))\n\ndef apply_smooth_tail_calibration_optionA(\n    preds_2d: np.ndarray,\n    params_by_target: dict,\n    *,\n    target_order=TAIL_TARGET_ORDER,\n    use_left_tail: bool = False,\n    use_right_tail: bool = True,\n    skip_targets=(),\n    clip_nonneg: bool = True,\n) -> np.ndarray:\n    yp = np.asarray(preds_2d, dtype=np.float64).copy()\n    skip_set = set(skip_targets) if skip_targets is not None else set()\n\n    for j, tname in enumerate(target_order):\n        if tname in skip_set:\n            continue\n        pj = params_by_target.get(tname)\n        if pj is None:\n            continue\n\n        alpha  = float(pj.get(\"alpha\", 1.0))\n        qL     = float(pj.get(\"qL\", np.nan))\n        qR     = float(pj.get(\"qR\", np.nan))\n        aL     = float(pj.get(\"aL\", 1.0)); bL = float(pj.get(\"bL\", 0.0))\n        aR     = float(pj.get(\"aR\", 1.0)); bR = float(pj.get(\"bR\", 0.0))\n        widthL = max(float(pj.get(\"widthL\", 1.0)), 1e-6)\n        widthR = max(float(pj.get(\"widthR\", 1.0)), 1e-6)\n\n        x = yp[:, j]\n        m = np.isfinite(x)\n        if not m.any():\n            continue\n        x0 = x[m]\n\n        delta = 0.0\n        if use_left_tail and np.isfinite(qL):\n            wL = _sigmoid_np((qL - x0) / widthL)\n            yL = aL * x0 + bL\n            delta = delta + wL * (yL - x0)\n        if use_right_tail and np.isfinite(qR):\n            wR = _sigmoid_np((x0 - qR) / widthR)\n            yR = aR * x0 + bR\n            delta = delta + wR * (yR - x0)\n\n        yp[m, j] = x0 + alpha * delta\n\n    if clip_nonneg:\n        yp = np.maximum(yp, 0.0)\n    return yp.astype(np.float32)\n\ndef apply_tail_calibration_df(\n    df_wide,\n    params_by_target: dict,\n    *,\n    target_order=TAIL_TARGET_ORDER,\n    use_left_tail=False,\n    use_right_tail=True,\n    skip_targets=(),\n    clip_nonneg=True,\n):\n    arr = df_wide[target_order].to_numpy(np.float32)\n    arr_cal = apply_smooth_tail_calibration_optionA(\n        arr, params_by_target,\n        target_order=target_order,\n        use_left_tail=use_left_tail,\n        use_right_tail=use_right_tail,\n        skip_targets=skip_targets,\n        clip_nonneg=clip_nonneg,\n    )\n    out = df_wide.copy()\n    out[target_order] = arr_cal\n    return out\n\ndef _postprocess_3_to_5(preds_direct_3: np.ndarray) -> np.ndarray:\n    # preds_direct_3: (N,3) [total, gdm, green]\n    pred_total = preds_direct_3[:, 0]\n    pred_gdm   = preds_direct_3[:, 1]\n    pred_green = preds_direct_3[:, 2]\n    pred_clover = np.maximum(0.0, pred_gdm - pred_green)\n    pred_dead   = np.maximum(0.0, pred_total - pred_gdm)\n    # order: ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    return np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1).astype(np.float32)\n\ndef _make_submission_from_preds(test_csv_path, preds_5, filenames, out_csv_path):\n    test_df = pd.read_csv(test_csv_path)\n\n    test_path_example = test_df[\"image_path\"].iloc[0]\n    if \"/\" in test_path_example:\n        prefix = test_path_example.rsplit(\"/\", 1)[0] + \"/\"\n        corrected_filenames = [prefix + fn for fn in filenames]\n    else:\n        corrected_filenames = filenames\n\n    preds_wide = pd.DataFrame(preds_5, columns=ALL_TARGET_COLS)\n    preds_wide.insert(0, \"image_path\", corrected_filenames)\n\n    params = make_oof_tail_calibration_params_optionA()\n    preds_wide = apply_tail_calibration_df(\n        preds_wide,\n        params,\n        use_left_tail=False,\n        use_right_tail=True,\n        skip_targets=[\"Dry_Clover_g\"],\n    )\n    preds_wide[\"Dry_Clover_g\"] *= 0.8\n    preds_wide = enforce_mass_balance(preds_wide)\n\n    preds_long = preds_wide.melt(\n        id_vars=[\"image_path\"],\n        value_vars=ALL_TARGET_COLS,\n        var_name=\"target_name\",\n        value_name=\"target\",\n    )\n\n    submission = pd.merge(\n        test_df[[\"sample_id\", \"image_path\", \"target_name\"]],\n        preds_long,\n        on=[\"image_path\", \"target_name\"],\n        how=\"left\",\n    )[[\"sample_id\", \"target\"]]\n\n    submission[\"target\"] = submission[\"target\"].fillna(0.0)\n    submission = submission.sort_values(\"sample_id\").reset_index(drop=True)\n    submission.to_csv(out_csv_path, index=False)\n    return submission\n\n\n# -------------------------\n# Test dataset + inference\n# -------------------------\nclass BiomassTestDataset(Dataset):\n    def __init__(self, img_dir):\n        self.img_dir = img_dir\n        self.paths = sorted([\n            os.path.join(img_dir, f) for f in os.listdir(img_dir)\n            if f.lower().endswith((\".jpg\", \".jpeg\", \".png\"))\n        ])\n        self.filenames = [os.path.basename(p) for p in self.paths]\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        path = self.paths[idx]\n        img = cv2.imread(path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        mid = img.shape[1] // 2\n        left = img[:, :mid].copy()\n        right = img[:, mid:].copy()\n        return left, right, self.filenames[idx]\n\ndef _collate_raw_np(batch):\n    lefts, rights, fns = zip(*batch)\n    return list(lefts), list(rights), list(fns)\n\n@torch.inference_mode()\ndef _predict_fold_tta_batched(\n    model: nn.Module,\n    loader: DataLoader,\n    device: torch.device,\n    img_size,\n    tta_steps: int,\n    use_amp: bool,\n    amp_dtype: torch.dtype,\n) -> np.ndarray:\n    model.eval()\n    N = len(loader.dataset)\n    preds_sum = np.zeros((N, 3), dtype=np.float32)\n    tta_steps = max(1, min(int(tta_steps), 4))\n\n    tta_tfms = [get_tta_transforms(img_size, mode=m) for m in range(tta_steps)]\n    nb = (device.type == \"cuda\")\n\n    for tfm in tta_tfms:\n        offset = 0\n        for left_list, right_list, _ in loader:\n            bs = len(left_list)\n            lt = torch.stack([tfm(image=im)[\"image\"] for im in left_list], dim=0)\n            rt = torch.stack([tfm(image=im)[\"image\"] for im in right_list], dim=0)\n            lt = lt.to(device, non_blocking=nb)\n            rt = rt.to(device, non_blocking=nb)\n\n            if nb and use_amp:\n                with torch.autocast(\"cuda\", dtype=amp_dtype):\n                    outs = model(lt, rt)\n            else:\n                outs = model(lt, rt)\n\n            total, gdm, green = outs[0].view(-1), outs[1].view(-1), outs[2].view(-1)\n            pred3 = torch.stack([total, gdm, green], dim=1).float().cpu().numpy()\n            preds_sum[offset:offset + bs] += pred3\n            offset += bs\n\n            del lt, rt, outs, total, gdm, green, pred3\n\n    preds_sum /= float(tta_steps)\n    return preds_sum.astype(np.float32)\n\ndef load_model(weight_path: str):\n    model = BiomassStudentModelv2(\n        model_name=CFG.MODEL_NAME,\n        pretrained=False,\n        head_dropout=float(CFG.HEAD_DROPOUT),\n        use_aux=bool(CFG.USE_AUX),\n    )\n    state = torch.load(weight_path, map_location=\"cpu\")\n    if isinstance(state, dict) and (\"model_state_dict\" in state or \"state_dict\" in state):\n        key = \"model_state_dict\" if \"model_state_dict\" in state else \"state_dict\"\n        sd = state[key]\n    else:\n        sd = state\n    model.load_state_dict(sd, strict=False)\n    return model\n\n\n# -------------------------\n# Run\n# -------------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\namp_dtype = torch.bfloat16 if (device.type == \"cuda\" and torch.cuda.is_bf16_supported()) else torch.float16\nuse_amp = (device.type == \"cuda\")\n\ntest_ds = BiomassTestDataset(CFG.TEST_IMAGE_DIR)\ntest_loader = DataLoader(\n    test_ds,\n    batch_size=int(CFG.BATCH_SIZE),\n    shuffle=False,\n    num_workers=int(CFG.NUM_WORKERS),\n    pin_memory=(device.type == \"cuda\"),\n    persistent_workers=(int(CFG.NUM_WORKERS) > 0),\n    collate_fn=_collate_raw_np,\n)\n\nmodel_dir = os.path.join(CFG.MODEL_DIR, CFG.PREFIX)\nfold_paths = []\nfor fold in range(int(CFG.N_FOLDS) * int(CFG.NSEEDS)):\n    p = os.path.join(model_dir, f\"best_model_fold{fold}.pth\")\n    if os.path.exists(p):\n        fold_paths.append(p)\n\nif len(fold_paths) == 0:\n    raise FileNotFoundError(f\"No fold weights found in: {model_dir}\")\n\nprint(f\"Device={device} | AMP={use_amp} ({amp_dtype})\")\nprint(f\"Found {len(fold_paths)} folds in {model_dir}\")\nprint(f\"Test images: {len(test_ds)}\")\n\npreds_sum = np.zeros((len(test_ds), 3), dtype=np.float32)\n\nfor i, wp in enumerate(fold_paths):\n    print(f\"Fold {i+1}/{len(fold_paths)}: {os.path.basename(wp)}\")\n    model = load_model(wp).to(device).eval()\n    preds_sum += _predict_fold_tta_batched(\n        model=model,\n        loader=test_loader,\n        device=device,\n        img_size=CFG.IMG_SIZE,\n        tta_steps=CFG.TTA_STEPS,\n        use_amp=use_amp,\n        amp_dtype=amp_dtype,\n    )\n    del model\n    gc.collect()\n    if device.type == \"cuda\":\n        torch.cuda.empty_cache()\n\npreds3 = preds_sum / float(len(fold_paths))\npreds5 = _postprocess_3_to_5(preds3)\n\nsub = _make_submission_from_preds(\n    test_csv_path=CFG.TEST_CSV,\n    preds_5=preds5,\n    filenames=test_ds.filenames,\n    out_csv_path=CFG.SUBMISSION_PATH,\n)\n\nprint(f\"\\n✓ Saved: {CFG.SUBMISSION_PATH}\")\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python dino_distill_infer.py","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Ensemble","metadata":{}},{"cell_type":"code","source":"# mix, siglip + dinov3 + mvp\n\nimport pandas as pd\nimport numpy as np\nimport os\n\n# =============================================================================\n\n# Target definitions required for Mass Balance\nALL_TARGETS = [\"Dry_Green_g\", \"Dry_Clover_g\", \"Dry_Dead_g\", \"GDM_g\", \"Dry_Total_g\"]\n\n# =============================================================================\n# HELPER FUNCTIONS\n# =============================================================================\n\ndef enforce_mass_balance(\n    df_wide: pd.DataFrame,\n    *,\n    fixed_clover: bool = False,\n    ordered_cols=(\"Dry_Green_g\", \"Dry_Clover_g\", \"Dry_Dead_g\", \"GDM_g\", \"Dry_Total_g\"),\n    # optional: how strongly to keep each component close (larger = keep closer)\n    # default None => unweighted (matches your old behavior)\n    weights=None,\n    # safety / postprocess\n    clip_min: float = 0.0,\n    clip_max: float | None = None,\n    # if you clip, you can re-impose exact equalities by recomputing dependent vars\n    # (useful when fixed_clover=True)\n    recompute_after_clip: bool = False,\n    eps: float = 1e-12,\n) -> pd.DataFrame:\n    \"\"\"\n    Enforce biological equalities:\n      1) Green + Clover - GDM = 0\n      2) Dead + GDM - Total  = 0\n\n    Modes\n    -----\n    fixed_clover=False:\n        Orthogonal projection of (G,C,D,GDM,T) onto the constraint subspace.\n        Optionally weighted via diagonal weights.\n\n    fixed_clover=True:\n        Clover is held fixed. We solve the constraints by adjusting the remaining\n        variables with minimal (weighted) change.\n\n    Notes\n    -----\n    - If you apply nonneg clipping after projection, exact equalities may no longer hold.\n      If you need exact equalities after clipping, set recompute_after_clip=True.\n    \"\"\"\n    cols = list(ordered_cols)\n    X = df_wide[cols].to_numpy(dtype=np.float64, copy=True)  # (N,5)\n    if X.ndim != 2 or X.shape[1] != 5:\n        raise ValueError(f\"Expected df_wide[{cols}] to be (N,5), got {X.shape}\")\n\n    # constraints: C y = 0, with y = [G,C,D,GDM,T]\n    C = np.array([\n        [1.0, 1.0, 0.0, -1.0,  0.0],   # G + C - GDM = 0\n        [0.0, 0.0, 1.0,  1.0, -1.0],   # D + GDM - T = 0\n    ], dtype=np.float64)  # (2,5)\n\n    # weights -> metric: minimize ||W^(1/2) (y-x)||^2 s.t. C y = 0\n    # if weights provided as \"keep-closer\" weights, we convert to inverse-variance metric.\n    if weights is None:\n        W = np.eye(5, dtype=np.float64)\n    else:\n        w = np.asarray(weights, dtype=np.float64).reshape(-1)\n        if w.size != 5:\n            raise ValueError(\"weights must have length 5 (for G,C,D,GDM,T)\")\n        w = np.maximum(w, eps)\n        # metric matrix for quadratic: (y-x)^T W (y-x)\n        W = np.diag(w)\n\n    # Precompute pieces for projection\n    # y = x - W^{-1} C^T (C W^{-1} C^T)^{-1} C x    (weighted orthogonal projection)\n    Winv = np.diag(1.0 / np.diag(W))\n\n    if not fixed_clover:\n        A = C @ Winv @ C.T  # (2,2)\n        # robust inverse for 2x2\n        try:\n            Ainv = np.linalg.inv(A)\n        except np.linalg.LinAlgError:\n            Ainv = np.linalg.pinv(A)\n\n        # apply to all rows vectorized:\n        # correction = Winv C^T Ainv (C x^T)  -> (5,N)\n        CXt = (C @ X.T)                      # (2,N)\n        corr = (Winv @ C.T @ (Ainv @ CXt))   # (5,N)\n        Y = (X.T - corr).T                   # (N,5)\n\n    else:\n        # Clover fixed: optimize over y_free = [G, D, GDM, T], Clover constant.\n        # Constraints become:\n        #   G - GDM = -Clover\n        #   D + GDM - T = 0\n        idx_free = np.array([0, 2, 3, 4], dtype=int)  # G, D, GDM, T\n        idx_c    = 1\n\n        Xf = X[:, idx_free]          # (N,4)\n        c0 = X[:, idx_c:idx_c+1]     # (N,1) Clover fixed\n\n        # Cf y_free = b, where b depends on clover\n        Cf = np.array([\n            [1.0, 0.0, -1.0,  0.0],   # G - GDM\n            [0.0, 1.0,  1.0, -1.0],   # D + GDM - T\n        ], dtype=np.float64)          # (2,4)\n\n        # b: first constraint equals -clover, second equals 0\n        b = np.concatenate([-c0, np.zeros_like(c0)], axis=1)  # (N,2)\n\n        # Weighted solve: y = x - W^{-1} Cf^T (Cf W^{-1} Cf^T)^{-1} (Cf x - b)\n        Wf = W[np.ix_(idx_free, idx_free)]\n        Winvf = np.diag(1.0 / np.diag(Wf))\n\n        A = Cf @ Winvf @ Cf.T\n        try:\n            Ainv = np.linalg.inv(A)\n        except np.linalg.LinAlgError:\n            Ainv = np.linalg.pinv(A)\n\n        r = (Cf @ Xf.T - b.T)                  # (2,N)\n        corr = (Winvf @ Cf.T @ (Ainv @ r))     # (4,N)\n        Yf = (Xf.T - corr).T                   # (N,4)\n\n        Y = X.copy()\n        Y[:, idx_free] = Yf\n        Y[:, idx_c] = X[:, idx_c]  # keep clover fixed exactly\n\n    # clip (optional)\n    if clip_min is not None:\n        Y = np.maximum(Y, float(clip_min))\n    if clip_max is not None:\n        Y = np.minimum(Y, float(clip_max))\n\n    # if you clipped and want exact equalities again, recompute dependent vars\n    # (this is deterministic and avoids projection drift after clamping)\n    if recompute_after_clip:\n        # safest \"dependency\" direction is:\n        #   GDM = G + C\n        #   T   = D + GDM\n        Y[:, 3] = Y[:, 0] + Y[:, 1]\n        Y[:, 4] = Y[:, 2] + Y[:, 3]\n        if clip_min is not None:\n            Y = np.maximum(Y, float(clip_min))\n        if clip_max is not None:\n            Y = np.minimum(Y, float(clip_max))\n\n    out = df_wide.copy()\n    out.loc[:, cols] = Y.astype(np.float32, copy=False)\n    return out\n\n\n\ndef robust_ensemble(file_paths, weights, clover_source=\"dino\"):\n    \"\"\"\n    Ensemble targets with weights across available models.\n    Special-case: Dry_Clover_g uses `clover_source` only (default: dino).\n    \"\"\"\n    print(\"--- Starting Ensemble ---\")\n    print(f\"Weights (raw): {weights}\")\n    print(f\"NOTE: Using {clover_source.upper()}-only for Dry_Clover_g\")\n\n    # ---- load all dfs into dict by name ----\n    dfs = {}\n    for name, path in file_paths.items():\n        if not os.path.exists(path):\n            raise FileNotFoundError(f\"Missing file: {path}\")\n        df = pd.read_csv(path).sort_values(\"sample_id\").reset_index(drop=True)\n        dfs[name] = df\n        print(f\"Loaded {name}: {len(df)} rows\")\n\n    # ---- alignment check (all sample_id identical & ordered) ----\n    base_ids = next(iter(dfs.values()))[\"sample_id\"]\n    for name, df in dfs.items():\n        if not df[\"sample_id\"].equals(base_ids):\n            raise ValueError(f\"Sample IDs do not match between submission files! Mismatch at: {name}\")\n\n    # ---- split sample_id -> image_id / target_name for each model ----\n    dfs_split = {}\n    for name, df in dfs.items():\n        df_split = df.copy()\n        df_split[[\"image_id\", \"target_name\"]] = df_split[\"sample_id\"].str.rsplit(\"__\", n=1, expand=True)\n        dfs_split[name] = df_split\n\n    # ---- determine which models are actually used (intersection of weights & files) ----\n    used_models = [m for m in weights.keys() if m in dfs_split]\n    if len(used_models) == 0:\n        raise ValueError(\"No usable models found. Check `weights` keys vs `FILES` keys.\")\n\n    # ---- validate clover source ----\n    if clover_source not in dfs_split:\n        raise ValueError(f\"clover_source='{clover_source}' not found in provided FILES.\")\n\n    # ---- normalize weights once (for non-clover targets) ----\n    w = np.array([float(weights[m]) for m in used_models], dtype=np.float64)\n    if not np.isfinite(w).all() or w.sum() <= 0:\n        raise ValueError(f\"Bad weights: {weights}\")\n    w = w / w.sum()\n    w_map = {m: float(wi) for m, wi in zip(used_models, w)}\n    print(f\"Weights (normalized over {used_models}): {w_map}\")\n\n    # ---- per-target ensembling ----\n    ensemble_results = []\n    for target in ALL_TARGETS:\n        if target == \"Dry_Clover_g\":\n            src_df = dfs_split[clover_source]\n            out = src_df[src_df[\"target_name\"] == target][[\"sample_id\", \"target\"]].copy()\n            ensemble_results.append(out)\n            continue\n\n        # Start from first model then merge others on sample_id\n        merged = None\n        for m in used_models:\n            mdf = dfs_split[m]\n            mt = mdf[mdf[\"target_name\"] == target][[\"sample_id\", \"target\"]].rename(\n                columns={\"target\": f\"{m}_pred\"}\n            )\n            merged = mt if merged is None else pd.merge(merged, mt, on=\"sample_id\", how=\"inner\")\n\n        # Weighted average across models\n        pred_cols = [f\"{m}_pred\" for m in used_models]\n        pred_mat = merged[pred_cols].to_numpy(np.float64)\n        w_vec = np.array([w_map[m] for m in used_models], dtype=np.float64)\n\n        merged[\"target\"] = (pred_mat * w_vec[None, :]).sum(axis=1).astype(np.float32)\n        ensemble_results.append(merged[[\"sample_id\", \"target\"]])\n\n    # ---- combine long ----\n    ensemble_df = pd.concat(ensemble_results, ignore_index=True)\n    print(\"Weighted average complete (special-case Clover).\")\n\n    # ---- long -> wide for mass balance ----\n    ensemble_df[[\"image_id\", \"target_name\"]] = ensemble_df[\"sample_id\"].str.rsplit(\"__\", n=1, expand=True)\n    wide_df = ensemble_df.pivot(index=\"image_id\", columns=\"target_name\", values=\"target\").reset_index()\n\n    # --- your post-hoc tweaks (kept) ---\n    if \"Dry_Clover_g\" in wide_df.columns:\n        wide_df[\"Dry_Clover_g\"] *= 0.8\n        clover = wide_df[\"Dry_Clover_g\"]\n        wide_df.loc[clover<0.3,\"Dry_Clover_g\"] *= 0.\n\n    if \"Dry_Dead_g\" in wide_df.columns:\n        dead = wide_df[\"Dry_Dead_g\"]\n        wide_df.loc[dead > 20, \"Dry_Dead_g\"] *= 1.1\n        wide_df.loc[dead < 10, \"Dry_Dead_g\"] *= 0.9\n\n    # ---- enforce constraints (Clover fixed) ----\n    print(\"Applying Mass Balance Constraints (Dry_Clover_g fixed)...\")\n    wide_balanced = enforce_mass_balance(wide_df, fixed_clover=True)\n\n    # ---- wide -> long ----\n    long_balanced = wide_balanced.melt(\n        id_vars=\"image_id\",\n        value_vars=ALL_TARGETS,\n        var_name=\"target_name\",\n        value_name=\"target\",\n    )\n    long_balanced[\"sample_id\"] = long_balanced[\"image_id\"] + \"__\" + long_balanced[\"target_name\"]\n\n    final_submission = long_balanced[[\"sample_id\", \"target\"]].sort_values(\"sample_id\").reset_index(drop=True)\n    return final_submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T05:35:28.524915Z","iopub.execute_input":"2026-01-16T05:35:28.525156Z","iopub.status.idle":"2026-01-16T05:35:28.583547Z","shell.execute_reply.started":"2026-01-16T05:35:28.525128Z","shell.execute_reply":"2026-01-16T05:35:28.583034Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# EXECUTION\n# =============================================================================\nif __name__ == \"__main__\":\n\n    ensemble_weights = {\n        \"siglip\": 0.35,\n        \"dino\":   0.35,\n        \"dino_distill\": 0.3,\n    }\n    \n    FILES = {\n        \"siglip\": \"submission_siglip.csv\",\n        \"dino\":   \"submission_mydino.csv\",\n        \"dino_distill\":    \"submission_distill.csv\",\n    }\n    \n    OUTPUT_FILE = \"submission.csv\"\n      \n    try:\n        submission = robust_ensemble(FILES, ensemble_weights, clover_source=\"dino\")\n        submission.to_csv(OUTPUT_FILE, index=False)\n        print(f\"\\nSuccess! Saved to {OUTPUT_FILE}\")\n        print(submission.head())\n\n        print(\"\\nStats:\")\n        print(submission[\"target\"].describe())\n\n    except Exception as e:\n        print(f\"\\nError during ensembling: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T05:35:28.524915Z","iopub.execute_input":"2026-01-16T05:35:28.525156Z","iopub.status.idle":"2026-01-16T05:35:28.583547Z","shell.execute_reply.started":"2026-01-16T05:35:28.525128Z","shell.execute_reply":"2026-01-16T05:35:28.583034Z"}},"outputs":[],"execution_count":null}]}