{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":14493832,"sourceType":"datasetVersion","datasetId":9257220},{"sourceId":293031629,"sourceType":"kernelVersion"},{"sourceId":711299,"sourceType":"modelInstanceVersion","modelInstanceId":540408,"modelId":553604},{"sourceId":719660,"sourceType":"modelInstanceVersion","modelInstanceId":547330,"modelId":560139},{"sourceId":726961,"sourceType":"modelInstanceVersion","modelInstanceId":552674,"modelId":565229},{"sourceId":726964,"sourceType":"modelInstanceVersion","modelInstanceId":552674,"modelId":565229},{"sourceId":727078,"sourceType":"modelInstanceVersion","modelInstanceId":552674,"modelId":565229},{"sourceId":727598,"sourceType":"modelInstanceVersion","modelInstanceId":552674,"modelId":565229}],"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":"\n## Credits: Eng Adam Al mohammedi, Baidalin, Adilzhan, CigarCat, Mattia Angeli, Khoa, samu2505 \n\n## dinov3 Lb: 70 now: +","metadata":{}},{"cell_type":"code","source":"# now import timm fresh\nimport timm, sys\nsys.path.append('/kaggle/input/iterative-stratification/iterative-stratification-master')\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n\n# Config\nimport os, gc, math, cv2, 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 time\nfrom timm.utils import ModelEmaV2\nfrom sklearn.model_selection import KFold, 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\ntorch.set_float32_matmul_precision(\"high\")\n\n\nclass CFG:\n    CREATE_SUBMISSION = True\n    USE_TQDM = True\n    PRETRAINED_DIR = '/kaggle/input/vit-pretrain/pytorch/v1/7/model'\n    PRETRAINED = False\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    SEED = 82947501\n    FOLDS_TO_TRAIN = [0, 1, 2, 3, 4]\n    TRAIN_CSV = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    TEST_IMAGE_DIR = '/kaggle/input/csiro-biomass/test'\n    TEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\n    SUBMISSION_DIR = './'\n    MODEL_DIR = '/kaggle/working/model' #'/kaggle/input/baseline-dinov3/pytorch/default/1/baseline-dinov3'\n    N_FOLDS = 4\n\n    MODEL_NAME = 'vit_huge_plus_patch16_dinov3.lvd1689m'\n    BACKBONE_PATH = '/kaggle/input/vit-huge-plus-patch16-dinov3-lvd1689m/vit_huge_plus_patch16_dinov3.lvd1689m_backbone.pth'\n    DINO_GRAD_CHECKPOINTING = False\n\n    IMG_SIZE = 512\n\n    VAL_TTA_TIMES = 1\n    TTA_STEPS = 1\n\n    BATCH_SIZE = 1\n    GRAD_ACC = 4\n    NUM_WORKERS = 4\n    EPOCHS = 50\n    FREEZE_EPOCHS = 0\n    WARMUP_EPOCHS = 3\n    LR_REST = 1e-3\n    LR_BACKBONE = 5e-4\n    WD = 1e-2\n    EMA_DECAY = 0.9\n    CLIP_GRAD_NORM = 1.0\n    PATIENCE = 8\n    TARGET_COLS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS = ['Dry_Clover_g', 'Dry_Dead_g']\n    ALL_TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    R2_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\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\nos.makedirs(CFG.MODEL_DIR, exist_ok=True)\nprint(f'Device : {CFG.DEVICE}')\nprint(f'Backbone: {CFG.MODEL_NAME} | Input: {CFG.IMG_SIZE}')\nprint(f'Freeze Epochs: {CFG.FREEZE_EPOCHS} | Warmup: {CFG.WARMUP_EPOCHS}')\nprint(f'EMA Decay: {CFG.EMA_DECAY} | Grad Acc: {CFG.GRAD_ACC}')\n\n\ndef normalize_species(species_str):\n    tokens = species_str.split('_')\n    normalized = [species_corrections.get(s, s) for s in tokens]\n    return '_'.join(sorted(normalized))\n\ndef get_species_set(species_str):\n    if pd.isna(species_str):\n        return set()\n    tokens = species_str.split('_')\n    normalized = [species_map.get(t, t) for t in tokens]\n    # Special case: \"Mixed\" = Ryegrass + Clover\n    if 'Mixed' in tokens:\n        return {'Ryegrass', 'Clover'}\n    return set(normalized)\ndef plt_show_loss(train_losses, val_scores):\n    fig, axes = plt.subplots(1, 2, figsize=(10, 4))  # two subplots side by side\n\n    # --- Train Loss subplot ---\n    axes[0].plot(train_losses, label='train_loss', color='blue')\n    axes[0].set_title(\"Train Loss\")\n    axes[0].grid(True)\n\n    # Add horizontal line at min loss\n    min_loss = min(train_losses)\n    axes[0].axhline(y=min_loss, color='red', linestyle='--', label='min loss')\n    axes[0].text(3, min_loss, f\"min={min_loss:.4f}\",\n                 va='bottom', ha='center', color='red')\n\n    axes[0].legend()\n\n    # --- Validation Score subplot ---\n    axes[1].plot(val_scores, label='val_score', color='orange')\n    axes[1].set_title(\"Validation Score\")\n    axes[1].grid(True)\n\n    # Add horizontal line at max val_score\n    max_val = max(val_scores)\n    axes[1].axhline(y=max_val, color='blue', linestyle='--', label='max val')\n    axes[1].text(3, max_val, f\"max={max_val:.4f}\",\n                 va='bottom', ha='center', color='blue')\n\n    axes[1].legend()\n\n    plt.tight_layout()\n    plt.show()\n\n# Metrics\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    r2_scores = []\n    for i in range(y_true.shape[1]):\n        yt = y_true[:, i];\n        yp = y_pred[:, i]\n        ss_res = np.sum((yt - yp) ** 2)\n        ss_tot = np.sum((yt - np.mean(yt)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n    r2_scores = np.array(r2_scores)\n    weighted = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted, r2_scores\n\n\ndef weighted_r2_score_global(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    flat_true = y_true.reshape(-1)\n    flat_pred = y_pred.reshape(-1)\n    w = np.tile(weights, y_true.shape[0])\n    mean_w = np.sum(w * flat_true) / np.sum(w)\n    ss_res = np.sum(w * (flat_true - flat_pred) ** 2)\n    ss_tot = np.sum(w * (flat_true - mean_w) ** 2)\n    global_r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n    avg_r2, per_r2 = weighted_r2_score(y_true, y_pred)\n    return global_r2, avg_r2, per_r2\n\n\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n\n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n\n        print(f'\\nTarget: {target}')\n        print(f'{\"Index\":<6} | {\"Image Path\":<40} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}')\n        print('-' * 90)\n\n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\n\n\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n\n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n\n        print(f'\\nTarget: {target}')\n        header = f'{\"Index\":<6} | {\"Image Path\":<40} | {\"State\":<6} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}'\n        print(header)\n        print('-' * len(header))\n\n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            state = val_df.iloc[idx]['State'] if 'State' in val_df.columns else 'NA'\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {str(state):<6} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\n\n\ndef compare_train_val(tr_df, val_df, targets, show_plots=True):\n    \"\"\"Quick comparison of target distributions and metadata between train and val splits.\"\"\"\n    print(\"\\n--- Train / Val Comparison ---\")\n\n    for t in targets:\n        tr = tr_df.get(t, pd.Series(dtype=float)).dropna()\n        val = val_df.get(t, pd.Series(dtype=float)).dropna()\n        print(f\"\\nTarget: {t}\")\n        print(f\"  Train: n={len(tr)} mean={tr.mean():.3f} std={tr.std():.3f} min={tr.min():.3f} max={tr.max():.3f}\")\n        print(\n            f\"  Val  : n={len(val)} mean={val.mean():.3f} std={val.std():.3f} min={val.min():.3f} max={val.max():.3f}\")\n        if show_plots:\n            try:\n                plt.figure(figsize=(6, 3))\n                sns.kdeplot(tr, label='train', fill=True)\n                sns.kdeplot(val, label='val', fill=True)\n                plt.legend()\n                plt.title(f'Distribution: {t}')\n                plt.show()\n            except Exception as e:\n                print('  Could not plot distributions for', t, '-', e)\n\n    # Compare Sampling_Date and State if present\n    if 'Sampling_Date' in tr_df.columns:\n        try:\n            tr_dates = pd.to_datetime(tr_df['Sampling_Date'], errors='coerce')\n            val_dates = pd.to_datetime(val_df['Sampling_Date'], errors='coerce')\n            print(\"\\nSampling_Date range:\")\n            print(f\"  Train: {tr_dates.min()} -> {tr_dates.max()} (missing {tr_dates.isna().sum()})\")\n            print(f\"  Val  : {val_dates.min()} -> {val_dates.max()} (missing {val_dates.isna().sum()})\")\n        except Exception as e:\n            print('  Could not parse Sampling_Date:', e)\n    if 'State' in tr_df.columns:\n        print(\"\\nState distribution (train vs val):\")\n        tr_state = tr_df['State'].value_counts(normalize=True)\n        val_state = val_df['State'].value_counts(normalize=True)\n        state_df = pd.concat([tr_state, val_state], axis=1, keys=['train', 'val']).fillna(0)\n\n        print(state_df)\n\n\ndef biomass_loss(outputs, labels, w=None):\n    total, gdm, green, clover, dead = outputs\n    mse = nn.MSELoss()\n    huber = nn.SmoothL1Loss(beta=5.0)  # Huber loss for robust regression (beta=5.0 as recommended)\n\n    l_green = huber(green.squeeze(), labels[:, 0])\n    l_dead = huber(dead.squeeze(), labels[:, 1])  # Use Huber loss for Dead\n    l_clover = huber(clover.squeeze(), labels[:, 2])\n    l_gdm = huber(gdm.squeeze(), labels[:, 3])\n    l_total = huber(total.squeeze(), labels[:, 4])\n\n    # Stack per-target losses in the SAME order as CFG.ALL_TARGET_COLS\n    losses = torch.stack([l_green, l_dead, l_clover, l_gdm, l_total])\n    # losses = torch.stack([l_green, l_dead, l_clover])\n\n    # Use provided weights, or default to CFG.R2_WEIGHTS\n    if w is None:\n        return losses.mean()\n    w = torch.as_tensor(w, device=losses.device, dtype=losses.dtype)\n    w = w / w.sum()\n    return (losses * w).sum()\n\n\n# Transforms\ndef get_train_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.Rotate(limit=(-10, 10), p=0.3, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101),\n        A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\n\ndef get_val_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\n\ndef get_tta_transforms(mode=0):\n    # mode 0: original\n    # mode 1: hflip\n    # mode 2: vflip\n    # mode 3: rotate90\n    transforms_list = [\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n    ]\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.append(A.RandomRotate90(p=1.0))  # RandomRotate90 with p=1.0 rotates 90, 180, 270 randomly?\n        # Albumentations RandomRotate90 rotates by 90, 180, 270.\n        # Reference uses transforms.RandomRotation([90, 90]) which is exactly 90 degrees.\n        # To match exactly 90 degrees in Albumentations, we might need Rotate(limit=(90,90), p=1.0)\n        # But RandomRotate90 is standard TTA. Let's use Rotate(limit=(90,90)) to be precise if that's what reference does.\n        # Reference: transforms.RandomRotation([90, 90]) -> rotates by exactly 90 degrees.\n        transforms_list.append(\n            A.Rotate(limit=(90, 90), p=1.0, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101))\n\n    transforms_list.extend([\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\n    return A.Compose(transforms_list, p=1.0)\n\n\ndef clean_image(img):\n    # 1. Safe Crop (Remove artifacts at the bottom)\n    h, w = img.shape[:2]\n    # Cut bottom 10% where artifacts often appear\n    img = img[0:int(h * 0.90), :]\n\n    # 2. Inpaint Date Stamp (Remove orange text)\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    # Define orange color range (adjust as needed)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n\n    # Dilate mask to cover text edges and reduce noise\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n\n    # Inpaint if mask is not empty\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n\n    return img\n\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df\n        self.transform = transform\n        self.img_dir = img_dir\n        self.paths = df['image_path'].values\n        self.labels = df[CFG.ALL_TARGET_COLS].values.astype(np.float32)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.img_dir, os.path.basename(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        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n        left = self.transform(image=left)['image']\n        right = self.transform(image=right)['image']\n        label = torch.from_numpy(self.labels[idx])\n        return left, right, label\n\n# Layers\nclass LocalMambaBlock(nn.Module):\n    \"\"\"\n    Lightweight Mamba-style block (Gated CNN) from the reference notebook.\n    Efficiently mixes tokens with linear complexity.\n    \"\"\"\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        # Depthwise conv mixes spatial information locally\n        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)\n        self.gate = nn.Linear(dim, dim)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x):\n        # x: (Batch, Tokens, Dim)\n        shortcut = x\n        x = self.norm(x)\n        # Gating mechanism\n        g = torch.sigmoid(self.gate(x))\n        x = x * g\n        # Spatial mixing via 1D Conv (requires transpose)\n        x = x.transpose(1, 2)  # -> (B, D, N)\n        x = self.dwconv(x)\n        x = x.transpose(1, 2)  # -> (B, N, D)\n        # Projection\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\n# Model\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=True, backbone_path=None):\n        super().__init__()\n        self.model_name = model_name\n        self.backbone_path = backbone_path\n\n        # 1. Load Backbone with global_pool='' to keep patch tokens\n        #    (B, 197, 1024) instead of (B, 1024)\n        self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='')\n\n        # 2. Enable Gradient Checkpointing (Crucial for ViT-Large memory!)\n        if hasattr(self.backbone, 'set_grad_checkpointing') and CFG.DINO_GRAD_CHECKPOINTING:\n            self.backbone.set_grad_checkpointing(True)\n            print(\"✓ Gradient Checkpointing enabled (saves ~50% VRAM)\")\n\n        nf = self.backbone.num_features\n\n        # 3. Mamba Fusion Neck\n        #    Mixes the concatenated tokens [Left, Right]\n        self.fusion = nn.Sequential(\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1),\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1)\n        )\n\n        # 4. Pooling & Heads\n        self.pool = nn.AdaptiveAvgPool1d(1)\n\n        # Heads (using the same logic as before, but on fused features)\n        self.head_green_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n        self.head_clover_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n        self.head_dead_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n\n        if pretrained:\n            self.load_pretrained()\n\n    def load_pretrained(self):\n        try:\n            # Load weights normally\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                # Handle common checkpoint wrappers (e.g. if saved with 'model' key)\n                if 'model' in sd:\n                    sd = sd['model']\n                elif 'state_dict' in sd:\n                    sd = sd['state_dict']\n            else:\n                # Original behavior: Download from internet\n                print(\"Downloading backbone weights...\")\n                sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool='').state_dict()\n\n            # Interpolate pos_embed if needed (for 256x256 vs 224x224)\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 pe_ck.shape != pe_m.shape:\n                    print(f\"Interpolating pos_embed: {pe_ck.shape} -> {pe_m.shape}\")\n                    # (Simple interpolation logic here or rely on timm's load if strict=False handles it well enough)\n                    # For robust interpolation, use the snippet provided in previous turn\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    def forward(self, left, right):\n        # 1. Extract Tokens (B, N, D)\n        #    Note: ViT usually returns [CLS, Patch1, Patch2...]\n        #    We remove CLS token for spatial mixing, or keep it. Let's keep it.\n        x_l = self.backbone(left)\n        x_r = self.backbone(right)\n\n        # 2. Concatenate Left and Right tokens along sequence dimension\n        #    (B, N, D) + (B, N, D) -> (B, 2N, D)\n        x_cat = torch.cat([x_l, x_r], dim=1)\n\n        # 3. Apply Mamba Fusion\n        #    This allows tokens from Left image to interact with tokens from Right image\n        x_fused = self.fusion(x_cat)\n\n        # 4. Global Pooling\n        #    (B, 2N, D) -> (B, D, 2N) -> (B, D, 1) -> (B, D)\n        x_pool = self.pool(x_fused.transpose(1, 2)).flatten(1)\n\n        # 5. Prediction Heads\n        green = self.head_green_raw(x_pool)\n        clover = self.head_clover_raw(x_pool)\n        dead = self.head_dead_raw(x_pool)\n\n        # Summation logic\n        gdm = green + clover\n        total = gdm + dead\n\n        return total, gdm, green, clover, dead\n\n\n# Utility Functions\ndef set_backbone_requires_grad(model: BiomassModel, requires_grad: bool):\n    for p in model.backbone.parameters():\n        p.requires_grad = requires_grad\n\ndef build_optimizer(model: BiomassModel):\n    # 1. Get backbone parameter IDs for exclusion\n    backbone_ids = {id(p) for p in model.backbone.parameters()}\n\n    # 2. Separate params into backbone vs. everything else (heads, fusion, etc.)\n    backbone_params = []\n    rest_params = []\n\n    for p in model.parameters():\n        if p.requires_grad:\n            if id(p) in backbone_ids:\n                backbone_params.append(p)\n            else:\n                rest_params.append(p)\n\n    return optim.AdamW([\n        {'params': backbone_params, 'lr': CFG.LR_BACKBONE, 'weight_decay': CFG.WD},\n        {'params': rest_params, 'lr': CFG.LR_REST, 'weight_decay': CFG.WD},\n    ])\n\n\ndef build_scheduler(optimizer):\n    def lr_lambda(epoch):\n        e = max(0, epoch - 1)\n        if e < CFG.WARMUP_EPOCHS:\n            return float(e + 1) / float(max(1, CFG.WARMUP_EPOCHS))\n        progress = (e - CFG.WARMUP_EPOCHS) / float(max(1, CFG.EPOCHS - CFG.WARMUP_EPOCHS))\n        return 0.5 * (1.0 + math.cos(math.pi * progress))\n\n    return LambdaLR(optimizer, lr_lambda)\n\n\n# Training and Validation Loops\n\nfrom contextlib import nullcontext\n\nUSE_BF16 = True\nAMP_DTYPE = torch.bfloat16 if USE_BF16 else torch.float16\n\nscaler = torch.amp.GradScaler(\n    'cuda',\n    enabled=(torch.cuda.is_available() and AMP_DTYPE == torch.float16)\n)\n\ndef autocast_ctx():\n    if not torch.cuda.is_available():\n        return nullcontext()\n    return torch.amp.autocast(device_type='cuda', dtype=AMP_DTYPE)\n\ndef _as_col(x: torch.Tensor, bs: int) -> torch.Tensor:\n    \"\"\"Force predictions to (B,1) regardless of model head returning (B,) or (B,1).\"\"\"\n    if x.ndim == 1:\n        return x.view(bs, 1)\n    if x.ndim == 2:\n        return x\n    return x.view(bs, 1)\n\ndef _ensure_2d_lab(lab: torch.Tensor, bs: int) -> torch.Tensor:\n    \"\"\"Force labels to (B,T).\"\"\"\n    if lab.ndim == 1:\n        return lab.view(bs, -1)\n    return lab\n\ndef _pred_pack(p_total, p_gdm, p_green, p_clover, p_dead, bs: int) -> torch.Tensor:\n    \"\"\"\n    Pack predictions in the exact order expected by your metric:\n    CFG.ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    \"\"\"\n    pg = _as_col(p_green,  bs)\n    pd = _as_col(p_dead,   bs)\n    pc = _as_col(p_clover, bs)\n    pgdm = _as_col(p_gdm,  bs)\n    pt = _as_col(p_total,  bs)\n    return torch.cat([pg, pd, pc, pgdm, pt], dim=1)  # (B,5)\n\n@torch.inference_mode()\ndef valid_epoch(eval_model, loader, device):\n    eval_model.eval()\n\n    n = len(loader.dataset)\n    n_targets = len(CFG.ALL_TARGET_COLS)\n\n    preds_cpu  = torch.empty((n, n_targets), dtype=torch.float32)\n    labels_cpu = torch.empty((n, n_targets), dtype=torch.float32)\n\n    total_loss = 0.0\n    offset = 0\n\n    for l, r, lab in loader:\n        bs = l.size(0)\n        l   = l.to(device, non_blocking=True)\n        r   = r.to(device, non_blocking=True)\n        lab = lab.to(device, non_blocking=True)\n\n        with autocast_ctx():\n            p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n            loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead),\n                                lab, w=CFG.LOSS_WEIGHTS)\n\n        total_loss += loss.detach().float().item() * bs\n\n        batch_pred = _pred_pack(p_total, p_gdm, p_green, p_clover, p_dead, bs).float().cpu()\n        batch_lab  = _ensure_2d_lab(lab, bs).float().cpu()\n\n        # Safety: ensure dims match (avoids silent 4/5 bugs)\n        if batch_pred.shape[1] != n_targets or batch_lab.shape[1] != n_targets:\n            raise RuntimeError(\n                f\"Target dim mismatch: pred={batch_pred.shape}, lab={batch_lab.shape}, \"\n                f\"CFG.ALL_TARGET_COLS={n_targets}\"\n            )\n\n        preds_cpu[offset:offset+bs]  = batch_pred\n        labels_cpu[offset:offset+bs] = batch_lab\n        offset += bs\n\n    pred_all    = preds_cpu.numpy()\n    true_labels = labels_cpu.numpy()\n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(true_labels, pred_all)\n\n    return total_loss / n, global_r2, avg_r2, per_r2, pred_all, true_labels\n\n\n@torch.inference_mode()\ndef valid_epoch_tta(eval_model, loaders, device):\n    \"\"\"\n    Fixed: always accumulate (N,5). Works even if a head outputs (B,1).\n    Assumes each loader iterates in identical order (shuffle=False).\n    \"\"\"\n    eval_model.eval()\n    assert len(loaders) > 0\n\n    n = len(loaders[0].dataset)\n    n_targets = len(CFG.ALL_TARGET_COLS)\n\n    labels_cpu = torch.empty((n, n_targets), dtype=torch.float32)\n    preds_sum  = torch.zeros((n, n_targets), dtype=torch.float32)\n\n    total_loss = 0.0\n\n    for tta_i, loader in enumerate(loaders):\n        offset = 0\n        tta_loss = 0.0\n\n        for l, r, lab in tqdm(loader,desc='validate'):\n            bs = l.size(0)\n            l   = l.to(device, non_blocking=True)\n            r   = r.to(device, non_blocking=True)\n            lab = lab.to(device, non_blocking=True)\n\n            with autocast_ctx():\n                p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n                loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead),\n                                    lab, w=CFG.LOSS_WEIGHTS)\n\n            tta_loss += loss.detach().float().item() * bs\n\n            batch_pred = _pred_pack(p_total, p_gdm, p_green, p_clover, p_dead, bs).float().cpu()\n            preds_sum[offset:offset+bs] += batch_pred\n\n            if tta_i == 0:\n                batch_lab = _ensure_2d_lab(lab, bs).float().cpu()\n                if batch_lab.shape[1] != n_targets:\n                    raise RuntimeError(f\"Label dim mismatch: {batch_lab.shape} vs {n_targets}\")\n                labels_cpu[offset:offset+bs] = batch_lab\n\n            offset += bs\n\n        total_loss += tta_loss / n\n\n    avg_preds   = (preds_sum / len(loaders)).numpy()\n    true_labels = labels_cpu.numpy()\n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(true_labels, avg_preds)\n\n    return total_loss / len(loaders), global_r2, avg_r2, per_r2, avg_preds, true_labels\n\n\ndef train_epoch(model, loader, opt, scheduler, device, ema: ModelEmaV2 | None = None):\n    model.train()\n    running = 0.0\n\n    opt.zero_grad(set_to_none=True)\n    itera = tqdm(loader, desc='train', leave=False) if CFG.USE_TQDM else loader\n\n    for i, (l, r, lab) in enumerate(itera):\n        bs = l.size(0)\n        l   = l.to(device, non_blocking=True)\n        r   = r.to(device, non_blocking=True)\n        lab = lab.to(device, non_blocking=True)\n\n        with autocast_ctx():\n            p_total, p_gdm, p_green, p_clover, p_dead = model(l, r)\n            loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead),\n                                lab, w=CFG.LOSS_WEIGHTS)\n            loss = loss / CFG.GRAD_ACC\n\n\n        if not torch.isfinite(loss):\n            print(\"Non-finite loss detected; skipping batch\")\n            opt.zero_grad(set_to_none=True)\n            continue\n\n        if scaler.is_enabled():\n            scaler.scale(loss).backward()\n        else:\n            loss.backward()\n\n        running += loss.detach().float().item() * bs * CFG.GRAD_ACC\n\n        do_step = ((i + 1) % CFG.GRAD_ACC == 0) or ((i + 1) == len(loader))\n        if do_step:\n            if scaler.is_enabled():\n                scaler.unscale_(opt)\n                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n                scaler.step(opt)\n                scaler.update()\n            else:\n                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n                opt.step()\n\n            if ema is not None:\n                ema.update(model.module if hasattr(model, \"module\") else model)\n\n            opt.zero_grad(set_to_none=True)\n\n    scheduler.step()\n    return running / len(loader.dataset)\n\n\n\n# === MAIN TRAINING LOOP ===\n\n# Helper for accurate GPU timings\ndef _sync():\n    if torch.cuda.is_available():\n        torch.cuda.synchronize()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T11:26:17.156935Z","iopub.execute_input":"2026-01-22T11:26:17.157639Z","iopub.status.idle":"2026-01-22T11:26:37.522779Z","shell.execute_reply.started":"2026-01-22T11:26:17.157593Z","shell.execute_reply":"2026-01-22T11:26:37.522017Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''# train model\n\nprint('Loading data...')\ndf_long = pd.read_csv(CFG.TRAIN_CSV)\ndf_wide = df_long.pivot(\n    index='image_path',\n    columns='target_name',\n    values='target'\n)\n\ndf_wide = df_wide.reset_index()\ndf_wide.columns.name = None\n\nmeta_cols = ['image_path', 'Sampling_Date', 'State', 'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm']\nmeta_df = df_long[meta_cols].drop_duplicates()\ndf_wide = df_wide.merge(meta_df, on='image_path', how='left')\n\n# Manual mapping for known inconsistencies\nspecies_corrections = {\n    'Barleygrass': 'BarleyGrass',\n}\ndf_wide['Species'] = df_wide['Species'].apply(normalize_species)\n\ndf = df_wide.copy()\n\n# 1. Total biomass bin (must be balanced)\ndf['total_bin'] = pd.qcut(df['Dry_Total_g'], q=10, labels=False, duplicates='drop')\n\n# 2. Season (Australian seasons)\ndf['Sampling_Date'] = pd.to_datetime(df['Sampling_Date'])\ndf['month'] = df['Sampling_Date'].dt.month\ndf['season'] = ((df['month'] + 9) % 12) // 3 + 1  # 1=Spring, 2=Summer, 3=Autumn, 4=Winter\n\n# 3. Binned continuous\ndf['ndvi_bin'] = pd.qcut(df['Pre_GSHH_NDVI'], q=5, labels=False, duplicates='drop')\ndf['height_bin'] = pd.qcut(df['Height_Ave_cm'], q=5, labels=False, duplicates='drop')\n\n# 4. State as categorical\ndf['state_enc'] = df['State'].astype('category').cat.codes\n\n# ==================== MULTI-HOT SPECIES ====================\n\nall_species = ['Ryegrass', 'Clover', 'Phalaris', 'Lucerne', 'Fescue', 'BarleyGrass',\n               'Capeweed', 'SilverGrass', 'SpearGrass', 'Bromegrass', 'CrumbWeed',\n               'SubcloverDalkeith', 'SubcloverLosa', 'WhiteClover']\n\n# Manual mapping for known variants\nspecies_map = {\n    'SubcloverDalkeith': 'Clover',\n    'SubcloverLosa': 'Clover',\n    'WhiteClover': 'Clover',\n}\nfor sp in all_species:\n    df[f'sp_{sp}'] = df['Species'].apply(lambda x: 1 if sp in get_species_set(x) else 0)\n\n# =====================================================================\n\n# Final multilabel array — now includes individual species presence!\nlabel_columns = ['total_bin', 'state_enc', 'season', 'ndvi_bin', 'height_bin'] + [f'sp_{sp}' for sp in all_species]\n\nlabels = df[label_columns].values.astype(int)\n\n# Multilabel Stratified K-Fold\nmskf = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n\ndf['fold'] = -1\nfor fold, (_, val_idx) in enumerate(mskf.split(df, labels)):\n    df.loc[val_idx, 'fold'] = fold\n\ncols = ['image_path', 'Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'Dry_Total_g', 'GDM_g'] + label_columns\ndf_val = df[df.fold == CFG.N_FOLDS].copy()\ndf_train = df[df.fold != CFG.N_FOLDS].copy()\n\ndf_val = df_val[cols]\ndf_train = df_train[cols]\n\nDL_KW = dict(\n    num_workers=CFG.NUM_WORKERS,\n    pin_memory=True,\n    persistent_workers=(CFG.NUM_WORKERS > 0),\n    prefetch_factor=4 if CFG.NUM_WORKERS > 0 else None,\n)\n\ntrain_dataset = BiomassDataset(\n    df_train,\n    get_train_transforms(),\n    CFG.TRAIN_IMAGE_DIR\n)\n\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=CFG.BATCH_SIZE,\n    shuffle=False,\n    drop_last=False,\n    **{k: v for k, v in DL_KW.items() if v is not None},\n)\n\nval_loaders = []\nfor mode in range(CFG.VAL_TTA_TIMES):\n    val_dataset = BiomassDataset(\n        df_val,\n        get_tta_transforms(mode),\n        CFG.TRAIN_IMAGE_DIR\n    )\n    \n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=CFG.BATCH_SIZE,\n        shuffle=False,\n        drop_last=False,\n        **{k: v for k, v in DL_KW.items() if v is not None},\n    )\n    val_loaders.append(val_loader)\n\n# One place for loader kwargs (fast)\n\n\n\n\nprint('Building model...')\nbackbone_path = getattr(CFG, 'BACKBONE_PATH', None)\nmodel = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED, backbone_path=backbone_path).to(CFG.DEVICE)\nbest_global_r2 = -np.inf\n# Load pretrained fold weights if available (for resuming or fine-tuning)\nif getattr(CFG, 'PRETRAINED_DIR', None) and os.path.isdir(CFG.PRETRAINED_DIR):\n    pretrained_path = os.path.join(CFG.PRETRAINED_DIR, f'best_model_vit_huge_plus_patch16_dinov3.lvd1689m.pth')\n    if os.path.exists(pretrained_path):\n        try:\n            state = torch.load(pretrained_path, map_location='cpu', weights_only=False)\n            if isinstance(state, dict) :\n            #if isinstance(state, dict) and ('model_state_dict' in state or 'state_dict' in state):\n                if '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                if 'epochs' in state:\n                    pretrained_epoch = state['epochs']\n                    print('pretrained_epoch',pretrained_epoch)\n                if 'global_r2' in state:\n                    best_global_r2 = state['global_r2']\n            else:\n                sd = state\n            model.load_state_dict(sd, strict=False)\n            model.to(CFG.DEVICE)\n            print(f'  ✓ Loaded pretrained weights from {pretrained_path}')\n        except Exception as e:\n            print(f'  ✗ Failed to load pretrained : {e}')\n    else:\n        print(f'  (No pretrained file for at {pretrained_path})')\nelse:\n    print('  (No PRETRAINED_DIR configured or directory missing)')\n\n# Single GPU: DO NOT wrap in DataParallel\nmodel = nn.DataParallel(model)  # <-- removed\n\n# Freeze/unfreeze backbone\nset_backbone_requires_grad(model.module, False)\n\noptimizer = build_optimizer(model.module)\nscheduler = build_scheduler(optimizer)\n\n# EMA on the real model\nema = ModelEmaV2(model, decay=CFG.EMA_DECAY)\n  \n\nbest_avg_r2 = -np.inf\npatience = 0\nbest_fold_preds = None\nbest_fold_true = None\ntrain_losses =[]\nval_losses = []\nsave_path = os.path.join(CFG.MODEL_DIR, f'best_model_{CFG.MODEL_NAME}.pth')\nfor epoch in range(1, CFG.EPOCHS + 1):\n    if epoch <= pretrained_epoch:\n        continue\n    if epoch == CFG.FREEZE_EPOCHS + 1:\n        patience = 0\n        set_backbone_requires_grad(model.module, True)\n        print(f'Epoch {epoch}: backbone unfrozen')\n\n    # ---- Train timing ----\n    _sync()\n    t0 = time.perf_counter()\n    tr_loss = train_epoch(model, train_loader, optimizer, scheduler, CFG.DEVICE, ema)\n    train_losses.append(tr_loss)\n    _sync()\n    t1 = time.perf_counter()\n\n    # choose eval model (EMA weights)\n    eval_model = ema.module if ema is not None else model\n\n    # ---- Val timing (TTA) ----\n    _sync()\n    t2 = time.perf_counter()\n    val_loss, global_r2, avg_r2, per_r2, preds_fold, true_fold = valid_epoch_tta(eval_model, val_loaders, CFG.DEVICE)\n    val_losses.append(val_loss)\n    _sync()\n    t3 = time.perf_counter()\n\n    time_tr = t1 - t0\n    time_val = t3 - t2\n    time_ep = t3 - t0\n\n    per_r2_str = ' | '.join([f'{CFG.ALL_TARGET_COLS[i][:5]}: {r2:.3f}' for i, r2 in enumerate(per_r2)])\n    lrs = [pg['lr'] for pg in optimizer.param_groups]\n    lr_str = ' '.join([f'lr{i}={lr:.3e}' for i, lr in enumerate(lrs)])\n\n    print(\n        f'Epoch {epoch:02d} | '\n        f'TLoss {tr_loss:.5f} | VLoss {val_loss:.5f} | '\n        f'avgR2 {avg_r2:.4f} | GlobalR² {global_r2:.4f} '\n        f'{\"[BEST]\" if global_r2 > best_global_r2 else \"\"} | '\n        f'{lr_str} | time_tr={time_tr:.1f}s time_val={time_val:.1f}s time_ep={time_ep:.1f}s'\n    )\n    print(f'  → {per_r2_str}')\n\n    if global_r2 > best_global_r2:\n        best_global_r2 = global_r2\n        best_avg_r2 = avg_r2\n\n        # Save EMA weights directly (no CPU clone) for speed\n        # NOTE: safest is to save state_dict tensors as-is; this is typically fine on DGX.\n        torch.save({\n            'model_state_dict':eval_model.state_dict(),\n            'model_name': CFG.MODEL_NAME,\n            'epochs': epoch,\n            'global_r2': best_global_r2\n        }, save_path)\n        print(f'  → SAVED EMA weights to {save_path} (GlobalR²: {best_global_r2:.4f})')\n\n        patience = 0\n        best_fold_preds = preds_fold\n        best_fold_true = true_fold\n    else:\n        patience += 1\n        if patience >= CFG.PATIENCE:\n            print(f'  → EARLY STOP (no improvement in {CFG.PATIENCE} epochs)')\n            break\n    plt_show_loss(train_losses, val_losses)\n\n    # keep memory tidy but avoid heavy cache/gc churn\n    del preds_fold, true_fold\n# Cleanup for this fold\n\ndel model, train_loader, val_loaders, optimizer, scheduler, ema, eval_model\n_sync()\ntorch.cuda.empty_cache()\ngc.collect()'''","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T11:26:37.524278Z","iopub.execute_input":"2026-01-22T11:26:37.524687Z","iopub.status.idle":"2026-01-22T11:26:37.534310Z","shell.execute_reply.started":"2026-01-22T11:26:37.524660Z","shell.execute_reply":"2026-01-22T11:26:37.533697Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport timm\nimport torch\nfrom albumentations import ToTensorV2\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nimport albumentations as A\n\nclass CFG:\n    CREATE_SUBMISSION = True\n    USE_TQDM = True\n    PRETRAINED_DIR = '/kaggle/input/vit-pretrain/pytorch/v1/7/model'\n    PRETRAINED = False\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    SEED = 82947501\n    FOLDS_TO_TRAIN = [0, 1, 2, 3, 4]\n    TRAIN_CSV = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    TEST_IMAGE_DIR = '/kaggle/input/csiro-biomass/test'\n    TEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\n    SUBMISSION_DIR = './'\n    MODEL_DIR = '/kaggle/working/model' #'/kaggle/input/baseline-dinov3/pytorch/default/1/baseline-dinov3'\n    N_FOLDS = 4\n\n    MODEL_NAME = 'vit_huge_plus_patch16_dinov3.lvd1689m'\n    BACKBONE_PATH = '/kaggle/input/vit-huge-plus-patch16-dinov3-lvd1689m/vit_huge_plus_patch16_dinov3.lvd1689m_backbone.pth'\n    DINO_GRAD_CHECKPOINTING = False\n\n    IMG_SIZE = 512\n\n    VAL_TTA_TIMES = 1\n    TTA_STEPS = 1\n\n    BATCH_SIZE = 1\n    GRAD_ACC = 4\n    NUM_WORKERS = 4\n    EPOCHS = 50\n    FREEZE_EPOCHS = 0\n    WARMUP_EPOCHS = 3\n    LR_REST = 1e-3\n    LR_BACKBONE = 5e-4\n    WD = 1e-2\n    EMA_DECAY = 0.9\n    CLIP_GRAD_NORM = 1.0\n    PATIENCE = 8\n    TARGET_COLS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS = ['Dry_Clover_g', 'Dry_Dead_g']\n    ALL_TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    R2_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\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\n# test model\n\nclass LocalMambaBlock(nn.Module):\n    \"\"\"\n    Lightweight Mamba-style block (Gated CNN) from the reference notebook.\n    Efficiently mixes tokens with linear complexity.\n    \"\"\"\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        # Depthwise conv mixes spatial information locally\n        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)\n        self.gate = nn.Linear(dim, dim)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x):\n        # x: (Batch, Tokens, Dim)\n        shortcut = x\n        x = self.norm(x)\n        # Gating mechanism\n        g = torch.sigmoid(self.gate(x))\n        x = x * g\n        # Spatial mixing via 1D Conv (requires transpose)\n        x = x.transpose(1, 2)  # -> (B, D, N)\n        x = self.dwconv(x)\n        x = x.transpose(1, 2)  # -> (B, N, D)\n        # Projection\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\n\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=True, backbone_path=None):\n        super().__init__()\n        self.model_name = model_name\n        self.backbone_path = backbone_path\n\n        # 1. Load Backbone with global_pool='' to keep patch tokens\n        #    (B, 197, 1024) instead of (B, 1024)\n        self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='')\n\n        nf = self.backbone.num_features\n\n        # 3. Mamba Fusion Neck\n        #    Mixes the concatenated tokens [Left, Right]\n        self.fusion = nn.Sequential(\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1),\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1)\n        )\n\n        # 4. Pooling & Heads\n        self.pool = nn.AdaptiveAvgPool1d(1)\n\n        # Heads (using the same logic as before, but on fused features)\n        self.head_green_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n        self.head_clover_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n        self.head_dead_raw = nn.Sequential(\n            nn.Linear(nf, nf // 2), nn.GELU(), nn.Dropout(0.2),\n            nn.Linear(nf // 2, 1), nn.Softplus()\n        )\n\n    def forward(self, left, right):\n        # 1. Extract Tokens (B, N, D)\n        #    Note: ViT usually returns [CLS, Patch1, Patch2...]\n        #    We remove CLS token for spatial mixing, or keep it. Let's keep it.\n        x_l = self.backbone(left)\n        x_r = self.backbone(right)\n\n        # 2. Concatenate Left and Right tokens along sequence dimension\n        #    (B, N, D) + (B, N, D) -> (B, 2N, D)\n        x_cat = torch.cat([x_l, x_r], dim=1)\n\n        # 3. Apply Mamba Fusion\n        #    This allows tokens from Left image to interact with tokens from Right image\n        x_fused = self.fusion(x_cat)\n\n        # 4. Global Pooling\n        #    (B, 2N, D) -> (B, D, 2N) -> (B, D, 1) -> (B, D)\n        x_pool = self.pool(x_fused.transpose(1, 2)).flatten(1)\n\n        # 5. Prediction Heads\n        green = self.head_green_raw(x_pool)\n        clover = self.head_clover_raw(x_pool)\n        dead = self.head_dead_raw(x_pool)\n\n        # Summation logic\n        gdm = green + clover\n        total = gdm + dead\n\n        return total, gdm, green, clover, dead\n\ndef get_val_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef clean_image(img):\n    # 1. Safe Crop (Remove artifacts at the bottom)\n    h, w = img.shape[:2]\n    # Cut bottom 10% where artifacts often appear\n    img = img[0:int(h * 0.90), :]\n\n    # 2. Inpaint Date Stamp (Remove orange text)\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    # Define orange color range (adjust as needed)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n\n    # Dilate mask to cover text edges and reduce noise\n    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)\n\n    # Inpaint if mask is not empty\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n\n    return img\n\ndef inference(model, loader, device):\n    model.eval()\n    itera = tqdm(loader, desc='train', leave=False) if CFG.USE_TQDM else loader\n    dry_clovers = []\n    dry_deads = []\n    dry_greens = []\n    gdms = []\n    dry_totals = []\n    for i, (l, r) in enumerate(itera):\n        bs = l.size(0)\n        l = l.to(device, non_blocking=True)\n        r = r.to(device, non_blocking=True)\n        with torch.no_grad():\n            p_total, p_gdm, p_green, p_clover, p_dead = model(l, r)\n            dry_clovers.append(p_clover.cpu().item() if p_clover.cpu().item() > 0 else 0)\n            dry_deads.append(p_dead.cpu().item() if p_dead.cpu().item() > 0 else 0)\n            dry_greens.append(p_green.cpu().item() if p_green.cpu().item() > 0 else 0)\n            gdms.append(p_gdm.cpu().item())\n            dry_totals.append(p_total.cpu().item())\n            \n    return dry_clovers, dry_deads, dry_greens, gdms, dry_totals\n\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df\n        self.transform = transform\n        self.img_dir = img_dir\n        self.paths = df['image_path'].unique().tolist()\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.img_dir, os.path.basename(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        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n        left = self.transform(image=left)['image']\n        right = self.transform(image=right)['image']\n        return left, right\n\n\ncheckpoint = torch.load(\n    f'/kaggle/input/vit-pretrain/pytorch/v1/7/model/best_model_vit_huge_plus_patch16_dinov3.lvd1689m.pth',\n    map_location=CFG.DEVICE, weights_only=False)\n\ntest_df = pd.read_csv('/kaggle/input/csiro-biomass/test.csv')\n\nDL_KW = dict(\n    num_workers=CFG.NUM_WORKERS,\n    pin_memory=True,\n    persistent_workers=(CFG.NUM_WORKERS > 0),\n    prefetch_factor=4 if CFG.NUM_WORKERS > 0 else None,\n)\n\ntest_datasets = BiomassDataset(\n    test_df,\n    get_val_transforms(),\n    '/kaggle/input/csiro-biomass/test'\n)\n\ntest_loader = DataLoader(\n    test_datasets,\n    batch_size=CFG.BATCH_SIZE,\n    shuffle=False,\n    num_workers=0\n)\nbackbone_path = None\nmodel = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED, backbone_path=backbone_path)\nmodel = nn.DataParallel(model).to(CFG.DEVICE)\nmodel.load_state_dict(checkpoint['model_state_dict'])\n\ndry_clovers, dry_deads, dry_greens, gdms, dry_totals = inference(model, test_loader, CFG.DEVICE)\nprint(len(test_df['sample_id'].values))\nprint(len([v for pair in zip(dry_clovers, dry_deads, dry_greens, dry_totals, gdms) for v in pair]))\nresult = pd.DataFrame({\n    'sample_id': test_df['sample_id'].values,\n    'target': [v for pair in zip(dry_clovers, dry_deads, dry_greens, dry_totals, gdms) for v in pair]\n})\nresult.to_csv('submission.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T04:55:49.692302Z","iopub.execute_input":"2026-01-23T04:55:49.692911Z","execution_failed":"2026-01-23T04:56:14.076Z"}},"outputs":[],"execution_count":null}]}