{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Image2Biomass — TRAIN notebook (Kaggle, PYTORCH + TIMM ViT, IMAGE-ONLY)\n#\n# Why image-only: official test.csv only has (sample_id, image_path, target_name), no NDVI/Height/Species.\n# We still use train-only fields (State/Sampling_Date) ONLY for grouped CV splits.\n#\n# Goal: swap the CNN baseline for the ViT backbones that most strong public notebooks use\n# (e.g., ViT-L/14 CLIP, DINO-like ViTs), via timm.\n\n# %% [0] Setup\nimport os, gc, random, math\nfrom dataclasses import dataclass\n\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\n\n# Optional: OpenCV for image cleaning (timestamp/artifact removal)\ntry:\n    import cv2\n    _HAS_CV2 = True\nexcept Exception as _e:\n    cv2 = None\n    _HAS_CV2 = False\n    print(\"[warn] cv2 not available; image cleaning disabled:\", repr(_e))\n\nfrom torchvision import transforms as tvt\n\nimport timm\nfrom timm.data import resolve_data_config, create_transform\n\n# ------------------\n# Paths / constants\n# ------------------\nDATA_DIR = \"/kaggle/input/csiro-biomass\"\nTRAIN_CSV = os.path.join(DATA_DIR, \"train.csv\")\nTEST_CSV  = os.path.join(DATA_DIR, \"test.csv\")\n\nOUT_DIR = \"/kaggle/working/i2b\"\nos.makedirs(OUT_DIR, exist_ok=True)\nMODELS_DIR = os.path.join(OUT_DIR, \"models\")\nos.makedirs(MODELS_DIR, exist_ok=True)\n\n# Reuse folds across runs to ensure apples-to-apples comparisons.\nREUSE_FOLDS = True\nFOLDS_PATH = os.path.join(OUT_DIR, \"folds.csv\")\n\nFOLD_SEED = 42\n# Optional: multi-seed training (average OOF in grams space). Set to [42] for single-seed.\nRUN_SEEDS = [42, 1337]  # add more seeds if you have time budget\nN_FOLDS = 5\n\n# ---- Model choices (start here)\n# Note: timm model names available on Kaggle can differ by timm version.\n# If a name is unknown, we'll auto-fallback to a known-good model.\n# Prefer smaller backbones first (helps reduce overfit + runtime).\n# NOTE: availability depends on timm version + whether weights are present offline.\nMODEL_CANDIDATES = [\n    \"vit_base_patch16_clip_224.openai\",\n    \"vit_base_patch16_224.augreg_in21k\",\n    \"vit_small_patch16_224.augreg_in21k\",\n]\nFALLBACK_MODEL_NAME = \"vit_large_patch14_clip_224.openai\"  # known-good on this setup\n\ndef pick_valid_timm_model(name: str, fallback: str = FALLBACK_MODEL_NAME) -> str:\n    name = str(name)\n    if timm.is_model(name):\n        return name\n\n    # Try to find close matches (prefix / contains)\n    matches = []\n    try:\n        matches = timm.list_models(f\"{name}*\")\n        if not matches:\n            matches = timm.list_models(f\"*{name}*\")\n    except Exception:\n        matches = []\n\n    if matches:\n        print(f\"[warn] Unknown timm model: {name}\")\n        print(\"       Using closest match:\", matches[0])\n        return matches[0]\n\n    # Heuristic: user often types swinv2_base_window12_384 (without suffix)\n    if name == \"swinv2_base_window12_384\":\n        for cand in [\n            \"swinv2_base_window12to24_192to384.ms_in22k_ft_in1k\",\n            \"swinv2_base_window12to16_192to256.ms_in22k_ft_in1k\",\n            \"swinv2_base_window12_192.ms_in22k_ft_in1k\",\n        ]:\n            if timm.is_model(cand):\n                print(f\"[warn] Unknown timm model: {name}\")\n                print(\"       Using swinv2 candidate:\", cand)\n                return cand\n\n    print(f\"[warn] Unknown timm model: {name}\")\n    print(\"       Falling back to:\", fallback)\n    return fallback\n\ndef pick_first_available_model(candidates, fallback: str) -> str:\n    for cand in candidates:\n        # exact\n        if timm.is_model(cand):\n            return cand\n        # close matches\n        try:\n            m = timm.list_models(f\"{cand}*\")\n            if m:\n                return m[0]\n            m = timm.list_models(f\"*{cand}*\")\n            if m:\n                return m[0]\n        except Exception:\n            pass\n    return fallback\n\nMODEL_NAME = pick_first_available_model(MODEL_CANDIDATES, fallback=FALLBACK_MODEL_NAME)\nprint(\"Selected MODEL_NAME:\", MODEL_NAME)\n\n# ViTs have a preferred input size (often 224 or 384). We'll pull it from timm config.\n# You can still set a larger TRAIN_RESIZE for texture preservation before random-crop.\n# For multi-crop validation, it's important to match timm's expected eval resize/crop_pct.\nTRAIN_RESIZE = 448\nVAL_RESIZE = None  # for multi-crop/tiling eval: if None, uses timm-derived resize based on crop_pct (with a CLIP-safe override)\n\n# Validation / TTA policy\n#\n# Key idea: checkpoint selection should correlate with the *final* scoring run.\n# Because our final OOF uses multi-crop TTA, we also evaluate each epoch with\n# the same multi-crop settings (small dataset, so this is still fast enough).\n#\n# If you want faster training, set VAL_EVAL_MCROPS=1.\nVAL_EVAL_MCROPS = 1\nVAL_EVAL_USE_TIMM = True  # keep eval+TTA consistent (both use our tiling path)\n\n# Final fold OOF uses TTA multi-crop/tiling.\nVAL_TTA_MCROPS = 5\nVAL_TTA_MODE = \"five\"   # \"five\" (corners+center) or \"grid\" (true tiling)\nVAL_TTA_RESIZE = 448     # must be >= IMG_SIZE; 320/384 preserve texture but cost time\n\n# Back-compat (do not edit below)\nVAL_SINGLE_CROP_USES_TIMM = VAL_EVAL_USE_TIMM\n\n# Training knobs\nBATCH_SIZE = 16          # raise until OOM (ViT-L/14 @224 often 16-32 on T4 w/ AMP)\nEPOCHS = 18              # run longer; best-epoch checkpointing will pick the peak\nLR = 3e-5\nMIN_LR = 1e-6            # don't decay LR all the way to 0\nWARMUP_EPOCHS = 1        # small warmup helps avoid wild early negative R2 swings\nWEIGHT_DECAY = 0.10\n\n# Regularization / augmentation (small dataset + shortcut risk)\nHEAD_DROPOUT = 0.50\nUSE_COLORJITTER = True\nCJ_BRIGHTNESS = 0.20\nCJ_CONTRAST   = 0.20\nCJ_SATURATION = 0.12\nCJ_HUE        = 0.02\nGRAD_CLIP = 1.0\nAMP = True\nNUM_WORKERS = 4\nPATIENCE = 6              # early stop if no val improvement (re-enabled now that split+TTA are stable)\nES_START_EPOCH = 8         # don't early-stop before this epoch (late-bloom folds are common)\n\nUSE_LOG1P = True         # strongly recommended\nENFORCE_CONSTRAINTS_EVAL = True   # enforce GDM=DG+DC and TOTAL=DD+GDM during eval/OOF\n\n# ------------------\n# Optional image cleaning (often helps on this comp)\n# ------------------\n# Goal: remove shortcuts (orange timestamps / bottom artifacts) so the model focuses on pasture texture.\n# Tune the HSV range if your stamp color differs.\nUSE_CLEANING = True\nCLEAN_CROP_BOTTOM_PCT = 0.10      # crop bottom 10% (color charts / cardboard)\nCLEAN_INPAINT_ORANGE = True       # inpaint orange timestamp text\nCLEAN_ORANGE_LOWER = (5, 150, 150)\nCLEAN_ORANGE_UPPER = (25, 255, 255)\nCLEAN_DILATE_ITER = 2\nCLEAN_CLAHE = False               # lighting normalization (LAB-CLAHE)\nCLEAN_WHITEBALANCE = False        # simple grey-world white balance\n\n# Columns\nIMG_COL = \"image_path\"\nTNAME_COL = \"target_name\"\nTVAL_COL = \"target\"\nDATE_COL = \"Sampling_Date\"\nSTATE_COL = \"State\"\nGROUP_KEY_COL = \"group_key\"\n\nTARGET_COLS = [\n    \"Dry_Green_g\",\n    \"Dry_Dead_g\",\n    \"Dry_Clover_g\",\n    \"GDM_g\",\n    \"Dry_Total_g\",\n]\nTARGET_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.2, 0.5], dtype=np.float32)\n\n# Independent-head setup: predict only (Dry_Green_g, Dry_Dead_g, Dry_Clover_g)\n# and derive the remaining targets:\n#   GDM_g = Dry_Green_g + Dry_Clover_g\n#   Dry_Total_g = Dry_Dead_g + GDM_g\nINDEP_TARGET_COLS = TARGET_COLS[:3]\nFULL_OUT_DIM = len(TARGET_COLS)\nINDEP_OUT_DIM = len(INDEP_TARGET_COLS)\n\n# ------------------\n# Repro / device\n# ------------------\n\ndef seed_everything(seed: int = 42):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\nseed_everything(FOLD_SEED)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"torch:\", torch.__version__)\nprint(\"cuda available:\", torch.cuda.is_available())\nif torch.cuda.is_available():\n    print(\"GPU:\", torch.cuda.get_device_name(0))\n\ntorch.backends.cudnn.benchmark = True  # speed\n\n# %% [1] Load CSVs (LONG)\ntrain_long = pd.read_csv(TRAIN_CSV)\n\ntest_long = pd.read_csv(TEST_CSV)\n# NOTE: test_long has no metadata columns by design; keep image-only.\n\nneed_train = {IMG_COL, TNAME_COL, TVAL_COL}\nmissing_train = sorted(list(need_train - set(train_long.columns)))\nif missing_train:\n    raise ValueError(f\"train.csv missing columns: {missing_train}\")\n\nneed_test = {IMG_COL, TNAME_COL}\nmissing_test = sorted(list(need_test - set(test_long.columns)))\nif missing_test:\n    raise ValueError(f\"test.csv missing columns: {missing_test}\")\n\n# %% [2] Pivot train to WIDE (one row per image)\n\ny_wide = (\n    train_long\n    .pivot_table(index=IMG_COL, columns=TNAME_COL, values=TVAL_COL, aggfunc=\"first\")\n    .reset_index()\n)\n\n# carry over split helpers\nmeta_cols = [c for c in [DATE_COL, STATE_COL] if c in train_long.columns]\nif meta_cols:\n    img_meta = train_long.groupby(IMG_COL, as_index=False)[meta_cols].first()\n    train_wide = img_meta.merge(y_wide, on=IMG_COL, how=\"inner\")\nelse:\n    train_wide = y_wide\n\nfor t in TARGET_COLS:\n    if t not in train_wide.columns:\n        train_wide[t] = np.nan\ntrain_wide = train_wide.dropna(subset=TARGET_COLS).reset_index(drop=True)\nprint(\"train_long:\", train_long.shape, \"train_wide:\", train_wide.shape)\n\n# %% [3] Grouped folds (train-only)\nfrom sklearn.model_selection import GroupKFold\nimport hashlib, json\n\n# ------------------\n# Fold scheme (pick ONE and keep it)\n# ------------------\n# Strict + practical default for this competition:\n# - group = State + Sampling_Date (prevents leakage across collection events)\n# - greedy assignment balances folds a bit vs plain GroupKFold\nFOLD_SCHEME = \"state_date\"   # \"state_date\", \"date\", \"state_month\", \"random\", \"state\"\nFOLD_ASSIGN = \"greedy\"       # \"greedy\" (balanced) or \"gkf\" (plain GroupKFold)\n\nGREEDY_USE_LOG1P = True\nFOLD_INFO_PATH = os.path.join(OUT_DIR, \"folds_info.json\")\n\ndef fold_signature(df: pd.DataFrame) -> str:\n    \"\"\"Stable signature of (image_path -> fold) so logs are comparable across runs.\"\"\"\n    s = (\n        df[[IMG_COL, \"fold\"]]\n        .sort_values(IMG_COL)\n        .astype(str)\n        .agg(\"|\".join, axis=1)\n        .str.cat(sep=\"\\n\")\n    )\n    return hashlib.md5(s.encode(\"utf-8\")).hexdigest()\n\n\ndef _to_date_str(series: pd.Series) -> pd.Series:\n    d = pd.to_datetime(series, errors=\"coerce\")\n    return d.dt.strftime(\"%Y-%m-%d\").fillna(\"NA\")\n\n\ndef make_group_key(df: pd.DataFrame) -> pd.Series:\n    if FOLD_SCHEME == \"random\":\n        return df[IMG_COL].astype(str)\n\n    if FOLD_SCHEME == \"date\":\n        return _to_date_str(df[DATE_COL]) if DATE_COL in df.columns else df[IMG_COL].astype(str)\n\n    if FOLD_SCHEME == \"state_date\":\n        if STATE_COL in df.columns and DATE_COL in df.columns:\n            return df[STATE_COL].astype(str).fillna(\"NA\") + \"_\" + _to_date_str(df[DATE_COL])\n        if STATE_COL in df.columns:\n            return df[STATE_COL].astype(str).fillna(\"NA\")\n        return df[IMG_COL].astype(str)\n\n    if FOLD_SCHEME == \"state_month\":\n        if STATE_COL in df.columns and DATE_COL in df.columns:\n            d = pd.to_datetime(df[DATE_COL], errors=\"coerce\")\n            ym = d.dt.to_period(\"M\").astype(str).fillna(\"NA\")\n            return df[STATE_COL].astype(str).fillna(\"NA\") + \"_\" + ym\n        if STATE_COL in df.columns:\n            return df[STATE_COL].astype(str).fillna(\"NA\")\n        return df[IMG_COL].astype(str)\n\n    if FOLD_SCHEME == \"state\":\n        return df[STATE_COL].astype(str).fillna(\"NA\") if STATE_COL in df.columns else df[IMG_COL].astype(str)\n\n    return df[IMG_COL].astype(str)\n\n\ndef make_greedy_folds(df: pd.DataFrame, n_folds: int) -> pd.Series:\n    \"\"\"Deterministic greedy group->fold assignment for regression.\n\n    Keeps groups intact but balances a proxy target across folds.\n    \"\"\"\n    y = df[TARGET_COLS].values.astype(np.float32)\n    proxy = (y * TARGET_WEIGHTS.reshape(1, -1)).sum(axis=1)\n    if GREEDY_USE_LOG1P:\n        proxy = np.log1p(np.maximum(proxy, 0.0))\n\n    tmp = df[[GROUP_KEY_COL]].copy()\n    tmp[\"proxy\"] = proxy\n\n    g = tmp.groupby(GROUP_KEY_COL)\n    group_sum = g[\"proxy\"].sum().sort_values(ascending=False)\n    group_cnt = g.size().loc[group_sum.index]\n\n    fold_sum = np.zeros(n_folds, dtype=np.float64)\n    fold_cnt = np.zeros(n_folds, dtype=np.int64)\n\n    group_to_fold = {}\n    for gr in group_sum.index:\n        best_f = int(min(range(n_folds), key=lambda f: (fold_sum[f], fold_cnt[f])))\n        group_to_fold[gr] = best_f\n        fold_sum[best_f] += float(group_sum.loc[gr])\n        fold_cnt[best_f] += int(group_cnt.loc[gr])\n\n    return df[GROUP_KEY_COL].map(group_to_fold).astype(int)\n\n\ndef make_folds(df: pd.DataFrame, n_folds: int) -> pd.DataFrame:\n    out = df.copy()\n    out[GROUP_KEY_COL] = make_group_key(out)\n\n    if FOLD_SCHEME == \"random\":\n        rng = np.random.default_rng(FOLD_SEED)\n        out[\"fold\"] = rng.integers(0, n_folds, size=len(out))\n        return out\n\n    n_groups = int(out[GROUP_KEY_COL].nunique())\n    if n_groups < n_folds:\n        print(f\"[warn] only {n_groups} groups (<{n_folds}); falling back to random folds\")\n        rng = np.random.default_rng(FOLD_SEED)\n        out[\"fold\"] = rng.integers(0, n_folds, size=len(out))\n        return out\n\n    out = out.sort_values([GROUP_KEY_COL, IMG_COL]).reset_index(drop=True)\n\n    if FOLD_ASSIGN == \"greedy\":\n        out[\"fold\"] = make_greedy_folds(out, n_folds)\n        return out\n\n    out[\"fold\"] = -1\n    gkf = GroupKFold(n_splits=n_folds)\n    groups = out[GROUP_KEY_COL].values\n    for fold, (_, val_idx) in enumerate(gkf.split(out, groups=groups)):\n        out.loc[val_idx, \"fold\"] = fold\n    return out\n\n# --- build / reuse folds\ntrain_wide[GROUP_KEY_COL] = make_group_key(train_wide)\n\nif REUSE_FOLDS and os.path.exists(FOLDS_PATH):\n    folds_df = pd.read_csv(FOLDS_PATH)\n    merged = train_wide.merge(folds_df[[IMG_COL, \"fold\"]], on=IMG_COL, how=\"left\")\n    if merged[\"fold\"].isna().any():\n        print(\"[warn] Some images missing in saved folds.csv; regenerating folds.\")\n        train_wide = make_folds(train_wide, N_FOLDS)\n        train_wide[[IMG_COL, GROUP_KEY_COL, \"fold\"]].to_csv(FOLDS_PATH, index=False)\n    else:\n        train_wide = merged\n        print(\"Reused folds from:\", FOLDS_PATH)\nelse:\n    train_wide = make_folds(train_wide, N_FOLDS)\n    train_wide[[IMG_COL, GROUP_KEY_COL, \"fold\"]].to_csv(FOLDS_PATH, index=False)\n\nfold_counts = train_wide[\"fold\"].value_counts().sort_index().to_dict()\nprint(\"FOLD_SCHEME:\", FOLD_SCHEME, \"| FOLD_ASSIGN:\", FOLD_ASSIGN)\nprint(\"n_groups:\", int(train_wide[GROUP_KEY_COL].nunique()))\nprint(\"fold_counts:\", fold_counts)\nprint(\"FOLD_SIG:\", fold_signature(train_wide))\n\ntry:\n    info = {\n        \"fold_scheme\": FOLD_SCHEME,\n        \"fold_assign\": FOLD_ASSIGN,\n        \"n_folds\": int(N_FOLDS),\n        \"n_groups\": int(train_wide[GROUP_KEY_COL].nunique()),\n        \"fold_counts\": {str(k): int(v) for k, v in fold_counts.items()},\n        \"fold_sig\": fold_signature(train_wide),\n        \"seed\": int(FOLD_SEED),\n    }\n    with open(FOLD_INFO_PATH, \"w\") as f:\n        json.dump(info, f, indent=2)\n    print(\"Wrote:\", FOLD_INFO_PATH)\nexcept Exception as e:\n    print(\"[warn] could not write folds_info.json:\", repr(e))\n\n# %% [4] Target transforms + metric\n\ndef y_transform(y: np.ndarray) -> np.ndarray:\n    y = y.astype(np.float32)\n    if USE_LOG1P:\n        return np.log1p(np.maximum(y, 0.0))\n    return y\n\n\ndef y_inverse(y: np.ndarray) -> np.ndarray:\n    y = y.astype(np.float32)\n    if USE_LOG1P:\n        return np.expm1(y)\n    return y\n\n\ndef enforce_constraints_grams(y_grams: np.ndarray) -> np.ndarray:\n    \"\"\"Enforce domain identities (optional):\n    GDM = Dry_Green + Dry_Clover\n    Dry_Total = Dry_Dead + GDM\n\n    Input/Output: y_grams shape (N,5) in TARGET_COLS order.\n    \"\"\"\n    out = np.maximum(y_grams, 0.0).copy()\n    dg = out[:, 0]\n    dd = out[:, 1]\n    dc = out[:, 2]\n\n    gdm = dg + dc\n    total = dd + gdm\n\n    out[:, 3] = gdm\n    out[:, 4] = total\n    return out\n\n\ndef weighted_r2_global(y_true: np.ndarray, y_pred: np.ndarray, w: np.ndarray) -> float:\n    # y_true/y_pred: (N, T), w: (T,)\n    w2 = w.reshape(1, -1)\n    w_sum = float(np.sum(w2)) * y_true.shape[0]\n    y_bar = float(np.sum(w2 * y_true) / (w_sum + 1e-12))\n    sse = float(np.sum(w2 * (y_true - y_pred) ** 2))\n    sst = float(np.sum(w2 * (y_true - y_bar) ** 2))\n    return float(1.0 - (sse / (sst + 1e-12)))\n\n\ndef score_wR2(y_true_grams: np.ndarray, y_pred_model_space: np.ndarray) -> float:\n    y_pred_grams = np.maximum(y_inverse(y_pred_model_space), 0.0)\n    if ENFORCE_CONSTRAINTS_EVAL:\n        y_pred_grams = enforce_constraints_grams(y_pred_grams)\n    return weighted_r2_global(y_true_grams, y_pred_grams, TARGET_WEIGHTS)\n\n# %% [5] Timm transforms (matches pretrained weights)\n\n# IMPORTANT: If internet is disabled, timm may fail to download pretrained weights.\n# In that case, add the model weights via \"Add Input\" -> \"Models\" on Kaggle, OR switch\n# to a model that exists in the environment.\n\nPRETRAINED = True\ntry:\n    _tmp_model = timm.create_model(MODEL_NAME, pretrained=True)\n    _tmp_model = None\n    print(f\"Loaded pretrained weights for: {MODEL_NAME}\")\nexcept Exception as e:\n    PRETRAINED = False\n    print(\"[warn] timm pretrained load failed (likely no internet / weights not present).\")\n    print(\"       We'll run with pretrained=False (pipeline debug only).\")\n    print(\"       To use pretrained offline: Add the model via Kaggle Models / kagglehub.\")\n    print(\"       Error:\", repr(e))\n    print(\"[warn] timm pretrained load failed (likely no internet / weights not present).\")\n    print(\"       Add the weights via Kaggle Models, or set pretrained=False as a fallback.\")\n    print(\"       Error:\", repr(e))\n\n# Resolve data config for proper normalization\n# Build a lightweight model instance to resolve the correct preprocessing config\n_cfg_model = timm.create_model(MODEL_NAME, pretrained=False)\ncfg = resolve_data_config({}, model=_cfg_model)\nIMG_SIZE = int(cfg.get(\"input_size\", (3, 224, 224))[1])\ndel _cfg_model\nprint(\"Model IMG_SIZE:\", IMG_SIZE)\nprint(\"timm crop_pct:\", cfg.get(\"crop_pct\"))\n\n# Derive eval resize. NOTE: for some CLIP-style models timm reports crop_pct=1.0.\n# In this competition, using a *larger* VAL_RESIZE (e.g., 320/384) tends to preserve texture\n# and also makes multi-crop tiling meaningful (otherwise 5-crop collapses to near-duplicates).\ntry:\n    _crop_pct = float(cfg.get(\"crop_pct\", 0.875))\nexcept Exception:\n    _crop_pct = 0.875\n\nTIMM_VAL_RESIZE = int(round(float(IMG_SIZE) / max(_crop_pct, 1e-6)))\n\nVAL_EVAL_RESIZE = int(TIMM_VAL_RESIZE)\n\n# For any tiling / multi-crop path we need a larger resize than IMG_SIZE to be meaningful.\nif VAL_TTA_RESIZE is None:\n    if _crop_pct >= 0.99:\n        VAL_TTA_RESIZE_EFFECTIVE = max(int(IMG_SIZE), 320)\n    else:\n        VAL_TTA_RESIZE_EFFECTIVE = int(TIMM_VAL_RESIZE)\nelse:\n    VAL_TTA_RESIZE_EFFECTIVE = int(VAL_TTA_RESIZE)\n\n# Dataset tiling path reads VAL_RESIZE; set it to the TTA resize.\nVAL_RESIZE = int(VAL_TTA_RESIZE_EFFECTIVE)\n\nprint(\"VAL_EVAL_RESIZE:\", VAL_EVAL_RESIZE, \"(timm would be:\", TIMM_VAL_RESIZE, \")\")\nprint(\"VAL_TTA_RESIZE:\", VAL_RESIZE)\nprint(\"VAL_EVAL_MCROPS:\", VAL_EVAL_MCROPS, \"VAL_EVAL_USE_TIMM:\", VAL_EVAL_USE_TIMM)\nprint(\"VAL_TTA_MCROPS:\", VAL_TTA_MCROPS, \"VAL_TTA_MODE:\", VAL_TTA_MODE)\n\n# Post-transform for custom crops (tensor + normalize only).\n# We use this when we do our own tiling (so we don't double-resize/crop).\nval_post = tvt.Compose([\n    tvt.ToTensor(),\n    tvt.Normalize(mean=cfg.get(\"mean\"), std=cfg.get(\"std\")),\n])\n\n# Optional: add mild ColorJitter before timm train_tfms\nCOLORJITTER_TFM = None\nif USE_COLORJITTER:\n    COLORJITTER_TFM = tvt.ColorJitter(\n        brightness=float(CJ_BRIGHTNESS),\n        contrast=float(CJ_CONTRAST),\n        saturation=float(CJ_SATURATION),\n        hue=float(CJ_HUE),\n    )\n\ntrain_tfms = create_transform(\n    input_size=(3, IMG_SIZE, IMG_SIZE),\n    is_training=True,\n    auto_augment=cfg.get(\"auto_augment\", None),\n    interpolation=cfg.get(\"interpolation\", \"bicubic\"),\n    mean=cfg.get(\"mean\"),\n    std=cfg.get(\"std\"),\n)\n\n# For more texture, do a larger resize then crop.\n# We wrap timm's transform by inserting a Resize before it.\nval_tfms = create_transform(\n    input_size=(3, IMG_SIZE, IMG_SIZE),\n    is_training=False,\n    interpolation=cfg.get(\"interpolation\", \"bicubic\"),\n    mean=cfg.get(\"mean\"),\n    std=cfg.get(\"std\"),\n)\n\n# %% [6] Dataset\n\n# --- optional image cleaning helpers ---\n\ndef _grey_world_white_balance(img_rgb: np.ndarray) -> np.ndarray:\n    \"\"\"Simple grey-world white balance in RGB.\"\"\"\n    img = img_rgb.astype(np.float32)\n    mean = img.reshape(-1, 3).mean(axis=0)\n    scale = float(mean.mean()) / (mean + 1e-6)\n    img = img * scale.reshape(1, 1, 3)\n    return np.clip(img, 0.0, 255.0).astype(np.uint8)\n\n\ndef _apply_clahe_rgb(img_rgb: np.ndarray, clip_limit: float = 2.0, grid: int = 8) -> np.ndarray:\n    \"\"\"CLAHE on L channel in LAB, then back to RGB.\"\"\"\n    lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)\n    l, a, b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=float(clip_limit), tileGridSize=(int(grid), int(grid)))\n    l2 = clahe.apply(l)\n    lab2 = cv2.merge((l2, a, b))\n    return cv2.cvtColor(lab2, cv2.COLOR_LAB2RGB)\n\n\ndef clean_image_pil(img: Image.Image) -> Image.Image:\n    \"\"\"Crop bottom artifacts + inpaint orange timestamp; optional lighting normalization.\"\"\"\n    if (not USE_CLEANING) or (not _HAS_CV2):\n        return img\n\n    arr = np.asarray(img)  # RGB uint8\n    if arr.ndim != 3 or arr.shape[2] != 3:\n        return img\n\n    h, w = arr.shape[:2]\n\n    # 1) Safe crop (remove bottom artifacts)\n    if CLEAN_CROP_BOTTOM_PCT and CLEAN_CROP_BOTTOM_PCT > 0:\n        cut = int(round(float(h) * (1.0 - float(CLEAN_CROP_BOTTOM_PCT))))\n        cut = max(1, min(int(h), int(cut)))\n        arr = arr[:cut, :, :]\n\n    # 2) Optional lighting normalization\n    if CLEAN_WHITEBALANCE:\n        arr = _grey_world_white_balance(arr)\n    if CLEAN_CLAHE:\n        arr = _apply_clahe_rgb(arr)\n\n    # 3) Inpaint orange timestamp text\n    if CLEAN_INPAINT_ORANGE:\n        hsv = cv2.cvtColor(arr, cv2.COLOR_RGB2HSV)\n        lower = np.array(CLEAN_ORANGE_LOWER, dtype=np.uint8)\n        upper = np.array(CLEAN_ORANGE_UPPER, dtype=np.uint8)\n        mask = cv2.inRange(hsv, lower, upper)\n\n        if CLEAN_DILATE_ITER and CLEAN_DILATE_ITER > 0:\n            kernel = np.ones((3, 3), np.uint8)\n            mask = cv2.dilate(mask, kernel, iterations=int(CLEAN_DILATE_ITER))\n\n        if int(mask.sum()) > 0:\n            arr = cv2.inpaint(arr, mask, 3, cv2.INPAINT_TELEA)\n\n    return Image.fromarray(arr)\n\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, training: bool, n_crops: int = 1, crop_mode: str = \"five\"):\n        self.df = df.reset_index(drop=True)\n        self.training = training\n        self.n_crops = int(max(1, n_crops))\n        self.crop_mode = str(crop_mode)\n\n    def __len__(self):\n        return len(self.df)\n\n    def _ensure_min_side(self, img: Image.Image, target_min: int) -> Image.Image:\n        w, h = img.size\n        min_side = min(w, h)\n        if min_side >= target_min:\n            return img\n        scale = float(target_min) / float(min_side)\n        new_w = int(round(w * scale))\n        new_h = int(round(h * scale))\n        return img.resize((new_w, new_h), resample=Image.BICUBIC)\n\n    def _resize_for_tiling(self, img: Image.Image, resize_min: int) -> Image.Image:\n        # Make sure we can take IMG_SIZE crops.\n        resize_min = int(resize_min)\n        target_min = max(resize_min, int(IMG_SIZE))\n        img = self._ensure_min_side(img, target_min)\n\n        # If resize_min > IMG_SIZE, scale so min-side ~= resize_min (preserve aspect ratio).\n        if resize_min and resize_min > IMG_SIZE:\n            w, h = img.size\n            scale = float(resize_min) / float(min(w, h))\n            if abs(scale - 1.0) > 1e-6:\n                new_w = int(round(w * scale))\n                new_h = int(round(h * scale))\n                img = img.resize((new_w, new_h), resample=Image.BICUBIC)\n        return img\n\n    def _crop_positions(self, w: int, h: int, crop: int, n: int, mode: str):\n        max_x = max(0, w - crop)\n        max_y = max(0, h - crop)\n\n        if n <= 1:\n            return [(max_x // 2, max_y // 2)]\n\n        if mode == \"five\":\n            pos = [\n                (0, 0),\n                (max_x, 0),\n                (0, max_y),\n                (max_x, max_y),\n                (max_x // 2, max_y // 2),\n            ]\n            return pos[:n]\n\n        # grid mode\n        g = int(math.ceil(math.sqrt(n)))\n        xs = np.linspace(0, max_x, g)\n        ys = np.linspace(0, max_y, g)\n        pos = []\n        for yy in ys:\n            for xx in xs:\n                pos.append((int(round(xx)), int(round(yy))))\n        return pos[:n]\n\n    def __getitem__(self, i: int):\n        rel = str(self.df.loc[i, IMG_COL])\n        path = os.path.join(DATA_DIR, rel)\n\n        img = Image.open(path).convert(\"RGB\")\n        img = clean_image_pil(img)\n        y = y_transform(self.df.loc[i, TARGET_COLS].values)\n        y = torch.from_numpy(y)\n\n        if self.training:\n            # Training: use timm's training pipeline (random crop/AA, etc.)\n            if TRAIN_RESIZE and TRAIN_RESIZE > IMG_SIZE:\n                img = self._resize_for_tiling(img, TRAIN_RESIZE)\n\n            # Mild photometric aug to reduce shortcut reliance (lighting/timestamps/etc.)\n            if COLORJITTER_TFM is not None:\n                img = COLORJITTER_TFM(img)\n\n            x = train_tfms(img)\n            return x, y\n\n        # Validation: optional tiling / multi-crop.\n        # For a single crop, optionally use timm's native eval pipeline (matches pretrained defaults).\n        if self.n_crops <= 1 and VAL_SINGLE_CROP_USES_TIMM:\n            x = val_tfms(img)\n            return x, y\n\n        # Multi-crop / tiling path (or custom single-crop): explicitly resize then crop tiles to IMG_SIZE.\n        # This preserves more texture/detail when VAL_RESIZE > IMG_SIZE, and makes multi-crop meaningful.\n        img = self._resize_for_tiling(img, VAL_RESIZE)\n        w, h = img.size\n\n        if self.n_crops <= 1:\n            # Center crop\n            x0 = max(0, (w - int(IMG_SIZE)) // 2)\n            y0 = max(0, (h - int(IMG_SIZE)) // 2)\n            crop = img.crop((x0, y0, x0 + int(IMG_SIZE), y0 + int(IMG_SIZE)))\n            x = val_post(crop)\n            return x, y\n\n        positions = self._crop_positions(w, h, int(IMG_SIZE), self.n_crops, self.crop_mode)\n        crops = [img.crop((x0, y0, x0 + int(IMG_SIZE), y0 + int(IMG_SIZE))) for (x0, y0) in positions]\n\n        # Apply tensor+normalize only (no extra resize/crop)\n        xs = torch.stack([val_post(c) for c in crops], dim=0)  # (K,3,H,W)\n        return xs, y\n\n\ndef make_loader(df: pd.DataFrame, training: bool, batch_size: int, n_crops: int = 1, crop_mode: str = \"five\") -> DataLoader:\n    # Training is always single-crop. Validation can be single-crop (eval) or multi-crop (TTA).\n    if training:\n        ds = BiomassDataset(df, training=True, n_crops=1, crop_mode=crop_mode)\n    else:\n        ds = BiomassDataset(df, training=False, n_crops=int(max(1, n_crops)), crop_mode=str(crop_mode))\n\n    return DataLoader(\n        ds,\n        batch_size=batch_size,\n        shuffle=training,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n        drop_last=training,\n        persistent_workers=(NUM_WORKERS > 0),\n        prefetch_factor=2,\n    )\n\n# %% [7] Model (timm backbone + regression head)\n\n# We predict ONLY the 3 independent components in model-space (log1p grams if USE_LOG1P):\n#   Dry_Green_g, Dry_Dead_g, Dry_Clover_g\n# Then derive:\n#   GDM_g = Dry_Green_g + Dry_Clover_g\n#   Dry_Total_g = Dry_Dead_g + GDM_g\n# This guarantees constraint-satisfying preds and avoids fitting redundant sums directly.\n\nclass ViTRegressor(nn.Module):\n    def __init__(self, model_name: str, out_dim: int = INDEP_OUT_DIM, pretrained: bool = True):\n        super().__init__()\n        self.out_dim = int(out_dim)\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool=\"avg\")\n        n = self.backbone.num_features\n        self.head = nn.Sequential(\n            nn.LayerNorm(n),\n            nn.Linear(n, 512),\n            nn.GELU(),\n            nn.Dropout(float(HEAD_DROPOUT)),\n            nn.Linear(512, self.out_dim),\n        )\n\n    def forward(self, x):\n        feat = self.backbone(x)\n        return self.head(feat)\n\n\ndef weighted_smooth_l1(pred, target, w, beta: float = 1.0):\n    \"\"\"SmoothL1 (Huber) per output, weighted to match the competition weights.\"\"\"\n    loss = F.smooth_l1_loss(pred, target, reduction=\"none\", beta=beta)\n    w = w.view(1, -1)\n    return (loss * w).mean()\n\n\ndef modelspace_to_grams_torch(yhat_model: torch.Tensor) -> torch.Tensor:\n    \"\"\"model-space -> grams >=0 (works for any last-dim size).\"\"\"\n    if USE_LOG1P:\n        return torch.expm1(yhat_model).clamp_min(0.0)\n    return yhat_model.clamp_min(0.0)\n\n\ndef grams_to_modelspace_torch(y_grams: torch.Tensor) -> torch.Tensor:\n    \"\"\"grams >=0 -> model-space (works for any last-dim size).\"\"\"\n    if USE_LOG1P:\n        return torch.log1p(y_grams.clamp_min(0.0))\n    return y_grams\n\n\ndef grams3_to_grams5_torch(g3: torch.Tensor) -> torch.Tensor:\n    \"\"\"g3: (B,3) grams for (DG, DD, DC) -> (B,5) grams in TARGET_COLS order.\"\"\"\n    gdm = g3[:, 0] + g3[:, 2]\n    total = g3[:, 1] + gdm\n    return torch.cat([g3, gdm.unsqueeze(1), total.unsqueeze(1)], dim=1)\n\n\ndef model3_to_model5_torch(yhat3_model: torch.Tensor) -> torch.Tensor:\n    \"\"\"(B,3) model-space -> (B,5) model-space via grams-space derivation.\"\"\"\n    g3 = modelspace_to_grams_torch(yhat3_model)       # (B,3) grams\n    g5 = grams3_to_grams5_torch(g3)                   # (B,5) grams\n    return grams_to_modelspace_torch(g5)              # (B,5) model-space\n\n\ndef model3_to_grams5_torch(yhat3_model: torch.Tensor) -> torch.Tensor:\n    \"\"\"(B,3) model-space -> (B,5) grams via constraint derivation.\"\"\"\n    g3 = modelspace_to_grams_torch(yhat3_model)\n    return grams3_to_grams5_torch(g3)\n\n# %% [8] Train loop (with best checkpoint by val weighted R2)\n\n\n@torch.no_grad()\ndef predict(model: nn.Module, loader: DataLoader):\n    \"\"\"Predict in *FULL* model-space (5 targets).\n\n    The network outputs only 3 independent targets in model-space. We convert to grams,\n    derive the 5 targets in grams-space, then convert back to model-space.\n\n    Supports either:\n    - x: (B,3,H,W)\n    - x: (B,K,3,H,W) for multi-crop/tiling.\n\n    For multi-crop, we average in grams-space (on the 3 independents), then derive 5.\n    \"\"\"\n    model.eval()\n    preds = []\n    for x, _ in loader:\n        if x.ndim == 5:\n            b, k, c, h, w = x.shape\n            x = x.view(b * k, c, h, w).to(device, non_blocking=True)\n            yhat3 = model(x)  # (b*k, 3) model-space\n\n            g3 = modelspace_to_grams_torch(yhat3).view(b, k, INDEP_OUT_DIM).mean(1)  # (b,3) grams\n            g5 = grams3_to_grams5_torch(g3)                                         # (b,5) grams\n            yhat5 = grams_to_modelspace_torch(g5)                                   # (b,5) model-space\n        else:\n            x = x.to(device, non_blocking=True)\n            yhat3 = model(x)                                                        # (b,3) model-space\n            yhat5 = model3_to_model5_torch(yhat3)                                   # (b,5) model-space\n\n        preds.append(yhat5.detach().cpu().numpy())\n\n    return np.concatenate(preds, axis=0)\n\n\n\ndef train_one_fold(fold: int, run_seed: int):\n    # Per-fold seeding keeps aug/shuffle deterministic within a run seed.\n    seed_everything(int(run_seed) + int(fold))\n    tr = train_wide[train_wide[\"fold\"] != fold].reset_index(drop=True)\n    va = train_wide[train_wide[\"fold\"] == fold].reset_index(drop=True)\n\n    tr_loader = make_loader(tr, training=True, batch_size=BATCH_SIZE)\n    # During training: eval is single-crop (fast + stable)\n    va_loader_eval = make_loader(va, training=False, batch_size=BATCH_SIZE, n_crops=VAL_EVAL_MCROPS, crop_mode=VAL_TTA_MODE)\n\n    # Create model\n    model = ViTRegressor(MODEL_NAME, out_dim=INDEP_OUT_DIM, pretrained=PRETRAINED).to(device)\n\n    # Multi-GPU (2x T4) if enabled\n    if torch.cuda.device_count() > 1:\n        print(\"Using\", torch.cuda.device_count(), \"GPUs via DataParallel\")\n        model = nn.DataParallel(model)\n\n    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n    # Warmup + cosine to MIN_LR (epoch-based)\n    def lr_lambda(epoch_idx: int) -> float:\n        # epoch_idx starts at 0 for LambdaLR\n        if WARMUP_EPOCHS and epoch_idx < WARMUP_EPOCHS:\n            return float(epoch_idx + 1) / float(WARMUP_EPOCHS)\n        if EPOCHS <= (WARMUP_EPOCHS or 0):\n            return 1.0\n        t = (epoch_idx - (WARMUP_EPOCHS or 0)) / float(EPOCHS - (WARMUP_EPOCHS or 0))\n        t = float(np.clip(t, 0.0, 1.0))\n        min_factor = float(MIN_LR / LR)\n        return min_factor + (1.0 - min_factor) * 0.5 * (1.0 + math.cos(math.pi * t))\n\n    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda)\n\n    # new AMP API (removes FutureWarning)\n    scaler = torch.amp.GradScaler('cuda', enabled=(AMP and device.type == 'cuda'))\n    no_improve = 0\n\n    best_wr2 = -1e9\n    best_path = os.path.join(MODELS_DIR, f\"vit_s{int(run_seed)}_fold{fold}.pth\")\n\n    y_true_grams = va[TARGET_COLS].values.astype(np.float32)\n    w_t = torch.tensor(TARGET_WEIGHTS, dtype=torch.float32, device=device)\n\n    for epoch in range(1, EPOCHS + 1):\n        model.train()\n        running = 0.0\n        n_seen = 0\n\n        for x, y in tr_loader:\n            # Support optional multi-crop batches: x can be (B,K,3,H,W)\n            if x.ndim == 5:\n                b, k, c, h, w = x.shape\n                x = x.view(b * k, c, h, w).to(device, non_blocking=True)\n                y = y.to(device, non_blocking=True).float()\n                batch_n = int(y.size(0))\n            else:\n                b, k = None, None\n                x = x.to(device, non_blocking=True)\n                y = y.to(device, non_blocking=True).float()\n                batch_n = int(y.size(0))\n\n            opt.zero_grad(set_to_none=True)\n            with torch.amp.autocast('cuda', enabled=scaler.is_enabled()):\n                yhat3 = model(x)  # (B,3) in model-space (log1p grams if USE_LOG1P)\n\n                if k is not None:\n                    # If training ever uses multi-crop batches, average in grams-space first.\n                    g3 = modelspace_to_grams_torch(yhat3).view(b, k, INDEP_OUT_DIM).mean(1)  # (B,3) grams\n                    pred5 = grams_to_modelspace_torch(grams3_to_grams5_torch(g3))            # (B,5) model-space\n                else:\n                    pred5 = model3_to_model5_torch(yhat3)                                   # (B,5) model-space\n\n                loss = weighted_smooth_l1(pred5, y, w_t, beta=1.0)\n\n            scaler.scale(loss).backward()\n            if GRAD_CLIP is not None and GRAD_CLIP > 0:\n                scaler.unscale_(opt)\n                torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)\n            scaler.step(opt)\n            scaler.update()\n\n            running += float(loss.item()) * batch_n\n            n_seen += batch_n\n\n        sched.step()\n        tr_loss = running / max(1, n_seen)\n\n        # val metric (single-crop eval)\n        va_pred = predict(model, va_loader_eval)\n        wr2 = score_wR2(y_true_grams, va_pred)\n\n        lr_now = opt.param_groups[0][\"lr\"]\n        print(f\"[seed {run_seed}] Fold {fold} | Epoch {epoch:02d}/{EPOCHS} | loss {tr_loss:.4f} | val_wR2 {wr2:.5f} | lr {lr_now:.2e}\")\n\n        if wr2 > best_wr2:\n            best_wr2 = wr2\n            no_improve = 0\n            # Save a SAFE checkpoint (tensors + simple python types only).\n            # This avoids PyTorch 2.6+ weights_only loading issues.\n            torch.save(\n                {\n                    \"model_name\": MODEL_NAME,\n                    \"pretrained\": PRETRAINED,\n                    \"out_dim\": int(INDEP_OUT_DIM),\n                    \"state_dict\": model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(),\n                    \"img_size\": IMG_SIZE,\n                },\n                best_path,\n            )\n        else:\n            no_improve += 1\n\n        if PATIENCE and epoch >= ES_START_EPOCH and no_improve >= PATIENCE:\n            print(f\"Fold {fold} early stop at epoch {epoch} (no improvement for {PATIENCE} evals)\")\n            break\n\n    # Load best and produce final val preds for OOF\n    ckpt = torch.load(best_path, map_location=\"cpu\", weights_only=True)\n    model2 = ViTRegressor(ckpt[\"model_name\"], out_dim=int(ckpt.get(\"out_dim\", INDEP_OUT_DIM)), pretrained=False).to(device)\n    model2.load_state_dict(ckpt[\"state_dict\"], strict=True)\n    model2.eval()\n\n    # Final fold predictions: run TTA multi-crop once per fold to build OOF.\n    va_loader_tta = make_loader(va, training=False, batch_size=BATCH_SIZE, n_crops=VAL_TTA_MCROPS, crop_mode=VAL_TTA_MODE)\n    va_pred_best = predict(model2, va_loader_tta)\n    wr2_best = score_wR2(y_true_grams, va_pred_best)\n    print(f\"[seed {run_seed}] Fold {fold}: BEST val_wR2={wr2_best:.5f} saved={best_path}\")\n\n    # Convert to grams for OOF storage\n    va_pred_grams = np.maximum(y_inverse(va_pred_best), 0.0)\n    if ENFORCE_CONSTRAINTS_EVAL:\n        va_pred_grams = enforce_constraints_grams(va_pred_grams)\n\n    del model, model2, tr_loader, va_loader_eval, va_loader_tta\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    return best_path, wr2_best, va.index.values, va_pred_grams\n\n# %% [9] CV train (optionally multi-seed ensemble)\n\ny_all_true_grams = train_wide[TARGET_COLS].values.astype(np.float32)\n\noof_seeds = []\nseed_fold_scores = {}\n\nfor run_seed in RUN_SEEDS:\n    print(\"=\" * 60)\n    print(f\"RUN_SEED={run_seed}\")\n    print(\"=\" * 60)\n\n    seed_everything(int(run_seed))\n\n    oof = np.zeros((len(train_wide), len(TARGET_COLS)), np.float32)\n    fold_scores = []\n\n    for fold in range(N_FOLDS):\n        best_path, best_wr2, va_idx, va_pred_grams = train_one_fold(fold, run_seed)\n        fold_scores.append(best_wr2)\n\n        sel = (train_wide[\"fold\"].values == fold)\n        oof[sel] = va_pred_grams\n\n    oof_seeds.append(oof)\n    seed_fold_scores[int(run_seed)] = [float(x) for x in fold_scores]\n    print(f\"[seed {run_seed}] mean val_wR2:\", float(np.mean(fold_scores)), \"folds:\", [float(x) for x in fold_scores])\n\n# Average OOF across seeds (already in grams space; constraints remain satisfied)\noof_avg = np.mean(np.stack(oof_seeds, axis=0), axis=0).astype(np.float32)\n\n# Save OOF\nfor j, c in enumerate(TARGET_COLS):\n    train_wide[f\"oof_{c}\"] = oof_avg[:, j]\n\noof_path = os.path.join(OUT_DIR, \"oof.csv\")\ntrain_wide.to_csv(oof_path, index=False)\n\n# Report ensemble score (computed from averaged OOF grams)\nens_fold_scores = []\nfor fold in range(N_FOLDS):\n    sel = (train_wide[\"fold\"].values == fold)\n    ens_fold_scores.append(weighted_r2_global(y_all_true_grams[sel], oof_avg[sel], TARGET_WEIGHTS))\n\nprint(\"Saved:\", oof_path)\nprint(\"Ensemble (averaged OOF) mean wR2:\", float(np.mean(ens_fold_scores)), \"folds:\", [float(x) for x in ens_fold_scores])\n\n# %% [10] Notes / next upgrades\n#\n# Run notes (Jan 2026):\n# - A prior ~0.6 run did NOT save folds.csv, so it is not reproducible. Going forward: keep REUSE_FOLDS=True.\n# - For this competition, \"state_date\" grouped folds are *much stricter* than random CV. Some groups (e.g., NSW/Tas dates)\n#   can form a difficult fold and drag the mean down.\n# - IMPORTANT: if VAL_EVAL is single-crop but final OOF uses multi-crop, checkpoint selection can become unstable.\n#   Current default evaluates each epoch with the same multi-crop settings as OOF/TTA.\n#\n# Latest strict-split run (state_date + greedy; epoch eval uses same multi-crop as TTA):\n#   settings: VAL_EVAL_MCROPS=5, VAL_EVAL_USE_TIMM=False, VAL_TTA_MCROPS=5, VAL_TTA_MODE=\"five\", VAL_TTA_RESIZE=320\n#   folds: [0.6178, 0.5939, 0.4537, 0.4858, 0.5622]  mean=0.5427\n#   -> big improvement vs the earlier 0.5009 run; remaining gap is mainly fold2/fold3 robustness.\n#\n# Prior strict-split baseline (for comparison):\n#   folds: [0.5718, 0.5388, 0.2166, 0.5437, 0.6334]  mean=0.5009\n#   -> fold2 was the problem fold in that run.\n#\n# Next upgrades to push back toward 0.6+:\n# (Applied in code defaults now)\n# - Revert to the best baseline TTA policy:\n#   VAL_EVAL_MCROPS=5, VAL_EVAL_USE_TIMM=False, VAL_TTA_MCROPS=5, VAL_TTA_MODE=\"five\", VAL_TTA_RESIZE=320\n# - Multi-seed ensembling support via RUN_SEEDS (average OOF in grams space).\n# - Early stopping guard: ES_START_EPOCH=8 (avoid killing late-bloom folds).\n#\n# Grid-tiling experiment (for reference):\n# - VAL_TTA_MODE=\"grid\", VAL_*_MCROPS=9, VAL_TTA_RESIZE=384  -> mean ~0.529 (worse than baseline)\n#\n# 1) Stronger tiling for the hard fold(s):\n#    - set VAL_TTA_MODE=\"grid\" and VAL_TTA_MCROPS=9 (3x3) or 16 (4x4)\n#    - optionally raise VAL_TTA_RESIZE to 384\n#\n# 2) Backbones to try (common in strong public notebooks):\n#    - vit_large_patch14_clip_224.openai (current)\n#    - vit_base_patch16_384 (lighter, often strong)\n#    - swinv2_base_window12_384 (often strong)\n#    - DINOv3 / SigLIP variants (usually require adding weights via Kaggle Models)\n#\n# 3) Augment for domain shift (small dataset):\n#    - add mild ColorJitter before train_tfms, and/or increase random erasing\n#\n# 4) Ensembling:\n#    - average 5-fold preds and optionally 2-3 backbones/seeds\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-20T05:08:22.505191Z","iopub.execute_input":"2026-01-20T05:08:22.505713Z","execution_failed":"2026-01-20T05:40:42.018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}