{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14091466,"sourceType":"datasetVersion","datasetId":8972849},{"sourceId":14107826,"sourceType":"datasetVersion","datasetId":8985968},{"sourceId":14110634,"sourceType":"datasetVersion","datasetId":8988221},{"sourceId":14112856,"sourceType":"datasetVersion","datasetId":8989960},{"sourceId":14124371,"sourceType":"datasetVersion","datasetId":8998984}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, gc, math, random, time\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\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)\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom sklearn.model_selection import StratifiedKFold\nfrom albumentations import CLAHE\n# =========================================================================================\n# 1. Global Configuration\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 = '.'\n    if not os.path.exists(OUT_DIR):\n        os.makedirs(OUT_DIR, exist_ok=True)\n\n    # model\n    MODEL_NAME = 'swinv2_large_window12to24_192to384.ms_in22k_ft_in1k'\n    IMG_SIZE = 768\n    IN_CHANS = 3\n    DUAL_STREAM = True\n    PRETRAINED_PATH = \"\"\n    \n    # folds\n    N_FOLDS = 5 # Used for splitting validation set only\n    SEED = 42\n\n    # train setup\n    EPOCHS = 80\n    BATCH_SIZE = 4\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\n    USE_LOG1P = False\n\n    # ema\n    USE_EMA = False\n    EMA_DECAY = 0.999\n\n    # deterministic\n    DETERMINISTIC = True\n\n    # inference\n    # ★★★ SET THIS TO True FOR SUBMISSION / False FOR TRAINING ★★★\n    INFERENCE_MODE = True\n    \n    # Path to your trained model weights\n    INFERENCE_MODEL_DIR = \"/kaggle/input/swin-large\" # Assuming model is saved in current dir\n    INFERENCE_BATCH_SIZE = 4\n    \n    # TTA Removed\n    \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    \n    # Explicitly use only Fold 0\n    INFERENCE_FOLDS = [0] \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    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        \n        # ★★★ 验证集同步加入，参数一致 ★★★\n        CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=1.0),\n        \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# 3. Datasets\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        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)\n\n        return left, right, target\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        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# =========================================================================================\n# 4. Model Architecture\n# =========================================================================================\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name=CFG.MODEL_NAME, pretrained=True, \n                 pretrained_path=None,  # <--- 新增参数：接收本地路径\n                 target_names=None, dual_stream=True, dropout=0.3):\n        super().__init__()\n        \n        print(f\"🏗️ Initializing Backbone: {model_name}\")\n        print(f\"📐 Target Image Size: {CFG.IMG_SIZE} x {CFG.IMG_SIZE}\")\n        \n        # 1. 创建骨架 (禁止自动下载)\n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=False,     \n            num_classes=0, \n            global_pool='avg', \n            in_chans=CFG.IN_CHANS,\n            img_size=CFG.IMG_SIZE # 关键：告诉模型我们要用大图\n        )\n        \n        # 2. ★★★ 智能加载权重 (Log-CPB插值) ★★★\n        if pretrained and pretrained_path and os.path.exists(pretrained_path):\n            print(f\"🔄 Smart Loading weights from: {pretrained_path}\")\n            try:\n                load_checkpoint(self.backbone, pretrained_path, strict=False)\n                print(\"✅ Weights loaded & Resized successfully!\")\n            except Exception as e:\n                print(f\"❌ Smart load failed: {e}\")\n                print(\"⚠️ 警告：正在使用随机初始化，模型可能无法收敛！\")\n        elif pretrained:\n            print(f\"🚨 CRITICAL ERROR: Pretrained=True but file not found at: {pretrained_path}\")\n            \n        # 3. ★★★ 初始冻结 ★★★\n        print(\"❄️ Freezing Backbone parameters...\")\n        for param in self.backbone.parameters():\n            param.requires_grad = False\n\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        for target_name in self.target_names:\n            head = nn.Sequential(\n                nn.Linear(self.n_combined_features, 512),\n                nn.LayerNorm(512), # Swin 推荐用 LayerNorm\n                nn.GELU(),         # Swin 推荐用 GELU\n                nn.Dropout(dropout),\n                nn.Linear(512, 1)\n            )\n            setattr(self, f'head_{target_name.lower().replace(\"_\", \"\")}', head)\n\n    def forward(self, l, r=None):\n        # 显式 no_grad 节省显存\n        with torch.no_grad():\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        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)\n            outputs.append(out)\n        \n        return torch.stack(outputs, dim=1)\n    \n    # 增加解冻方法\n    def unfreeze_backbone(self):\n        print(\"🔥 Unfreezing Backbone...\")\n        for param in self.backbone.parameters():\n            param.requires_grad = True\n        # 移除 forward 里的 no_grad\n        self.forward = self._forward_unfrozen\n\n    def _forward_unfrozen(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        outputs = []\n        for target_name in self.target_names:\n            head = getattr(self, f'head_{target_name.lower().replace(\"_\", \"\")}')\n            outputs.append(head(x).squeeze(1))\n        return torch.stack(outputs, dim=1)\n# =========================================================================================\n# 5. Loss & Metrics\n# =========================================================================================\nclass WeightedMSELoss(nn.Module):\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        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    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        pT, pGDM, pGR = pred\n        tT, tGDM, tGR = target[:,0], target[:,1], target[:,2]\n\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        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        v1 = torch.relu(PG - PT)\n        v2 = torch.relu(PR - PG)\n        loss_cons = (v1 + v2).mean() * self.cons_w\n\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\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}\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    to_np = lambda x: x.detach().cpu().numpy()\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\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    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# =========================================================================================\n# 6. EMA & Utilities\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\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# 7. Training Logic\n# =========================================================================================\ndef train_one_fold(df, fold, lr=None, batch_size=None, wd=None, warmup_epochs=None, dropout=None, aug_strength=None):\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===== TRAINING FOLD {fold} =====\")\n    print(f\"  LR={train_lr:.2e}, BATCH_SIZE={train_batch_size}, WD={train_wd:.4f}\")\n    \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    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, batch_size=train_batch_size, shuffle=True, num_workers=CFG.NUM_WORKERS,\n        pin_memory=True, drop_last=True, worker_init_fn=seed_worker, generator=g, persistent_workers=(CFG.NUM_WORKERS > 0))\n    val_dl = DataLoader(\n        val_ds, batch_size=train_batch_size*2, shuffle=False, num_workers=CFG.NUM_WORKERS,\n        pin_memory=True, worker_init_fn=seed_worker, generator=g, persistent_workers=(CFG.NUM_WORKERS > 0))\n\n    model = BiomassModel(CFG.MODEL_NAME, pretrained=True, pretrained_path=CFG.PRETRAINED_PATH, 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    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    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    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)\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\n            train_loss += loss.item()\n\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)\n                    loss = criterion(pred, y)\n                val_loss += loss.item()\n                y_pred_all.append(pred.detach().cpu())\n                y_true_all.append(y.detach().cpu())\n\n        y_pred_all = torch.cat(y_pred_all, dim=0)\n        y_true_all = torch.cat(y_true_all, dim=0)\n\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        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}: Val R2={wr2:.5f}\")\n\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: 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()\n            print(f\"  -> Best updated. Save {best_path}\")\n\n    return best_metric\n\n# =========================================================================================\n# 8. Inference Utilities\n# =========================================================================================\ndef load_model(model_path, model_name=None, target_names=None, dual_stream=None, dropout=0.3, device=None):\n    if model_name is None: model_name = CFG.MODEL_NAME\n    if target_names is None: target_names = CFG.TARGETS\n    if dual_stream is None: dual_stream = CFG.DUAL_STREAM\n    if device is None: device = CFG.DEVICE\n    \n    model = BiomassModel(model_name=model_name, pretrained=False, target_names=target_names,\n                         dual_stream=dual_stream, dropout=dropout).to(device)\n    checkpoint = torch.load(model_path, map_location=device)\n    model.load_state_dict(checkpoint)\n    model.eval()\n    return model\n\ndef predict_test_set(model, test_df, image_dir, transform, batch_size=32, dual_stream=None, device=None, num_workers=2):\n    if dual_stream is None: dual_stream = CFG.DUAL_STREAM\n    if device is None: device = CFG.DEVICE\n    \n    dataset = TestDataset(test_df, image_dir, transform, dual_stream=dual_stream)\n    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n    all_predictions = []\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            # Single Forward Pass (No TTA)\n            with torch.cuda.amp.autocast(enabled=CFG.USE_AMP):\n                pred = model(left, right)\n            \n            all_predictions.append(pred.cpu().numpy())\n    \n    return np.concatenate(all_predictions, axis=0)\n\ndef ensemble_predict(test_df, model_paths, image_dir, model_name=None, target_names=None, dual_stream=None, dropout=0.3, batch_size=32, device=None, num_workers=2, weights=None):\n    \"\"\"Integrated inference wrapper to satisfy dependency\"\"\"\n    if weights is None: weights = [1.0] * len(model_paths)\n    weights = np.array(weights)\n    weights = weights / weights.sum()\n    transform = get_valid_tf(CFG.IMG_SIZE)\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(model_path, model_name=model_name, target_names=target_names, dual_stream=dual_stream, dropout=dropout, device=device)\n        pred = predict_test_set(model=model, test_df=test_df, image_dir=image_dir, transform=transform, batch_size=batch_size, dual_stream=dual_stream, device=device, num_workers=num_workers)\n        all_predictions.append(pred * weights[i])\n        del model; gc.collect(); torch.cuda.empty_cache()\n    \n    return np.sum(all_predictions, axis=0)\n\ndef predict_and_save_submission(test_csv_path, output_path, model_paths, image_dir=None, model_name=None, target_names=None, dual_stream=None, dropout=0.3, batch_size=None, device=None, num_workers=None, weights=None, use_log1p=None):\n    # --- 补全默认参数 ---\n    if image_dir is None: image_dir = CFG.TEST_IMAGE_DIR\n    if use_log1p is None: use_log1p = CFG.USE_LOG1P\n    if batch_size is None: batch_size = CFG.INFERENCE_BATCH_SIZE\n    if num_workers is None: num_workers = CFG.NUM_WORKERS\n    if target_names is None: target_names = CFG.TARGETS # ★ 关键修复\n    \n    # 1. 准备数据\n    test_df_long = pd.read_csv(test_csv_path)\n    unique_images = test_df_long['image_path'].unique()\n    test_df = pd.DataFrame({'image_path': unique_images})\n    \n    if isinstance(model_paths, str): model_paths = [model_paths]\n    \n    print(f\"Predicting on {len(test_df)} images using {len(model_paths)} model(s) (No TTA)...\")\n    \n    # 2. 运行预测\n    predictions = ensemble_predict(\n        test_df=test_df, model_paths=model_paths, image_dir=image_dir, \n        model_name=model_name, target_names=target_names, dual_stream=dual_stream, dropout=dropout, \n        batch_size=batch_size, device=device, num_workers=num_workers, weights=weights\n    )\n    \n    if use_log1p: predictions = np.expm1(predictions)\n    \n    # 3. 构造结果 DataFrame\n    pred_df = pd.DataFrame(predictions, columns=target_names, index=test_df.index)\n    pred_df['image_path'] = test_df['image_path'].values\n    \n    # 4. 优化后的填表逻辑 (使用 Map 加速，防止 KeyError)\n    submission_rows = []\n    \n    # 构建快速查询字典\n    pred_map = {}\n    for idx, row in pred_df.iterrows():\n        # iloc[0] = Total, iloc[1] = GDM, iloc[2] = Green\n        pred_map[row['image_path']] = (row.iloc[0], row.iloc[1], row.iloc[2])\n\n    for _, row in test_df_long.iterrows():\n        image_path = row['image_path']\n        target_name = row['target_name']\n        \n        # 查表\n        if image_path in pred_map:\n            total, gdm, green = pred_map[image_path]\n        else:\n            total, gdm, green = 0.0, 0.0, 0.0\n        \n        # 逻辑推算\n        if target_name == 'Dry_Total_g': value = total\n        elif target_name == 'GDM_g': value = gdm\n        elif target_name == 'Dry_Green_g': value = green\n        elif target_name == 'Dry_Dead_g': value = total - gdm\n        #三叶草手动-1g这个东西非常少\n        elif target_name == 'Dry_Clover_g': value = gdm - green - 1\n        else: value = 0.0\n        \n        submission_rows.append({'sample_id': row['sample_id'], 'target': max(0.0, value)})\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    print(\"--- Top 5 rows ---\")\n    print(submission_df.head())\n    \n    return submission_df\n\n# =========================================================================================\n# 9. Main Pipeline\n# =========================================================================================\ndef long_to_wide_for_training(df_long, targets=('Dry_Total_g','GDM_g','Dry_Green_g'), strict=True, aggfunc='first'):\n    INDEX_COLS = ['image_path', 'Sampling_Date', 'State', 'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm']\n    wide = df_long.pivot_table(index=INDEX_COLS, columns='target_name', values='target', aggfunc=aggfunc).reset_index()\n    wide['image_id'] = wide['image_path'].apply(lambda x: os.path.splitext(os.path.basename(x))[0])\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    return wide[keep_cols].copy()\n\ndef add_stratified_folds(df, n_folds=5, label_col='Dry_Total_g', bins=5, seed=42):\n    df = df.copy()\n    y = df[label_col].values\n    df['_strat'] = pd.qcut(y, q=bins, labels=False, duplicates='drop')\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    return df.drop(columns=['_strat'])\n\ndef main():\n    if CFG.INFERENCE_MODE:\n        # Only look for Fold 0\n        path = os.path.join(CFG.INFERENCE_MODEL_DIR, 'best_model_fold0.pth')\n        \n        if not os.path.exists(path):\n            print(f\"Error: Model not found at {path}. Did you train Fold 0?\")\n            return\n\n        predict_and_save_submission(CFG.TEST_CSV, CFG.SUBMISSION_OUTPUT, [path])\n    \n    else:\n        print(\"Training Mode: ONLY training FOLD 0 (Single Model)\")\n        df = pd.read_csv(CFG.TRAIN_CSV)\n        df = long_to_wide_for_training(df, targets=CFG.TARGETS)\n        df = add_stratified_folds(df, n_folds=CFG.N_FOLDS)\n        \n        # 只训练第 0 折\n        train_one_fold(df, 0)\n\nif __name__ == '__main__':\n    main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-12T09:45:47.744670Z","iopub.execute_input":"2025-12-12T09:45:47.745324Z","iopub.status.idle":"2025-12-12T09:46:42.433325Z","shell.execute_reply.started":"2025-12-12T09:45:47.745296Z","shell.execute_reply":"2025-12-12T09:46:42.432661Z"}},"outputs":[],"execution_count":null}]}