{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13586352,"sourceType":"datasetVersion","datasetId":8630412}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# [LB:0.59] CSIRO | Training & Inference Baseline\n\n## Approach\n\n- **Dual-stream architecture**: Split images into left and right halves for dual-stream processing\n- **Backbone**: ConvNeXt Tiny (timm)\n- **Output**: 3-head output (Dry_Total_g, GDM_g, Dry_Green_g)\n- **Training**: K-fold cross-validation (5 folds, stratified by Dry_Total_g)\n- **Evaluation**: Weighted R² score (CSIRO official evaluation metric)\n- **Ensemble**: Automatic ensemble of all fold models\n- **TTA**: Test-time augmentation support (horizontal/vertical flips)\n\n## Implementation Summary\n\n- Set `CFG.INFERENCE_MODE = True` to switch to inference mode\n- **Training**: Automatically trains all folds and saves `best_model_fold0-4.pth`\n- **Inference**: Automatically detects all fold models and ensembles them\n- `CFG.INFERENCE_MODEL_DIR` allows specifying the model loading directory\n- Set `CFG.USE_TTA = True` to enable TTA\n\n## Leaderboard Scores\n\n- **Public LB**: 0.59 (2025/11/2)","metadata":{}},{"cell_type":"code","source":"# ===============================================================\n# Training for CSIRO Biomass (Dual-stream, 3-head)\n# k-fold / AMP / Cosine / EMA / Constraint penalty\n# ===============================================================\nimport os, gc, math, random, time\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport cv2\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR\nfrom albumentations import (\n    Compose, Resize, HorizontalFlip, VerticalFlip, RandomRotate90,\n    ShiftScaleRotate, RandomBrightnessContrast, HueSaturationValue,\n    RandomResizedCrop, CoarseDropout, Normalize\n)\nfrom albumentations.pytorch import ToTensorV2\nimport timm\n\n# -------------------------\n# 1) Config\n# -------------------------\nclass CFG:\n    # paths\n    SELECT_BEST_BY = 'r2' \n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    TRAIN_CSV = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    OUT_DIR = '.'           # 重み・OOFの保存先\n    os.makedirs(OUT_DIR, exist_ok=True)\n\n    # model\n    MODEL_NAME = 'convnext_tiny'\n    IMG_SIZE = 512\n    IN_CHANS = 3\n    DUAL_STREAM = True  # True: 左右二流, False: 単一ストリーム（全画像）\n\n    # folds\n    N_FOLDS = 5\n    SEED = 42\n\n    # train setup\n    EPOCHS = 80\n    BATCH_SIZE = 16\n    NUM_WORKERS = 2\n    LR = 3e-4\n    WD = 0.05\n    WARMUP_EPOCHS = 1\n    GRAD_ACCUM = 1\n    MAX_NORM = 1.0\n    USE_AMP = True\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n    # loss / targets\n    TARGETS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    TRAIN_FIVE_OUTPUT_LOSS = False   # 5出力（Green, Dead, Clover, GDM, Total）で学習するか\n    USE_LOG1P = False             # log1p変換で安定化\n\n    # ema\n    USE_EMA = False\n    EMA_DECAY = 0.999\n\n    # deterministic\n    DETERMINISTIC = True\n\n    # inference\n    INFERENCE_MODE =  True  # False=学習, True=推論\n    INFERENCE_MODEL_DIR = \"/kaggle/input/convnext-tiny\"  # NoneならOUT_DIRを使用、指定時はそのディレクトリからモデルを読み込み\n    INFERENCE_BATCH_SIZE = 32\n    USE_TTA = False  # Test-time augmentation\n    TEST_CSV = os.path.join(BASE_PATH, 'test.csv')\n    TEST_IMAGE_DIR = os.path.join(BASE_PATH, 'test')\n    SUBMISSION_OUTPUT = os.path.join(OUT_DIR, 'submission.csv')\n    INFERENCE_FOLDS = None  # Noneなら全fold自動検出、リスト指定時はそのfoldのみ使用\n\nprint(f\"Device: {CFG.DEVICE}\")\n\ndef set_seed(seed=42, deterministic=True):\n    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    if deterministic:\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\nset_seed(CFG.SEED, CFG.DETERMINISTIC)\n\n# -------------------------\n# 2) Augmentations\n# -------------------------\ndef get_train_tf(img_size, aug_strength=1.0):\n    \"\"\"\n    Get training augmentations with adjustable strength.\n    \n    Args:\n        img_size: Image size\n        aug_strength: Augmentation strength multiplier (1.0 = default, 0.0 = no augmentation, >1.0 = stronger)\n    \"\"\"\n    # Scale augmentation parameters by strength\n    shift_limit = 0.02 * aug_strength\n    scale_limit = 0.1 * aug_strength\n    rotate_limit = int(10 * aug_strength)\n    hue_shift = int(10 * aug_strength)\n    sat_shift = int(10 * aug_strength)\n    val_shift = int(10 * aug_strength)\n    brightness_limit = 0.15 * aug_strength\n    contrast_limit = 0.15 * aug_strength\n    dropout_p = min(0.3 * aug_strength, 1.0)\n    \n    return Compose([\n        RandomResizedCrop(size=(img_size, img_size), scale=(0.85, 1.0), ratio=(0.95, 1.05), p=1.0),\n        HorizontalFlip(p=0.5),\n        VerticalFlip(p=0.2),\n        RandomRotate90(p=0.2),\n        ShiftScaleRotate(shift_limit=shift_limit, scale_limit=scale_limit, rotate_limit=rotate_limit, \n                        border_mode=cv2.BORDER_REFLECT_101, p=0.5),\n        HueSaturationValue(hue_shift_limit=hue_shift, sat_shift_limit=sat_shift, val_shift_limit=val_shift, p=0.3),\n        RandomBrightnessContrast(brightness_limit=brightness_limit, contrast_limit=contrast_limit, p=0.3),\n        CoarseDropout(max_holes=4, max_height=int(img_size*0.08), max_width=int(img_size*0.08),\n                      min_holes=1, fill_value=0, p=dropout_p),\n        Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n        ToTensorV2()\n    ], additional_targets={'image_right': 'image'} if CFG.DUAL_STREAM else {})\n\ndef get_valid_tf(img_size):\n    return Compose([\n        Resize(img_size, img_size),\n        Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n        ToTensorV2()\n    ], additional_targets={'image_right': 'image'} if CFG.DUAL_STREAM else {})\n\n# -------------------------\n# 3) Dataset (dual-stream: left/right)\n# -------------------------\nclass TrainDataset(Dataset):\n    def __init__(self, df, image_dir, tf, use_log1p=True):\n        self.df = df.reset_index(drop=True)\n        self.paths = self.df['image_path'].values\n        self.y = self.df[CFG.TARGETS].values.astype(np.float32)\n        self.image_dir = image_dir\n        self.tf = tf\n        self.use_log1p = use_log1p\n\n    def __len__(self): return len(self.df)\n\n    def __getitem__(self, idx):\n        raw_path = self.paths[idx]\n        # Use provided path if it exists; otherwise, fall back to joining with image_dir and basename\n        candidate = raw_path if os.path.exists(raw_path) else os.path.join(self.image_dir, os.path.basename(raw_path))\n        img = cv2.imread(candidate)\n        if img is None:\n            img = np.zeros((1000,2000,3), np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        if CFG.DUAL_STREAM:\n            h, w, _ = img.shape\n            mid = w//2\n            left = img[:, :mid]\n            right = img[:, mid:]\n            t = self.tf(image=left, image_right=right)\n            left = t['image']\n            right = t['image_right']\n        else:\n            t = self.tf(image=img)\n            left = t['image']\n            right = left  # インターフェース維持：モデルが単一ストリーム時は右を無視\n\n        target = self.y[idx].copy()\n        if self.use_log1p:\n            target = np.log1p(target)  # 安定化\n        target = torch.from_numpy(target)  # [3]\n\n        return left, right, target\n\n# -------------------------\n# 4) Model (dual-stream, separate heads per target)\n# -------------------------\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name='convnext_tiny', pretrained=True, target_names=None, dual_stream=True, dropout=0.3):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg', in_chans=CFG.IN_CHANS)\n        self.target_names = target_names if target_names is not None else ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n        self.num_outputs = len(self.target_names)\n        self.dual_stream = dual_stream\n        self.dropout = dropout\n        nf = self.backbone.num_features\n        self.n_combined_features = nf * 2 if self.dual_stream else nf\n        \n        # 各ターゲットごとに独立したheadを作成\n        for target_name in self.target_names:\n            head = nn.Sequential(\n                nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n                nn.ReLU(),\n                nn.Dropout(dropout),\n                nn.Linear(self.n_combined_features // 2, 1)\n            )\n            setattr(self, f'head_{target_name.lower().replace(\"_\", \"\")}', head)\n\n    def forward(self, l, r=None):\n        fl = self.backbone(l)\n        if self.dual_stream:\n            fr = self.backbone(r)\n            x = torch.cat([fl, fr], dim=1)\n        else:\n            x = fl\n        \n        # 各headから出力を取得して結合\n        outputs = []\n        for target_name in self.target_names:\n            head = getattr(self, f'head_{target_name.lower().replace(\"_\", \"\")}')\n            out = head(x).squeeze(1)  # [B,1] -> [B]\n            outputs.append(out)\n        \n        return torch.stack(outputs, dim=1)  # [B, num_outputs]\n\n# -------------------------\n# 5) Loss / Metrics\n# -------------------------\nclass WeightedMSELoss(nn.Module):\n    \"\"\"\n    Weighted MSE loss for 3 targets (Total, GDM, Green).\n    \"\"\"\n    def __init__(self, weights):\n        super(WeightedMSELoss, self).__init__()\n        self.weights = torch.tensor(weights, dtype=torch.float32)\n\n    def forward(self, predictions, targets):\n        self.weights = self.weights.to(predictions.device)\n        # predictions/targets: [B, 3]\n        mse_per_target = (predictions - targets) ** 2\n        weighted_mse = mse_per_target * self.weights.unsqueeze(0)\n        return weighted_mse.mean()\n\nclass ConstraintLoss(nn.Module):\n    \"\"\"\n    主損失: L1(予測, 正解)\n    補助: 物理制約違反ペナルティ\n      - monotonic: Total >= GDM >= Green  （違反分だけReLUで加算）\n      - non-neg:   各出力 >= 0            （違反分だけReLUで加算）\n    ログ空間で学習している場合、逆変換 expm1 してから制約評価。\n    \"\"\"\n    def __init__(self, l1_w=1.0, cons_w=0.1, nonneg_w=0.05, use_log1p=True):\n        super().__init__()\n        self.l1 = nn.L1Loss()\n        self.l1_w = l1_w\n        self.cons_w = cons_w\n        self.nonneg_w = nonneg_w\n        self.use_log1p = use_log1p\n\n    def forward(self, pred, target):\n        # pred/target: tuple of three tensors (B,)\n        pT, pGDM, pGR = pred\n        tT, tGDM, tGR = target[:,0], target[:,1], target[:,2]\n\n        # 主損失：ログ空間ならそのままL1\n        loss_main = self.l1(pT, tT) + self.l1(pGDM, tGDM) + self.l1(pGR, tGR)\n        loss_main = self.l1_w * loss_main / 3.0\n\n        # 制約は実数空間で評価\n        if self.use_log1p:\n            PT = torch.expm1(pT)\n            PG = torch.expm1(pGDM)\n            PR = torch.expm1(pGR)\n        else:\n            PT, PG, PR = pT, pGDM, pGR\n\n        zero = torch.zeros_like(PT)\n        # monotonic violation\n        v1 = torch.relu(PG - PT)   # want PT >= PG\n        v2 = torch.relu(PR - PG)   # want PG >= PR\n        loss_cons = (v1 + v2).mean() * self.cons_w\n\n        # non-negative violation\n        n1 = torch.relu(-PT); n2 = torch.relu(-PG); n3 = torch.relu(-PR)\n        loss_nonneg = (n1 + n2 + n3).mean() * self.nonneg_w\n\n        return loss_main + loss_cons + loss_nonneg\n\ndef rmse_torch(y_pred, y_true):\n    return torch.sqrt(torch.mean((y_pred - y_true) ** 2))\n\ndef metric_rmse(pred_tuple, target, use_log1p=True):\n    pT, pGDM, pGR = pred_tuple\n    tT, tGDM, tGR = target[:,0], target[:,1], target[:,2]\n    if use_log1p:\n        pT, pGDM, pGR = [torch.expm1(x) for x in (pT, pGDM, pGR)]\n        tT, tGDM, tGR = [torch.expm1(x) for x in (tT, tGDM, tGR)]\n    rmse_T = rmse_torch(pT, tT)\n    rmse_G = rmse_torch(pGDM, tGDM)\n    rmse_R = rmse_torch(pGR, tGR)\n    return (rmse_T + rmse_G + rmse_R) / 3.0, (rmse_T, rmse_G, rmse_R)\n\n\n# =========================\n# CSIRO weighted R2 metric\n# =========================\nCSIRO_WEIGHTS = {\n    'Dry_Green_g': 0.10,\n    'Dry_Dead_g':  0.10,\n    'Dry_Clover_g':0.10,\n    'GDM_g':       0.20,\n    'Dry_Total_g': 0.50,\n}\n\n# 5出力の順序（学習用）\nFIVE_TARGET_ORDER = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n\ndef _r2_1d(y_true: np.ndarray, y_pred: np.ndarray) -> float:\n    y_true = np.asarray(y_true, dtype=float)\n    y_pred = np.asarray(y_pred, dtype=float)\n    m = np.isfinite(y_true) & np.isfinite(y_pred)\n    y_true = y_true[m]; y_pred = y_pred[m]\n    if y_true.size == 0: return np.nan\n    rss = np.sum((y_true - y_pred) ** 2)\n    tss = np.sum((y_true - y_true.mean()) ** 2)\n    if tss <= 0:\n        return 1.0 if np.allclose(y_true, y_pred) else 0.0\n    return 1.0 - rss / tss\n\ndef _five_from_three(total, gdm, green):\n    clover = gdm - green\n    dead   = total - gdm\n    return {\n        'Dry_Green_g':  green,\n        'Dry_Clover_g': clover,\n        'Dry_Dead_g':   dead,\n        'GDM_g':        gdm,\n        'Dry_Total_g':  total,\n    }\n\ndef csiro_weighted_r2_from_three_tensors(\n    p_total, p_gdm, p_green,\n    t_total, t_gdm, t_green,\n    use_log1p=True\n):\n    \"\"\"\n    3ターゲットのTensor([N])から、5ターゲットに拡張して重み付きR²を返す。\n    戻り: overall(float), per_target(dict)\n    \"\"\"\n    to_np = lambda x: x.detach().cpu().numpy()\n    # log1p学習なら実数空間に戻してから評価\n    if use_log1p:\n        p_total, p_gdm, p_green = [torch.expm1(x) for x in (p_total, p_gdm, p_green)]\n        t_total, t_gdm, t_green = [torch.expm1(x) for x in (t_total, t_gdm, t_green)]\n\n    y_true = _five_from_three(to_np(t_total), to_np(t_gdm), to_np(t_green))\n    y_pred = _five_from_three(to_np(p_total), to_np(p_gdm), to_np(p_green))\n\n    per = {k: _r2_1d(y_true[k], y_pred[k]) for k in y_true.keys()}\n    wsum = sum(CSIRO_WEIGHTS.values())\n    overall = float(np.nansum([CSIRO_WEIGHTS[k]/wsum * per[k] for k in per.keys()]))\n    return overall, per\n\n\ndef csiro_weighted_r2_from_five_tensors(\n    pred5: torch.Tensor,\n    true5: torch.Tensor,\n    columns=FIVE_TARGET_ORDER,\n    use_log1p=True\n):\n    \"\"\"\n    5ターゲット（列順 columns）から重み付きR²を返す。\n    戻り: overall(float), per_target(dict)\n    \"\"\"\n    if use_log1p:\n        pred5 = torch.expm1(pred5)\n        true5 = torch.expm1(true5)\n\n    to_np = lambda x: x.detach().cpu().numpy()\n    p = to_np(pred5)\n    t = to_np(true5)\n\n    per = {}\n    for j, name in enumerate(columns):\n        per[name] = _r2_1d(t[:, j], p[:, j])\n    wsum = sum(CSIRO_WEIGHTS.values())\n    overall = float(np.nansum([CSIRO_WEIGHTS[k]/wsum * per[k] for k in columns]))\n    return overall, per\n\n\ndef build_five_from_three_tensors(pred_tuple, target, use_log1p=True):\n    \"\"\"\n    3出力のテンソルから、学習用の5出力テンソルを構築する。\n    - 入力: pred_tuple = (pT, pGDM, pGR) 各[B]\n            target: [B,3] (列順は Total, GDM, Green)\n    - 変換は実数空間で行う（log1p学習時は expm1 で戻す）\n    - 出力: (pred5[B,5], target5[B,5]) with order = FIVE_TARGET_ORDER\n    \"\"\"\n    pT, pGDM, pGR = pred_tuple\n    tT, tGDM, tGR = target[:,0], target[:,1], target[:,2]\n\n    if use_log1p:\n        PT = torch.expm1(pT); PG = torch.expm1(pGDM); PR = torch.expm1(pGR)\n        TT = torch.expm1(tT); TG = torch.expm1(tGDM); TR = torch.expm1(tGR)\n    else:\n        PT, PG, PR = pT, pGDM, pGR\n        TT, TG, TR = tT, tGDM, tGR\n\n    # 5出力（Green, Dead, Clover, GDM, Total）\n    pred_dead = PT - PG\n    pred_clover = PG - PR\n    tgt_dead = TT - TG\n    tgt_clover = TG - TR\n\n    pred_map = {\n        'Dry_Green_g': PR,\n        'Dry_Dead_g': pred_dead,\n        'Dry_Clover_g': pred_clover,\n        'GDM_g': PG,\n        'Dry_Total_g': PT,\n    }\n    tgt_map = {\n        'Dry_Green_g': TR,\n        'Dry_Dead_g': tgt_dead,\n        'Dry_Clover_g': tgt_clover,\n        'GDM_g': TG,\n        'Dry_Total_g': TT,\n    }\n\n    pred5 = torch.stack([pred_map[k] for k in FIVE_TARGET_ORDER], dim=1)\n    tgt5 = torch.stack([tgt_map[k] for k in FIVE_TARGET_ORDER], dim=1)\n    return pred5, tgt5\n\n\n# -------------------------\n# 6) EMA\n# -------------------------\nclass ModelEMA:\n    def __init__(self, model, decay=0.999):\n        self.ema = BiomassModel(CFG.MODEL_NAME, pretrained=False, target_names=model.target_names, \n                                dual_stream=model.dual_stream, dropout=model.dropout).to(CFG.DEVICE)\n        self.ema.load_state_dict(model.state_dict())\n        self.ema.eval()\n        self.decay = decay\n        for p in self.ema.parameters():\n            p.requires_grad_(False)\n\n    @torch.no_grad()\n    def update(self, model):\n        d = self.decay\n        msd = model.state_dict()\n        for k, v in self.ema.state_dict().items():\n            v.copy_(v * d + (1. - d) * msd[k])\n\n# -------------------------\n# 7) Utilities\n# -------------------------\ndef kfold_split(df, n_folds=5, seed=42):\n    from sklearn.model_selection import KFold\n    df = df.sample(frac=1.0, random_state=seed).reset_index(drop=True)\n    kf = KFold(n_splits=n_folds, shuffle=True, random_state=seed)\n    df['fold'] = -1\n    for i, (_, val_idx) in enumerate(kf.split(df)):\n        df.loc[val_idx, 'fold'] = i\n    return df\n\ndef save_checkpoint(model, path):\n    sd = model.state_dict()\n    torch.save(sd, path)\n\n# -------------------------\n# 8) Train One Fold\n# -------------------------\ndef train_one_fold(df, fold, lr=None, batch_size=None, wd=None, warmup_epochs=None, dropout=None, aug_strength=None):\n    \"\"\"\n    Train one fold with optional parameter overrides.\n    \n    Args:\n        df: DataFrame with fold column\n        fold: Fold number to train\n        lr: Learning rate override (default: CFG.LR)\n        batch_size: Batch size override (default: CFG.BATCH_SIZE)\n        wd: Weight decay override (default: CFG.WD)\n        warmup_epochs: Warmup epochs override (default: CFG.WARMUP_EPOCHS)\n        dropout: Dropout rate override (default: 0.3)\n        aug_strength: Augmentation strength override (default: 1.0)\n    \"\"\"\n    # Use overrides if provided, otherwise use CFG defaults\n    train_lr = lr if lr is not None else CFG.LR\n    train_batch_size = batch_size if batch_size is not None else CFG.BATCH_SIZE\n    train_wd = wd if wd is not None else CFG.WD\n    train_warmup_epochs = warmup_epochs if warmup_epochs is not None else CFG.WARMUP_EPOCHS\n    train_dropout = dropout if dropout is not None else 0.3\n    train_aug_strength = aug_strength if aug_strength is not None else 1.0\n    \n    print(f\"\\n===== FOLD {fold} / {CFG.N_FOLDS} =====\")\n    print(f\"  LR={train_lr:.2e}, BATCH_SIZE={train_batch_size}, WD={train_wd:.4f}\")\n    print(f\"  WARMUP_EPOCHS={train_warmup_epochs}, DROPOUT={train_dropout:.2f}, AUG_STRENGTH={train_aug_strength:.2f}\")\n    trn_df = df[df['fold'] != fold].reset_index(drop=True)\n    val_df = df[df['fold'] == fold].reset_index(drop=True)\n\n    train_tf = get_train_tf(CFG.IMG_SIZE, aug_strength=train_aug_strength)\n    valid_tf = get_valid_tf(CFG.IMG_SIZE)\n\n    trn_ds = TrainDataset(trn_df, CFG.TRAIN_IMAGE_DIR, train_tf, use_log1p=CFG.USE_LOG1P)\n    val_ds = TrainDataset(val_df, CFG.TRAIN_IMAGE_DIR, valid_tf, use_log1p=CFG.USE_LOG1P)\n\n    # Deterministic worker seeding\n    def seed_worker(worker_id):\n        s = torch.initial_seed() % 2**32\n        np.random.seed(s)\n        random.seed(s)\n    g = torch.Generator()\n    g.manual_seed(CFG.SEED)\n\n    trn_dl = DataLoader(\n        trn_ds,\n        batch_size=train_batch_size,\n        shuffle=True,\n        num_workers=CFG.NUM_WORKERS,\n        pin_memory=True,\n        drop_last=True,\n        worker_init_fn=seed_worker,\n        generator=g,\n        persistent_workers=(CFG.NUM_WORKERS > 0),\n    )\n    val_dl = DataLoader(\n        val_ds,\n        batch_size=train_batch_size*2,\n        shuffle=False,\n        num_workers=CFG.NUM_WORKERS,\n        pin_memory=True,\n        worker_init_fn=seed_worker,\n        generator=g,\n        persistent_workers=(CFG.NUM_WORKERS > 0),\n    )\n\n    model = BiomassModel(CFG.MODEL_NAME, pretrained=True, target_names=CFG.TARGETS, dual_stream=CFG.DUAL_STREAM, dropout=train_dropout).to(CFG.DEVICE)\n    optimizer = AdamW(model.parameters(), lr=train_lr, weight_decay=train_wd)\n    # True warmup + cosine, stepped per optimizer step\n    steps_per_epoch = max(1, math.ceil(len(trn_dl) / CFG.GRAD_ACCUM))\n    warmup_steps = max(1, train_warmup_epochs * steps_per_epoch)\n    total_steps = max(1, CFG.EPOCHS * steps_per_epoch)\n    cosine_steps = max(1, total_steps - warmup_steps)\n    warmup = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_steps)\n    cosine = CosineAnnealingLR(optimizer, T_max=cosine_steps, eta_min=train_lr*1e-2)\n    scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps])\n    scaler = torch.cuda.amp.GradScaler(enabled=CFG.USE_AMP)\n    ema = ModelEMA(model, decay=CFG.EMA_DECAY) if CFG.USE_EMA else None\n\n    # 損失の重み（CFG.TARGETSの順序に基づく）\n    weights = [CSIRO_WEIGHTS[k] for k in CFG.TARGETS]\n    criterion = WeightedMSELoss(weights=weights)\n\n    select_is_r2 = CFG.SELECT_BEST_BY.lower() == 'r2'\n    best_metric = -float('inf') if select_is_r2 else float('inf')\n    best_preds = None\n\n    global_step = 0\n    for epoch in range(1, CFG.EPOCHS+1):\n        model.train()\n        train_loss = 0.0\n\n        for i, (l, r, y) in enumerate(tqdm(trn_dl, desc=f\"Train ep{epoch}\")):\n            l = l.to(CFG.DEVICE, non_blocking=True); r = r.to(CFG.DEVICE, non_blocking=True); y = y.to(CFG.DEVICE, non_blocking=True)\n            with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                pred = model(l, r)  # [B,K]\n                loss = criterion(pred, y) / CFG.GRAD_ACCUM\n            scaler.scale(loss).backward()\n\n            if (i+1) % CFG.GRAD_ACCUM == 0:\n                if CFG.MAX_NORM is not None:\n                    scaler.unscale_(optimizer)\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.MAX_NORM)\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad(set_to_none=True)\n                if ema: ema.update(model)\n                scheduler.step()\n                global_step += 1\n\n            train_loss += loss.item()\n\n        # ---- validation (EMA優先) ----\n        model.eval()\n        eval_model = ema.ema if ema else model\n        val_loss = 0.0\n        y_pred_all, y_true_all = [], []\n\n        with torch.no_grad():\n            for l, r, y in tqdm(val_dl, desc=\"Valid\"):\n                l = l.to(CFG.DEVICE, non_blocking=True); r = r.to(CFG.DEVICE, non_blocking=True); y = y.to(CFG.DEVICE, non_blocking=True)\n                with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                    pred = eval_model(l, r)  # [B,K]\n                    loss = criterion(pred, y)\n                val_loss += loss.item()\n                y_pred_all.append(pred.detach().cpu())  # [B,K]\n                y_true_all.append(y.detach().cpu())\n\n        y_pred_all = torch.cat(y_pred_all, dim=0)  # [N,K]\n        y_true_all = torch.cat(y_true_all, dim=0)  # [N,K]\n\n        # --- 既存: RMSE（参考） ---\n        # RMSE（参考）: 常に Total/GDM/Green の3指標で計算\n        idxT = CFG.TARGETS.index('Dry_Total_g')\n        idxG = CFG.TARGETS.index('GDM_g')\n        idxR = CFG.TARGETS.index('Dry_Green_g')\n        rmse_mean, (rmse_T, rmse_G, rmse_R) = metric_rmse(\n            (y_pred_all[:,idxT], y_pred_all[:,idxG], y_pred_all[:,idxR]),\n            torch.stack([y_true_all[:,idxT], y_true_all[:,idxG], y_true_all[:,idxR]], dim=1),\n            use_log1p=CFG.USE_LOG1P\n        )\n\n        # --- 新規: コンペ評価(重み付きR²) ---\n        if len(CFG.TARGETS) == 5:\n            wr2, per_r2 = csiro_weighted_r2_from_five_tensors(\n                pred5=y_pred_all,\n                true5=y_true_all,\n                columns=CFG.TARGETS,\n                use_log1p=CFG.USE_LOG1P\n            )\n        else:\n            wr2, per_r2 = csiro_weighted_r2_from_three_tensors(\n                p_total=y_pred_all[:,0], p_gdm=y_pred_all[:,1], p_green=y_pred_all[:,2],\n                t_total=y_true_all[:,0], t_gdm=y_true_all[:,1], t_green=y_true_all[:,2],\n                use_log1p=CFG.USE_LOG1P\n            )\n\n        print(f\"[Fold {fold}] Epoch {epoch}: \"\n              f\"train_loss={train_loss/len(trn_dl):.4f}  \"\n              f\"val_loss={val_loss/len(val_dl):.4f}  \"\n              f\"RMSE_mean={rmse_mean:.4f} (T:{rmse_T:.4f} G:{rmse_G:.4f} R:{rmse_R:.4f})  \"\n              f\"WeightedR2={wr2:.5f}  \"\n              f\"R2(Total:{per_r2['Dry_Total_g']:.3f} GDM:{per_r2['GDM_g']:.3f} \"\n              f\"Green:{per_r2['Dry_Green_g']:.3f} Dead:{per_r2['Dry_Dead_g']:.3f} \"\n              f\"Clover:{per_r2['Dry_Clover_g']:.3f})\"\n        )\n\n        # ---- best の判定（CFG.SELECT_BEST_BY）----\n        current_metric = wr2 if select_is_r2 else rmse_mean.item()\n        improved = (current_metric > best_metric) if select_is_r2 else (current_metric < best_metric)\n        if epoch == 1:\n            improved = True\n\n        if improved:\n            best_path = os.path.join(CFG.OUT_DIR, f'best_model_fold{fold}.pth')\n            save_checkpoint(eval_model, best_path)\n            best_metric = current_metric\n            best_preds = y_pred_all.numpy()  # OOF保存用（ベスト）\n            print(f\"  -> Best updated ({CFG.SELECT_BEST_BY.upper()}). Save {best_path}\")\n\n    \n    # 逆変換してOOF保存\n    # Save OOF from best epoch\n    oof_preds = best_preds if best_preds is not None else y_pred_all.numpy()\n    if CFG.USE_LOG1P:\n        oof_preds = np.expm1(oof_preds)\n    oof_df = val_df[['image_path'] + CFG.TARGETS].copy()\n    for i,t in enumerate(CFG.TARGETS):\n        oof_df[f'pred_{t}'] = oof_preds[:, i]\n    oof_df.to_csv(os.path.join(CFG.OUT_DIR, f'oof_fold{fold}.csv'), index=False)\n    return best_metric\n\n\n\n# ============================================\n# 学習用のピボット & 層化Fold作成ユーティリティ\n# ============================================\nimport os\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import StratifiedKFold\n\nALL_TARGET_COLS = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'GDM_g', 'Dry_Total_g']\nINDEX_COLS = ['image_path', 'Sampling_Date', 'State', 'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm']\n\ndef _dup_check_for_pivot(df_long, index_cols=INDEX_COLS, name_col='target_name'):\n    keys = index_cols + [name_col]\n    dup_mask = df_long.duplicated(keys, keep=False)\n    return df_long.loc[dup_mask, keys].value_counts().reset_index(name='count')\n\ndef long_to_wide_for_training(\n    df_long: pd.DataFrame,\n    targets=('Dry_Total_g','GDM_g','Dry_Green_g'),\n    strict=True,\n    aggfunc='first'\n) -> pd.DataFrame:\n    \"\"\"\n    長形式(train.csv) → 学習用の広形式（画像1行＋3ターゲット列）。\n    余剰の2ターゲット(Dry_Dead/Clover)は保持してもOKだが、学習では未使用。\n    \"\"\"\n    # 1) 重複チェック（必要なら平均等で潰す）\n    if strict:\n        dups = _dup_check_for_pivot(df_long)\n        if len(dups):\n            raise ValueError(\n                f\"Pivot keys have duplicates ({len(dups)} rows). \"\n                f\"Set strict=False or aggfunc='mean'.\\n{dups.head()}\"\n            )\n\n    # 2) ピボット\n    wide = df_long.pivot_table(\n        index=INDEX_COLS,\n        columns='target_name',\n        values='target',\n        aggfunc=aggfunc\n    ).reset_index()\n\n    # 3) 学習で使う3列が無ければエラー\n    for t in targets:\n        if t not in wide.columns:\n            raise KeyError(f\"Required target column missing after pivot: {t}\")\n\n    # 4) image_id 付与（任意）\n    wide['image_id'] = wide['image_path'].apply(lambda x: os.path.splitext(os.path.basename(x))[0])\n\n    # 5) 学習に最低限必要なカラムだけ抜き出し（メタは残してOK）\n    keep_cols = list(INDEX_COLS) + list(targets) + ['image_id']\n    keep_cols = [c for c in keep_cols if c in wide.columns]\n    wide = wide[keep_cols].copy()\n    return wide\n\ndef add_stratified_folds(\n    df: pd.DataFrame,\n    n_folds=5,\n    label_col='Dry_Total_g',\n    bins=5,\n    seed=42\n) -> pd.DataFrame:\n    \"\"\"\n    Dry_Total_g をビニングして層化KFold。ターゲット分布を各foldで均等化しやすい。\n    \"\"\"\n    df = df.copy()\n    # 欠損・定数対応（qcutが失敗しないように）\n    y = df[label_col].values\n    # 一意値が少なければ bin 数を調整\n    uniq = np.unique(y)\n    bins = min(bins, max(2, len(uniq)))\n    # 量子化。重複binに配慮して duplicates='drop'\n    df['_strat'] = pd.qcut(y, q=bins, labels=False, duplicates='drop')\n\n    skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=seed)\n    df['fold'] = -1\n    for f, (_, val_idx) in enumerate(skf.split(df, df['_strat'])):\n        df.loc[val_idx, 'fold'] = f\n    df = df.drop(columns=['_strat'])\n    return df\n\ndef save_training_csv_for_existing_pipeline(df_wide: pd.DataFrame, out_path: str):\n    \"\"\"\n    既存の学習スクリプトが読む `train.csv` 形式に合わせて保存。\n    （= image_path と 3ターゲット列が含まれていればOK）\n    \"\"\"\n    cols_needed = ['image_path', 'Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    missing = [c for c in cols_needed if c not in df_wide.columns]\n    if missing:\n        raise KeyError(f\"Columns missing for training: {missing}\")\n    df_wide.to_csv(out_path, index=False)\n    print(f\"Saved training CSV for pipeline: {out_path}  shape={df_wide.shape}\")\n\n\n# -------------------------\n# 9) Main: k-fold\n# -------------------------\n\ndef main(lr=None, batch_size=None, wd=None, warmup_epochs=None, dropout=None, aug_strength=None):\n    \"\"\"\n    Main training/inference function with optional parameter overrides.\n    \n    Args:\n        lr: Learning rate override (default: CFG.LR)\n        batch_size: Batch size override (default: CFG.BATCH_SIZE)\n        wd: Weight decay override (default: CFG.WD)\n        warmup_epochs: Warmup epochs override (default: CFG.WARMUP_EPOCHS)\n        dropout: Dropout rate override (default: 0.3)\n        aug_strength: Augmentation strength override (default: 1.0)\n    \n    Returns:\n        dict: Results containing 'cv_mean', 'cv_std', and 'fold_scores' (training mode)\n              or DataFrame with predictions (inference mode)\n    \"\"\"\n    # 推論モードの場合\n    if CFG.INFERENCE_MODE:\n        print(\"=\" * 50)\n        print(\"INFERENCE MODE\")\n        print(\"=\" * 50)\n        \n        # モデル読み込みディレクトリを決定\n        model_dir = CFG.INFERENCE_MODEL_DIR if CFG.INFERENCE_MODEL_DIR is not None else CFG.OUT_DIR\n        print(f\"Loading models from: {model_dir}\")\n        \n        # 使用するfoldを決定\n        if CFG.INFERENCE_FOLDS is None:\n            # 全foldを自動検出\n            model_paths = []\n            for fold in range(CFG.N_FOLDS):\n                model_path = os.path.join(model_dir, f'best_model_fold{fold}.pth')\n                if os.path.exists(model_path):\n                    model_paths.append(model_path)\n                else:\n                    print(f\"Warning: Model not found: {model_path}\")\n            \n            if len(model_paths) == 0:\n                raise ValueError(f\"No model files found in {model_dir}\")\n            \n            print(f\"Found {len(model_paths)} model(s): {[os.path.basename(p) for p in model_paths]}\")\n        else:\n            # 指定されたfoldのみ使用\n            model_paths = []\n            for fold in CFG.INFERENCE_FOLDS:\n                model_path = os.path.join(model_dir, f'best_model_fold{fold}.pth')\n                if os.path.exists(model_path):\n                    model_paths.append(model_path)\n                else:\n                    raise FileNotFoundError(f\"Model not found: {model_path}\")\n            print(f\"Using {len(model_paths)} model(s) from folds: {CFG.INFERENCE_FOLDS}\")\n        \n        # 推論実行\n        submission_df = predict_and_save_submission(\n            test_csv_path=CFG.TEST_CSV,\n            output_path=CFG.SUBMISSION_OUTPUT,\n            model_paths=model_paths,\n            image_dir=CFG.TEST_IMAGE_DIR,\n            model_name=CFG.MODEL_NAME,\n            target_names=CFG.TARGETS,\n            dual_stream=CFG.DUAL_STREAM,\n            dropout=dropout if dropout is not None else 0.3,\n            batch_size=CFG.INFERENCE_BATCH_SIZE,\n            device=CFG.DEVICE,\n            num_workers=CFG.NUM_WORKERS,\n            tta=CFG.USE_TTA,\n            weights=None,\n            use_log1p=CFG.USE_LOG1P\n        )\n        \n        print(\"=\" * 50)\n        print(f\"Inference completed. Submission saved to: {CFG.SUBMISSION_OUTPUT}\")\n        print(\"=\" * 50)\n        \n        return submission_df\n    \n    # 学習モード（既存の処理）\n    print(\"=\" * 50)\n    print(\"TRAINING MODE\")\n    print(\"=\" * 50)\n    df = pd.read_csv(CFG.TRAIN_CSV)\n\n    # 2) 学習用に広形式へ（ターゲット列を作る）\n    target_list = FIVE_TARGET_ORDER if CFG.TRAIN_FIVE_OUTPUT_LOSS else ['Dry_Total_g','GDM_g','Dry_Green_g']\n    # CFG に反映（Dataset/OOF保存で使用する列順）\n    CFG.TARGETS = target_list\n    df = long_to_wide_for_training(\n        df,\n        targets=tuple(target_list),\n        strict=True,          # 重複があればエラーにしたくないなら False に\n        aggfunc='first'       # 重複を平均で潰すなら 'mean'\n    )\n    \n    # 3) fold 付与（Dry_Total_g 層化。評価寄与が大きいので相性○）\n    df = add_stratified_folds(\n        df,\n        n_folds=5,\n        label_col='Dry_Total_g',\n        bins=5,\n        seed=42\n    )\n    # 期待カラム: image_path と CFG.TARGETS\n    # 必要に応じてここで欠損処理や外れ値clipを実施\n    assert set(CFG.TARGETS).issubset(df.columns), f\"train.csvに{CFG.TARGETS}が必要です\"\n\n\n    # fold作成（GroupKFoldを使いたい場合は plot_id 等で置き換え）\n    if 'fold' not in df.columns:\n        df = kfold_split(df, n_folds=CFG.N_FOLDS, seed=CFG.SEED)\n\n    df.to_csv(os.path.join(CFG.OUT_DIR, 'train_folds.csv'), index=False)\n    print(\"Folds saved:\", os.path.join(CFG.OUT_DIR, 'train_folds.csv'))\n\n    bests = []\n    for f in range(CFG.N_FOLDS):\n        best_metric = train_one_fold(df, f, lr=lr, batch_size=batch_size, wd=wd, \n                                      warmup_epochs=warmup_epochs, dropout=dropout, aug_strength=aug_strength)\n        bests.append(best_metric)\n        gc.collect()\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n\n    metric_name = 'Weighted R2' if CFG.SELECT_BEST_BY.lower() == 'r2' else 'RMSE'\n    cv_mean = np.mean(bests)\n    cv_std = np.std(bests)\n    print(f\"\\n=== CV {metric_name} (mean±std) ===\")\n    print(f\"{cv_mean:.5f} ± {cv_std:.5f}\")\n    \n    return {\n        'cv_mean': cv_mean,\n        'cv_std': cv_std,\n        'fold_scores': bests\n    }\n\n# -------------------------\n# 10) Inference Functions\n# -------------------------\n\nclass TestDataset(Dataset):\n    \"\"\"Test dataset for inference (no labels).\"\"\"\n    def __init__(self, df, image_dir, tf, dual_stream=True):\n        self.df = df.reset_index(drop=True)\n        self.paths = self.df['image_path'].values\n        self.image_dir = image_dir\n        self.tf = tf\n        self.dual_stream = dual_stream\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        raw_path = self.paths[idx]\n        # Use provided path if it exists; otherwise, fall back to joining with image_dir and basename\n        candidate = raw_path if os.path.exists(raw_path) else os.path.join(self.image_dir, os.path.basename(raw_path))\n        img = cv2.imread(candidate)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        if self.dual_stream:\n            h, w, _ = img.shape\n            mid = w // 2\n            left = img[:, :mid]\n            right = img[:, mid:]\n            t = self.tf(image=left, image_right=right)\n            left = t['image']\n            right = t['image_right']\n        else:\n            t = self.tf(image=img)\n            left = t['image']\n            right = left  # インターフェース維持：モデルが単一ストリーム時は右を無視\n\n        return left, right, idx\n\n\ndef load_model(model_path, model_name=None, target_names=None, dual_stream=None, dropout=0.3, device=None):\n    \"\"\"\n    Load a trained model from checkpoint.\n    \n    Args:\n        model_path: Path to the model checkpoint (.pth file)\n        model_name: Model name (default: CFG.MODEL_NAME)\n        target_names: Target names (default: CFG.TARGETS)\n        dual_stream: Whether to use dual stream (default: CFG.DUAL_STREAM)\n        dropout: Dropout rate (default: 0.3)\n        device: Device to load model on (default: CFG.DEVICE)\n    \n    Returns:\n        Loaded model in eval mode\n    \"\"\"\n    if model_name is None:\n        model_name = CFG.MODEL_NAME\n    if target_names is None:\n        target_names = CFG.TARGETS\n    if dual_stream is None:\n        dual_stream = CFG.DUAL_STREAM\n    if device is None:\n        device = CFG.DEVICE\n    \n    model = BiomassModel(\n        model_name=model_name,\n        pretrained=False,\n        target_names=target_names,\n        dual_stream=dual_stream,\n        dropout=dropout\n    ).to(device)\n    \n    checkpoint = torch.load(model_path, map_location=device)\n    model.load_state_dict(checkpoint)\n    model.eval()\n    \n    return model\n\n\ndef predict_single_image(model, image_path, transform, dual_stream=None, device=None, tta=False):\n    \"\"\"\n    Predict on a single image with optional TTA.\n    \n    Args:\n        model: Trained model\n        image_path: Path to image file\n        transform: Transform to apply\n        dual_stream: Whether to use dual stream (default: CFG.DUAL_STREAM)\n        device: Device to use (default: CFG.DEVICE)\n        tta: Whether to use test-time augmentation (horizontal/vertical flips)\n    \n    Returns:\n        numpy array of predictions [num_targets]\n    \"\"\"\n    if dual_stream is None:\n        dual_stream = CFG.DUAL_STREAM\n    if device is None:\n        device = CFG.DEVICE\n    \n    # Load image\n    img = cv2.imread(image_path)\n    if img is None:\n        img = np.zeros((1000, 2000, 3), np.uint8)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    \n    predictions = []\n    \n    if not tta:\n        # Single prediction\n        if dual_stream:\n            h, w, _ = img.shape\n            mid = w // 2\n            left = img[:, :mid]\n            right = img[:, mid:]\n            t = transform(image=left, image_right=right)\n            left = t['image']\n            right = t['image_right']\n        else:\n            t = transform(image=img)\n            left = t['image']\n            right = left\n        \n        left = left.unsqueeze(0).to(device)\n        right = right.unsqueeze(0).to(device)\n        \n        with torch.no_grad():\n            with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                pred = model(left, right)  # [1, num_targets]\n        \n        predictions.append(pred.cpu().numpy()[0])\n    else:\n        # TTA: original, hflip, vflip, hvflip\n        augs = [\n            {'hflip': False, 'vflip': False},\n            {'hflip': True, 'vflip': False},\n            {'hflip': False, 'vflip': True},\n            {'hflip': True, 'vflip': True},\n        ]\n        \n        for aug in augs:\n            img_aug = img.copy()\n            if aug['hflip']:\n                img_aug = cv2.flip(img_aug, 1)\n            if aug['vflip']:\n                img_aug = cv2.flip(img_aug, 0)\n            \n            if dual_stream:\n                h, w, _ = img_aug.shape\n                mid = w // 2\n                left = img_aug[:, :mid]\n                right = img_aug[:, mid:]\n                t = transform(image=left, image_right=right)\n                left = t['image']\n                right = t['image_right']\n            else:\n                t = transform(image=img_aug)\n                left = t['image']\n                right = left\n            \n            left = left.unsqueeze(0).to(device)\n            right = right.unsqueeze(0).to(device)\n            \n            with torch.no_grad():\n                with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                    pred = model(left, right)  # [1, num_targets]\n            \n            predictions.append(pred.cpu().numpy()[0])\n        \n        # Average predictions\n        predictions = np.mean(predictions, axis=0)\n    \n    return predictions[0] if isinstance(predictions, list) else predictions\n\n\ndef predict_test_set(\n    model,\n    test_df,\n    image_dir,\n    transform,\n    batch_size=32,\n    dual_stream=None,\n    device=None,\n    num_workers=2,\n    tta=False\n):\n    \"\"\"\n    Predict on entire test set.\n    \n    Args:\n        model: Trained model\n        test_df: DataFrame with 'image_path' column\n        image_dir: Directory containing test images\n        transform: Transform to apply\n        batch_size: Batch size for inference\n        dual_stream: Whether to use dual stream (default: CFG.DUAL_STREAM)\n        device: Device to use (default: CFG.DEVICE)\n        num_workers: Number of workers for DataLoader\n        tta: Whether to use test-time augmentation\n    \n    Returns:\n        numpy array of predictions [N, num_targets]\n    \"\"\"\n    if dual_stream is None:\n        dual_stream = CFG.DUAL_STREAM\n    if device is None:\n        device = CFG.DEVICE\n    \n    dataset = TestDataset(test_df, image_dir, transform, dual_stream=dual_stream)\n    dataloader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=num_workers,\n        pin_memory=True,\n        persistent_workers=(num_workers > 0),\n    )\n    \n    all_predictions = []\n    model.eval()\n    \n    with torch.no_grad():\n        for left, right, indices in tqdm(dataloader, desc=\"Predicting\"):\n            left = left.to(device, non_blocking=True)\n            right = right.to(device, non_blocking=True)\n            \n            if tta:\n                # TTA: original, hflip, vflip, hvflip\n                preds_tta = []\n                for hflip, vflip in [(False, False), (True, False), (False, True), (True, True)]:\n                    left_aug = torch.flip(left, dims=[3] if hflip else []) if hflip else left\n                    left_aug = torch.flip(left_aug, dims=[2] if vflip else []) if vflip else left_aug\n                    right_aug = torch.flip(right, dims=[3] if hflip else []) if hflip else right\n                    right_aug = torch.flip(right_aug, dims=[2] if vflip else []) if vflip else right_aug\n                    \n                    with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                        pred = model(left_aug, right_aug)  # [B, num_targets]\n                    preds_tta.append(pred)\n                \n                pred = torch.stack(preds_tta, dim=0).mean(dim=0)  # Average over TTA\n            else:\n                with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                    pred = model(left, right)  # [B, num_targets]\n            \n            all_predictions.append(pred.cpu().numpy())\n    \n    return np.concatenate(all_predictions, axis=0)\n\n\ndef ensemble_predict(\n    test_df,\n    model_paths,\n    image_dir,\n    model_name=None,\n    target_names=None,\n    dual_stream=None,\n    dropout=0.3,\n    batch_size=32,\n    device=None,\n    num_workers=2,\n    tta=False,\n    weights=None\n):\n    \"\"\"\n    Ensemble predictions from multiple models (typically different folds).\n    \n    Args:\n        test_df: DataFrame with 'image_path' column\n        model_paths: List of paths to model checkpoints\n        image_dir: Directory containing test images\n        model_name: Model name (default: CFG.MODEL_NAME)\n        target_names: Target names (default: CFG.TARGETS)\n        dual_stream: Whether to use dual stream (default: CFG.DUAL_STREAM)\n        dropout: Dropout rate (default: 0.3)\n        batch_size: Batch size for inference\n        device: Device to use (default: CFG.DEVICE)\n        num_workers: Number of workers for DataLoader\n        tta: Whether to use test-time augmentation\n        weights: Weights for each model (default: equal weights)\n    \n    Returns:\n        numpy array of ensemble predictions [N, num_targets]\n    \"\"\"\n    if model_name is None:\n        model_name = CFG.MODEL_NAME\n    if target_names is None:\n        target_names = CFG.TARGETS\n    if dual_stream is None:\n        dual_stream = CFG.DUAL_STREAM\n    if device is None:\n        device = CFG.DEVICE\n    if weights is None:\n        weights = [1.0] * len(model_paths)\n    \n    # Normalize weights\n    weights = np.array(weights)\n    weights = weights / weights.sum()\n    \n    transform = get_valid_tf(CFG.IMG_SIZE)\n    \n    all_predictions = []\n    \n    for i, model_path in enumerate(model_paths):\n        print(f\"Loading model {i+1}/{len(model_paths)}: {model_path}\")\n        model = load_model(\n            model_path,\n            model_name=model_name,\n            target_names=target_names,\n            dual_stream=dual_stream,\n            dropout=dropout,\n            device=device\n        )\n        \n        pred = predict_test_set(\n            model=model,\n            test_df=test_df,\n            image_dir=image_dir,\n            transform=transform,\n            batch_size=batch_size,\n            dual_stream=dual_stream,\n            device=device,\n            num_workers=num_workers,\n            tta=tta\n        )\n        \n        all_predictions.append(pred * weights[i])\n        \n        # Clean up\n        del model\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n    \n    ensemble_pred = np.sum(all_predictions, axis=0)\n    return ensemble_pred\n\n\ndef predict_and_save_submission(\n    test_csv_path,\n    output_path,\n    model_paths,\n    image_dir=None,\n    model_name=None,\n    target_names=None,\n    dual_stream=None,\n    dropout=0.3,\n    batch_size=None,\n    device=None,\n    num_workers=None,\n    tta=None,\n    weights=None,\n    use_log1p=None\n):\n    \"\"\"\n    Predict on test set and save submission file in the required format.\n    \n    Args:\n        test_csv_path: Path to test.csv (long format)\n        output_path: Path to save submission CSV\n        model_paths: List of paths to model checkpoints (or single path)\n        image_dir: Directory containing test images (default: inferred from BASE_PATH)\n        model_name: Model name (default: CFG.MODEL_NAME)\n        target_names: Target names (default: CFG.TARGETS)\n        dual_stream: Whether to use dual stream (default: CFG.DUAL_STREAM)\n        dropout: Dropout rate (default: 0.3)\n        batch_size: Batch size for inference\n        device: Device to use (default: CFG.DEVICE)\n        num_workers: Number of workers for DataLoader\n        tta: Whether to use test-time augmentation\n        weights: Weights for each model (default: equal weights)\n        use_log1p: Whether model was trained with log1p (default: CFG.USE_LOG1P)\n    \n    Returns:\n        DataFrame with predictions\n    \"\"\"\n    if image_dir is None:\n        image_dir = CFG.TEST_IMAGE_DIR\n    if model_name is None:\n        model_name = CFG.MODEL_NAME\n    if target_names is None:\n        target_names = CFG.TARGETS\n    if dual_stream is None:\n        dual_stream = CFG.DUAL_STREAM\n    if device is None:\n        device = CFG.DEVICE\n    if use_log1p is None:\n        use_log1p = CFG.USE_LOG1P\n    if batch_size is None:\n        batch_size = CFG.INFERENCE_BATCH_SIZE\n    if num_workers is None:\n        num_workers = CFG.NUM_WORKERS\n    if tta is None:\n        tta = CFG.USE_TTA\n    \n    # Load test CSV (long format)\n    test_df_long = pd.read_csv(test_csv_path)\n    \n    # Get unique image paths\n    unique_images = test_df_long['image_path'].unique()\n    test_df = pd.DataFrame({'image_path': unique_images})\n    \n    # Convert model_paths to list if single path\n    if isinstance(model_paths, str):\n        model_paths = [model_paths]\n    \n    # Get predictions\n    print(f\"Predicting on {len(test_df)} images using {len(model_paths)} model(s)...\")\n    predictions = ensemble_predict(\n        test_df=test_df,\n        model_paths=model_paths,\n        image_dir=image_dir,\n        model_name=model_name,\n        target_names=target_names,\n        dual_stream=dual_stream,\n        dropout=dropout,\n        batch_size=batch_size,\n        device=device,\n        num_workers=num_workers,\n        tta=tta,\n        weights=weights\n    )\n    \n    # Convert log1p predictions back to original scale\n    if use_log1p:\n        predictions = np.expm1(predictions)\n    \n    # Create prediction DataFrame\n    pred_df = pd.DataFrame(\n        predictions,\n        columns=target_names,\n        index=test_df.index\n    )\n    pred_df['image_path'] = test_df['image_path'].values\n    \n    # Convert to long format for submission\n    submission_rows = []\n    for _, row in test_df_long.iterrows():\n        image_path = row['image_path']\n        target_name = row['target_name']\n        \n        # Find corresponding prediction\n        pred_row = pred_df[pred_df['image_path'] == image_path].iloc[0]\n        \n        # Get the 3 main targets\n        total = pred_row['Dry_Total_g']\n        gdm = pred_row['GDM_g']\n        green = pred_row['Dry_Green_g']\n        \n        # Derive all 5 targets\n        if target_name == 'Dry_Total_g':\n            value = total\n        elif target_name == 'GDM_g':\n            value = gdm\n        elif target_name == 'Dry_Green_g':\n            value = green\n        elif target_name == 'Dry_Dead_g':\n            value = total - gdm\n        elif target_name == 'Dry_Clover_g':\n            value = gdm - green\n        else:\n            value = 0.0\n        \n        submission_rows.append({\n            'sample_id': row['sample_id'],\n            'target': max(0.0, value)  # Ensure non-negative\n        })\n    \n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv(output_path, index=False)\n    print(f\"Submission saved to: {output_path}\")\n    \n    return submission_df\nmain()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T00:44:28.236187Z","iopub.execute_input":"2025-11-03T00:44:28.236822Z","iopub.status.idle":"2025-11-03T00:44:35.213107Z","shell.execute_reply.started":"2025-11-03T00:44:28.236798Z","shell.execute_reply":"2025-11-03T00:44:35.212386Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}