{"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":706618,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":536503,"modelId":494969}],"dockerImageVersionId":31240,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\n基于DINOv2的生物质预测模型\n支持完整图像输入和可选的左右分割\n\"\"\"\n\nimport os\nimport math\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom timm.utils import ModelEmaV2\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom tqdm import tqdm\nimport cv2\nimport gc\n\n\n# ================== 配置 ==================\nclass CFG:\n    # 路径配置\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    MODEL_DIR = '/kaggle/working/models'\n    EXP_ID = 'freeze20_splitLR_augTrue'\n    # 模型配置\n    MODEL_NAME = 'vit_large_patch14_dinov2.lvd142m'  # DINOv2 Large\n    # MODEL_NAME = 'vit_giant_patch14_dinov2.lvd142m'  # 如果需要更大模型\n    PRETRAINED = True\n    IMG_SIZE = 518  # DINOv2标准尺寸 (可以是14的倍数: 518, 392等)\n    \n    # 图像分割配置\n    SPLIT_LEFT_RIGHT = True  # 是否分割左右图像\n    USE_CROP_BOTTOM = False   # 是否裁剪底部\n    USE_AUGMENTATION = True  # 是否使用数据增强\n    \n    # 训练配置\n    SEED = 82947501\n    N_FOLDS = 5\n    FOLDS_TO_TRAIN = [0]  # 测试时只训练fold 0\n    EPOCHS = 20\n    FREEZE_EPOCHS = 20\n    WARMUP_EPOCHS = 3\n    BATCH_SIZE = 4\n    GRAD_ACC = 4\n    NUM_WORKERS = 4\n    \n    # 优化器配置（与原代码一致）\n    LR_BACKBONE = 5e-4\n    LR_REST = 1e-3\n    WD = 1e-2\n    EMA_DECAY = 0.9\n    PATIENCE = 5\n    \n    # 目标列配置\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    \n    # 损失和评估权重（与原代码一致）\n    LOSS_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.0, 0.0])  # 只监督Green, Dead, Clover\n    R2_WEIGHTS = np.array([0.1, 0.1, 0.1, 0.2, 0.5])    # 评估所有5个目标\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n\n# ================== 数据集 ==================\ndef get_transforms(is_train=False):\n    \"\"\"获取数据变换\"\"\"\n    transforms_list = [A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE)]\n    \n    if is_train and CFG.USE_AUGMENTATION:\n        transforms_list.extend([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n        ])\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\nclass BiomassDataset(Dataset):\n    \"\"\"生物质数据集\"\"\"\n    def __init__(self, df, transform, img_dir, split_lr=False):\n        self.df = df\n        self.transform = transform\n        self.img_dir = img_dir\n        self.split_lr = split_lr\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        \n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        \n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # 可选：裁剪底部\n        if CFG.USE_CROP_BOTTOM:\n            h, w = img.shape[:2]\n            img = img[0:int(h*0.90), :]\n        \n        label = torch.from_numpy(self.labels[idx])\n        \n        if self.split_lr:\n            # 分割左右图像\n            h, w, _ = img.shape\n            mid = w // 2\n            left = img[:, :mid]\n            right = img[:, mid:]\n            \n            left = self.transform(image=left)['image']\n            right = self.transform(image=right)['image']\n            return left, right, label\n        else:\n            # 完整图像\n            # 需要调整图像尺寸以匹配模型输入\n            img = self.transform(image=img)['image']\n            return img, label\n\n\n# ================== 模型 ==================\nclass BiomassModel(nn.Module):\n    \"\"\"生物质预测模型\"\"\"\n    def __init__(self, model_name, pretrained=True, split_lr=False):\n        super().__init__()\n        self.model_name = model_name\n        self.split_lr = split_lr\n        \n        # 加载backbone（保留patch tokens）\n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=False, \n            num_classes=0, \n            global_pool=''  # 保留所有tokens\n        )\n            # 2. 手动加载Input里的DINOv2权重\n        if pretrained:\n            # 你的权重路径（从截图里的INPUT结构来的）\n            dinov2_weight_path = '/kaggle/input/timm-models/pytorch/vit_large_patch14_dinov2.lvd142m/1/pytorch_model.bin'\n\n            # 加载权重\n            state_dict = torch.load(dinov2_weight_path, map_location=CFG.DEVICE)\n            \n            # 适配HuggingFace格式的权重（因为你的权重是从HF下载的）\n            new_state_dict = {}\n            for k, v in state_dict.items():\n                # HuggingFace的DINOv2权重键名是\"encoder.xxx\"，timm需要去掉\"encoder.\"前缀\n                if k.startswith('encoder.'):\n                    k = k.replace('encoder.', '')  # 例如 \"encoder.blocks.0\" → \"blocks.0\"\n                new_state_dict[k] = v\n            \n            # 加载到backbone\n            self.backbone.load_state_dict(new_state_dict, strict=False)\n            print(f\"✅ 成功加载DINOv2权重：{dinov2_weight_path}\")\n        # 启用梯度检查点以节省显存\n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(True)\n        \n        nf = self.backbone.num_features\n        \n        # 如果分割左右图像，需要融合特征\n        if split_lr:\n            # 简单的concatenation后处理\n            self.fusion = nn.Sequential(\n                nn.LayerNorm(nf * 2),\n                nn.Linear(nf * 2, nf),\n                nn.GELU(),\n                nn.Dropout(0.1)\n            )\n            head_input_dim = nf\n        else:\n            self.fusion = None\n            head_input_dim = nf\n        \n        # 全局池化\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        \n        # 三个预测头（使用Softplus确保非负）\n        self.head_green = self._make_head(head_input_dim)\n        self.head_clover = self._make_head(head_input_dim)\n        self.head_dead = self._make_head(head_input_dim)\n    \n    def _make_head(self, input_dim):\n        \"\"\"创建预测头\"\"\"\n        return nn.Sequential(\n            nn.Linear(input_dim, input_dim // 2),\n            nn.GELU(),\n            nn.Dropout(0.2),\n            nn.Linear(input_dim // 2, 1),\n            nn.Softplus()  # 确保输出非负\n        )\n    \n    def forward(self, *inputs):\n        \"\"\"\n        前向传播\n        如果split_lr=True: inputs = (left, right)\n        否则: inputs = (img,)\n        \"\"\"\n        if self.split_lr:\n            left, right = inputs\n            # 提取特征\n            x_l = self.backbone(left)  # (B, N, D)\n            x_r = self.backbone(right)  # (B, N, D)\n            \n            # 池化\n            x_l = self.pool(x_l.transpose(1, 2)).flatten(1)  # (B, D)\n            x_r = self.pool(x_r.transpose(1, 2)).flatten(1)  # (B, D)\n            \n            # 融合\n            x = torch.cat([x_l, x_r], dim=1)  # (B, 2D)\n            x = self.fusion(x)  # (B, D)\n        else:\n            img = inputs[0]\n            # 提取特征\n            x = self.backbone(img)  # (B, N, D)\n            # 池化\n            x = self.pool(x.transpose(1, 2)).flatten(1)  # (B, D)\n        \n        # 预测三个基础值\n        green = self.head_green(x)\n        clover = self.head_clover(x)\n        dead = self.head_dead(x)\n        \n        # 根据约束关系推导其他值\n        gdm = green + clover\n        total = gdm + dead\n        \n        return total, gdm, green, clover, dead\n\n\n# ================== 损失函数 ==================\ndef biomass_loss(outputs, labels, weights=None):\n    \"\"\"\n    生物质损失函数（Huber Loss）\n    Args:\n        outputs: (total, gdm, green, clover, dead)\n        labels: (B, 5) tensor with [green, dead, clover, gdm, total]\n        weights: 损失权重\n    \"\"\"\n    total, gdm, green, clover, dead = outputs\n    \n    # Huber loss (Smooth L1)\n    huber = nn.SmoothL1Loss(beta=5.0)\n    \n    # 计算每个目标的损失\n    l_green = huber(green.squeeze(), labels[:, 0])\n    l_dead = huber(dead.squeeze(), labels[:, 1])\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    # 堆叠损失\n    losses = torch.stack([l_green, l_dead, l_clover, l_gdm, l_total])\n    \n    # 应用权重\n    if weights is None:\n        return losses.mean()\n    \n    weights = torch.as_tensor(weights, device=losses.device, dtype=losses.dtype)\n    weights = weights / weights.sum()\n    return (losses * weights).sum()\n\n\n# ================== 评估指标 ==================\ndef weighted_r2_score_global(y_true, y_pred, weights):\n    \"\"\"\n    计算全局加权R²（图2中的公式）\n    Args:\n        y_true: (N, 5) numpy array\n        y_pred: (N, 5) numpy array\n        weights: (5,) numpy array\n    Returns:\n        global_r2: 全局加权R²\n        avg_r2: 平均R²\n        per_r2: 每个目标的R²\n    \"\"\"\n    # 展平\n    flat_true = y_true.reshape(-1)\n    flat_pred = y_pred.reshape(-1)\n    w = np.tile(weights, y_true.shape[0])\n    \n    # 计算加权均值\n    mean_w = np.sum(w * flat_true) / np.sum(w)\n    \n    # 计算SS_res和SS_tot\n    ss_res = np.sum(w * (flat_true - flat_pred) ** 2)\n    ss_tot = np.sum(w * (flat_true - mean_w) ** 2)\n    \n    # 全局R²\n    global_r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n    \n    # 每个目标的R²\n    per_r2 = []\n    for i in range(y_true.shape[1]):\n        yt = y_true[:, i]\n        yp = y_pred[:, i]\n        ss_res_i = np.sum((yt - yp) ** 2)\n        ss_tot_i = np.sum((yt - np.mean(yt)) ** 2)\n        r2_i = 1 - ss_res_i / ss_tot_i if ss_tot_i > 0 else 0.0\n        per_r2.append(r2_i)\n    \n    per_r2 = np.array(per_r2)\n    avg_r2 = np.sum(per_r2 * weights) / np.sum(weights)\n    \n    return global_r2, avg_r2, per_r2\n\n\n# ================== 训练和验证 ==================\nscaler = torch.amp.GradScaler('cuda', enabled=torch.cuda.is_available())\n\n\ndef train_epoch(model, loader, optimizer, scheduler, device, ema=None, split_lr=False):\n    \"\"\"训练一个epoch\"\"\"\n    model.train()\n    running_loss = 0.0\n    optimizer.zero_grad()\n    \n    amp_ctx = torch.amp.autocast(device_type='cuda') if torch.cuda.is_available() else torch.cuda.amp.autocast(enabled=False)\n    \n    for i, batch in enumerate(tqdm(loader, desc='Train', leave=False)):\n        if split_lr:\n            left, right, labels = batch\n            left = left.to(device, non_blocking=True)\n            right = right.to(device, non_blocking=True)\n            labels = labels.to(device, non_blocking=True)\n            \n            with amp_ctx:\n                outputs = model(left, right)\n                loss = biomass_loss(outputs, labels, CFG.LOSS_WEIGHTS) / CFG.GRAD_ACC\n        else:\n            img, labels = batch\n            img = img.to(device, non_blocking=True)\n            labels = labels.to(device, non_blocking=True)\n            \n            with amp_ctx:\n                outputs = model(img)\n                loss = biomass_loss(outputs, labels, CFG.LOSS_WEIGHTS) / CFG.GRAD_ACC\n        \n        scaler.scale(loss).backward()\n        running_loss += loss.item() * labels.size(0) * CFG.GRAD_ACC\n        \n        if (i + 1) % CFG.GRAD_ACC == 0 or (i + 1) == len(loader):\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            scaler.step(optimizer)\n            scaler.update()\n            if ema is not None:\n                ema.update(model.module if hasattr(model, 'module') else model)\n            optimizer.zero_grad()\n    \n    scheduler.step()\n    return running_loss / len(loader.dataset)\n\n\n@torch.no_grad()\ndef valid_epoch(model, loader, device, split_lr=False):\n    \"\"\"验证一个epoch\"\"\"\n    model.eval()\n    running_loss = 0.0\n    all_preds = []\n    all_labels = []\n    \n    amp_ctx = torch.amp.autocast(device_type='cuda') if torch.cuda.is_available() else torch.cuda.amp.autocast(enabled=False)\n    \n    for batch in tqdm(loader, desc='Valid', leave=False):\n        if split_lr:\n            left, right, labels = batch\n            left = left.to(device, non_blocking=True)\n            right = right.to(device, non_blocking=True)\n            labels = labels.to(device, non_blocking=True)\n            \n            with amp_ctx:\n                total, gdm, green, clover, dead = model(left, right)\n                loss = biomass_loss((total, gdm, green, clover, dead), labels, CFG.LOSS_WEIGHTS)\n        else:\n            img, labels = batch\n            img = img.to(device, non_blocking=True)\n            labels = labels.to(device, non_blocking=True)\n            \n            with amp_ctx:\n                total, gdm, green, clover, dead = model(img)\n                loss = biomass_loss((total, gdm, green, clover, dead), labels, CFG.LOSS_WEIGHTS)\n        \n        running_loss += loss.item() * labels.size(0)\n        \n        # 收集预测和标签\n        preds = np.stack([\n            green.cpu().numpy().ravel(),\n            dead.cpu().numpy().ravel(),\n            clover.cpu().numpy().ravel(),\n            gdm.cpu().numpy().ravel(),\n            total.cpu().numpy().ravel()\n        ], axis=1)\n        \n        all_preds.append(preds)\n        all_labels.append(labels.cpu().numpy())\n    \n    # 计算R²\n    all_preds = np.concatenate(all_preds, axis=0)\n    all_labels = np.concatenate(all_labels, axis=0)\n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(all_labels, all_preds, CFG.R2_WEIGHTS)\n    \n    return running_loss / len(loader.dataset), global_r2, avg_r2, per_r2\n\n\n# ================== 优化器和调度器 ==================\ndef build_optimizer(model):\n    \"\"\"构建优化器\"\"\"\n    backbone_ids = {id(p) for p in model.backbone.parameters()}\n    \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    \"\"\"构建学习率调度器（余弦退火+warmup）\"\"\"\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\ndef set_backbone_requires_grad(model, requires_grad):\n    \"\"\"设置backbone的梯度\"\"\"\n    for p in model.backbone.parameters():\n        p.requires_grad = requires_grad\n\n\n# ================== 主训练循环 ==================\ndef train_fold(fold, tr_df, val_df):\n    \"\"\"训练单个fold\"\"\"\n    print(f'\\n{\"=\"*70}')\n    print(f'Training Fold {fold}')\n    print(f'{\"=\"*70}')\n    \n    # 创建数据集\n    tr_set = BiomassDataset(tr_df, get_transforms(is_train=True), \n                           CFG.TRAIN_IMAGE_DIR, split_lr=CFG.SPLIT_LEFT_RIGHT)\n    val_set = BiomassDataset(val_df, get_transforms(is_train=False), \n                            CFG.TRAIN_IMAGE_DIR, split_lr=CFG.SPLIT_LEFT_RIGHT)\n    \n    tr_loader = DataLoader(tr_set, batch_size=CFG.BATCH_SIZE, shuffle=True,\n                          num_workers=CFG.NUM_WORKERS, pin_memory=True, drop_last=True)\n    val_loader = DataLoader(val_set, batch_size=CFG.BATCH_SIZE, shuffle=False,\n                           num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    \n    # 创建模型\n    base_model = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED, \n                             split_lr=CFG.SPLIT_LEFT_RIGHT).to(CFG.DEVICE)\n    model = nn.DataParallel(base_model)\n    \n    # 冻结backbone\n    set_backbone_requires_grad(base_model, False)\n    \n    # 优化器和调度器\n    optimizer = build_optimizer(base_model)\n    scheduler = build_scheduler(optimizer)\n    ema = ModelEmaV2(base_model, decay=CFG.EMA_DECAY)\n    \n    # 训练\n    best_global_r2 = -np.inf\n    patience_counter = 0\n    \n    for epoch in range(1, CFG.EPOCHS + 1):\n        # 解冻backbone\n        if epoch == CFG.FREEZE_EPOCHS + 1:\n            set_backbone_requires_grad(base_model, True)\n            print(f'Epoch {epoch}: Backbone unfrozen')\n        \n        # 训练\n        tr_loss = train_epoch(model, tr_loader, optimizer, scheduler, \n                            CFG.DEVICE, ema, split_lr=CFG.SPLIT_LEFT_RIGHT)\n        \n        # 验证（使用EMA模型）\n        eval_model = ema.module if ema is not None else base_model\n        val_loss, global_r2, avg_r2, per_r2 = valid_epoch(\n            eval_model, val_loader, CFG.DEVICE, split_lr=CFG.SPLIT_LEFT_RIGHT\n        )\n        \n        # 打印结果\n        per_r2_str = ' | '.join([f'{CFG.ALL_TARGET_COLS[i][:5]}: {r2:.3f}' \n                                 for i, r2 in enumerate(per_r2)])\n        print(f'Epoch {epoch:02d} | TrLoss: {tr_loss:.5f} | ValLoss: {val_loss:.5f}')\n        print(f'  GlobalR²: {global_r2:.4f} | AvgR²: {avg_r2:.4f} {\"[BEST]\" if global_r2 > best_global_r2 else \"\"}')\n        print(f'  {per_r2_str}')\n        \n        # 保存最佳模型\n        if global_r2 > best_global_r2:\n            best_global_r2 = global_r2\n            patience_counter = 0\n            \n            os.makedirs(CFG.MODEL_DIR, exist_ok=True)\n            save_path = os.path.join(\n                CFG.MODEL_DIR, \n                f'best_model_{CFG.EXP_ID}_fold{fold}.pth'  # 核心：嵌入EXP_ID\n            )\n            torch.save(eval_model.state_dict(), save_path)\n            print(f'  → Model saved to {save_path}')\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE:\n                print(f'  → Early stopping (patience={CFG.PATIENCE})')\n                break\n        \n        torch.cuda.empty_cache()\n        gc.collect()\n    \n    # 清理\n    del model, base_model, optimizer, scheduler, ema\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return best_global_r2\n\n\n# ================== 主程序 ==================\ndef main():\n    \"\"\"主程序\"\"\"\n    print(f'Device: {CFG.DEVICE}')\n    print(f'Model: {CFG.MODEL_NAME}')\n    print(f'Image size: {CFG.IMG_SIZE}')\n    print(f'Split left-right: {CFG.SPLIT_LEFT_RIGHT}')\n    print(f'Use augmentation: {CFG.USE_AUGMENTATION}')\n    print(f'Use crop bottom: {CFG.USE_CROP_BOTTOM}')\n    \n    # 加载数据\n    print('\\nLoading data...')\n    df_long = pd.read_csv(CFG.TRAIN_CSV)\n    df_wide = df_long.pivot(index='image_path', columns='target_name', \n                            values='target').reset_index()\n    \n    # 合并元数据\n    if 'Sampling_Date' in df_long.columns and 'State' in df_long.columns:\n        meta_df = df_long[['image_path', 'Sampling_Date', 'State']].drop_duplicates()\n        df_wide = df_wide.merge(meta_df, on='image_path', how='left')\n    \n    df_wide = df_wide[['image_path', 'Sampling_Date', 'State'] + CFG.ALL_TARGET_COLS]\n    print(f'Total images: {len(df_wide)}')\n    \n    # 5-Fold交叉验证\n    sgkf = StratifiedGroupKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\n    groups = df_wide['Sampling_Date']\n    y_stratify = df_wide['State']\n    \n    fold_results = []\n    \n    for fold, (tr_idx, val_idx) in enumerate(sgkf.split(df_wide, y_stratify, groups=groups)):\n        if fold not in CFG.FOLDS_TO_TRAIN:\n            print(f'\\nSkipping fold {fold}')\n            continue\n        \n        tr_df = df_wide.iloc[tr_idx].reset_index(drop=True)\n        val_df = df_wide.iloc[val_idx].reset_index(drop=True)\n        \n        best_r2 = train_fold(fold, tr_df, val_df)\n        fold_results.append({'fold': fold, 'global_r2': best_r2})\n    \n    # 打印总结\n    print(f'\\n{\"=\"*70}')\n    print('Training Summary')\n    print(f'{\"=\"*70}')\n    for result in fold_results:\n        print(f\"Fold {result['fold']}: Global R² = {result['global_r2']:.4f}\")\n\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\n基于DINOv2的生物质预测模型 - 推理代码（修复维度错误版）\n核心：100%适配训练代码 + 修复OpenCV resize维度错误\n\"\"\"\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom tqdm import tqdm\nimport cv2\nimport gc\n\n\n# ================== 配置（完全复用训练代码）==================\nclass CFG:\n    # 路径配置（和训练一致）\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    TEST_CSV = os.path.join(BASE_PATH, 'test.csv')\n    TEST_IMAGE_DIR = os.path.join(BASE_PATH, 'test')\n    MODEL_DIR = '/kaggle/working/models'\n    EXP_ID = 'freeze20_splitLR_augTrue'  # 和训练代码的EXP_ID完全一致\n    \n    # 模型配置（和训练一致）\n    MODEL_NAME = 'vit_large_patch14_dinov2.lvd142m'\n    PRETRAINED = False  # 推理时不需要预训练权重\n    IMG_SIZE = 518\n    \n    # 图像分割配置（和训练完全一致）\n    SPLIT_LEFT_RIGHT = True  # 必须和训练一致\n    USE_CROP_BOTTOM = False   # 必须和训练一致\n    USE_AUGMENTATION = False  # 推理时关闭增强\n    \n    # 推理配置\n    TTA_STEPS = 1  # TTA步数（1=无TTA）\n    FOLDS_TO_TRAIN = [0]  # 要推理的fold列表（和训练一致）\n    N_FOLDS = 5\n    BATCH_SIZE = 1  # TTA推理强制batch_size=1\n    NUM_WORKERS = 4\n    \n    # 目标列配置（完全复用训练代码）\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    \n    # 控制开关\n    CREATE_SUBMISSION = True\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n\n# ================== 1. 数据变换（完全复用训练代码）==================\ndef get_transforms(is_train=False):\n    \"\"\"和训练代码完全一致的变换逻辑\"\"\"\n    transforms_list = [A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE)]\n    \n    if is_train and CFG.USE_AUGMENTATION:\n        transforms_list.extend([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n        ])\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\n# ================== 2. TTA变换（适配训练代码的transforms）==================\ndef get_tta_transforms(num_transforms):\n    \"\"\"TTA变换：基于训练代码的transforms扩展\"\"\"\n    normalize = A.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225]\n    )\n    to_tensor = ToTensorV2()\n    \n    all_tta_transforms = [\n        # View 1: Original（和训练代码的验证变换一致）\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            normalize,\n            to_tensor\n        ]),\n        # View 2: Horizontal Flip\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.HorizontalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n        # View 3: Vertical Flip\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.VerticalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n        # View 4: Both Flips\n        A.Compose([\n            A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n            A.HorizontalFlip(p=1.0),\n            A.VerticalFlip(p=1.0),\n            normalize,\n            to_tensor\n        ]),\n    ]\n    tta_transforms = all_tta_transforms[:num_transforms]\n    return tta_transforms\n\nprint(f\"✓ TTA transforms defined ({CFG.TTA_STEPS} views)\")\n\n\n# ================== 3. 测试数据集（完全适配训练代码的Dataset）==================\nclass BiomassTestDataset(Dataset):\n    \"\"\"测试数据集：100%对齐训练代码的图像处理逻辑\"\"\"\n    def __init__(self, img_dir, split_lr=False):\n        self.img_dir = img_dir\n        self.split_lr = split_lr  # 和训练代码的SPLIT_LEFT_RIGHT一致\n        # 读取测试文件夹中的图片（和训练代码的路径读取逻辑一致）\n        self.paths = sorted([\n            os.path.join(img_dir, f) for f in os.listdir(img_dir)\n            if f.lower().endswith(('.jpg', '.jpeg', '.png'))\n        ])\n        self.filenames = [os.path.basename(p) for p in self.paths]\n    \n    def __len__(self):\n        return len(self.paths)\n    \n    def __getitem__(self, idx):\n        # 图像读取逻辑（和训练代码完全一致）\n        path = self.paths[idx]\n        img = cv2.imread(path)\n        \n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        \n        # BGR转RGB（和训练代码一致）\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # 裁剪底部（和训练代码的USE_CROP_BOTTOM一致）\n        if CFG.USE_CROP_BOTTOM:\n            h, w = img.shape[:2]\n            img = img[0:int(h*0.90), :]\n        \n        # 分割左右图像（和训练代码的SPLIT_LEFT_RIGHT一致）\n        if self.split_lr:\n            h, w, _ = img.shape\n            mid = w // 2\n            left = img[:, :mid].copy()\n            right = img[:, mid:].copy()\n            \n            # 返回原始numpy数组（uint8格式，避免维度问题）\n            return left, right, self.filenames[idx]\n        else:\n            return img, self.filenames[idx]\n\nprint(\"✓ Test dataset class defined\")\n\n\n# ================== 4. 模型结构（完全复用训练代码）==================\nclass BiomassModel(nn.Module):\n    \"\"\"和训练代码完全一致的模型结构\"\"\"\n    def __init__(self, model_name, pretrained=True, split_lr=False):\n        super().__init__()\n        self.model_name = model_name\n        self.split_lr = split_lr\n        \n        # 加载backbone（和训练代码一致）\n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=pretrained, \n            num_classes=0, \n            global_pool=''\n        )\n        \n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(True)\n        \n        nf = self.backbone.num_features\n        \n        # 左右图融合（和训练代码一致）\n        if split_lr:\n            self.fusion = nn.Sequential(\n                nn.LayerNorm(nf * 2),\n                nn.Linear(nf * 2, nf),\n                nn.GELU(),\n                nn.Dropout(0.1)\n            )\n            head_input_dim = nf\n        else:\n            self.fusion = None\n            head_input_dim = nf\n        \n        # 全局池化（和训练代码一致）\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        \n        # 预测头（和训练代码一致）\n        self.head_green = self._make_head(head_input_dim)\n        self.head_clover = self._make_head(head_input_dim)\n        self.head_dead = self._make_head(head_input_dim)\n    \n    def _make_head(self, input_dim):\n        \"\"\"和训练代码一致的预测头\"\"\"\n        return nn.Sequential(\n            nn.Linear(input_dim, input_dim // 2),\n            nn.GELU(),\n            nn.Dropout(0.2),\n            nn.Linear(input_dim // 2, 1),\n            nn.Softplus()\n        )\n    \n    def forward(self, *inputs):\n        \"\"\"和训练代码完全一致的前向传播\"\"\"\n        if self.split_lr:\n            left, right = inputs\n            x_l = self.backbone(left)\n            x_r = self.backbone(right)\n            \n            x_l = self.pool(x_l.transpose(1, 2)).flatten(1)\n            x_r = self.pool(x_r.transpose(1, 2)).flatten(1)\n            \n            x = torch.cat([x_l, x_r], dim=1)\n            x = self.fusion(x)\n        else:\n            img = inputs[0]\n            x = self.backbone(img)\n            x = self.pool(x.transpose(1, 2)).flatten(1)\n        \n        green = self.head_green(x)\n        clover = self.head_clover(x)\n        dead = self.head_dead(x)\n        \n        gdm = green + clover\n        total = gdm + dead\n        \n        return total, gdm, green, clover, dead\n\n\n# ================== 5. TTA推理（修复维度错误）==================\n@torch.no_grad()\ndef predict_with_tta(model, left_np, right_np, tta_transforms):\n    \"\"\"单模型TTA推理：修复维度错误 + 适配训练代码\"\"\"\n    all_tta_preds = []\n    \n    # 确保输入是uint8格式的3维数组（H, W, C）\n    left_np = left_np.astype(np.uint8) if left_np.dtype != np.uint8 else left_np\n    right_np = right_np.astype(np.uint8) if right_np.dtype != np.uint8 else right_np\n\n    for tfm in tta_transforms:\n        # 应用TTA变换（输入是纯3维数组，无batch维度）\n        left_tensor = tfm(image=left_np)['image'].unsqueeze(0).to(CFG.DEVICE)\n        right_tensor = tfm(image=right_np)['image'].unsqueeze(0).to(CFG.DEVICE)\n        \n        # 模型前向传播（和训练代码一致）\n        total, gdm, green, clover, dead = model(left_tensor, right_tensor)\n        \n        # 提取和训练代码一致的3个基础目标\n        p_total = total.cpu().item()\n        p_gdm = gdm.cpu().item()\n        p_green = green.cpu().item()\n        \n        all_tta_preds.append([p_total, p_gdm, p_green])\n    \n    # TTA结果平均\n    final_pred = np.mean(all_tta_preds, axis=0)\n    return final_pred\n\n\n# ================== 6. 主推理函数（修复维度错误）==================\ndef run_inference():\n    \"\"\"主推理函数：适配训练代码 + 修复维度错误\"\"\"\n    print(\"\\n\" + \"=\"*70)\n    print(\"STARTING INFERENCE\")\n    print(\"=\"*70)\n    print(f\"EXP_ID: {CFG.EXP_ID}\")\n    print(f\"Split left-right: {CFG.SPLIT_LEFT_RIGHT}\")\n    print(f\"Device: {CFG.DEVICE}\")\n    \n    # 创建测试数据集（和训练代码的split_lr一致）\n    dataset = BiomassTestDataset(CFG.TEST_IMAGE_DIR, split_lr=CFG.SPLIT_LEFT_RIGHT)\n    loader = DataLoader(\n        dataset,\n        batch_size=1,\n        shuffle=False,\n        num_workers=CFG.NUM_WORKERS,\n        pin_memory=True\n    )\n    \n    # 获取TTA变换\n    tta_transforms = get_tta_transforms(CFG.TTA_STEPS)\n    \n    # 初始化预测结果累加器（和训练代码的目标维度一致）\n    accumulated_preds = np.zeros((len(dataset), 3), dtype=np.float32)\n    \n    # 要推理的fold列表（和训练代码一致）\n    folds_to_use = getattr(CFG, 'FOLDS_TO_TRAIN', list(range(CFG.N_FOLDS)))\n    if not folds_to_use:\n        folds_to_use = list(range(CFG.N_FOLDS))\n\n    print(f\"Folds requested for inference: {folds_to_use}\")\n    filenames = dataset.filenames.copy()\n    successful_folds = 0\n    \n    # 遍历指定的fold（适配训练代码的模型保存规则）\n    for fold in folds_to_use:\n        print(f\"\\nProcessing Fold {fold}...\")\n        \n        # 创建模型（和训练代码一致）\n        model = BiomassModel(\n            CFG.MODEL_NAME,\n            pretrained=CFG.PRETRAINED,\n            split_lr=CFG.SPLIT_LEFT_RIGHT\n        ).to(CFG.DEVICE)\n        \n        # 加载训练好的模型权重（适配训练代码的EXP_ID命名规则）\n        weight_path = os.path.join(\n            CFG.MODEL_DIR, \n            f'best_model_{CFG.EXP_ID}_fold{fold}.pth'  # 关键：和训练代码的保存路径一致\n        )\n        \n        if not os.path.exists(weight_path):\n            print(f\"Warning: Model file {weight_path} not found! Skipping fold {fold}.\")\n            del model\n            torch.cuda.empty_cache(); gc.collect()\n            continue\n        \n        # 加载权重（和训练代码一致）\n        state = torch.load(weight_path, map_location='cpu')\n        if isinstance(state, dict) and ('model_state_dict' in state or 'state_dict' in state):\n             key = 'model_state_dict' if 'model_state_dict' in state else 'state_dict'\n             sd = state[key]\n        else:\n             sd = state\n        \n        model.load_state_dict(sd)\n        model.eval()\n        \n        # 单fold推理（修复维度错误）\n        for i, batch in enumerate(tqdm(loader, desc=f\"Fold {fold}\")):\n            # 正确解析batch：分离left/right/filename，去除batch维度\n            left, right, filename = batch\n            # 关键修复：从tensor转换为numpy，并挤压batch维度（从[1,H,W,C]→[H,W,C]）\n            left_np = left.squeeze(0).cpu().numpy()  # 挤压batch维度\n            right_np = right.squeeze(0).cpu().numpy()\n            \n            # TTA推理\n            pred = predict_with_tta(model, left_np, right_np, tta_transforms)\n            accumulated_preds[i] += pred\n        \n        successful_folds += 1\n        del model\n        torch.cuda.empty_cache(); gc.collect()\n    \n    if successful_folds == 0:\n        raise FileNotFoundError(f\"No model weights found for requested folds: {folds_to_use}\")\n    \n    # 多fold结果平均\n    final_predictions = accumulated_preds / successful_folds\n    \n    print(f\"\\nInference complete. Successfully used {successful_folds} fold(s)\")\n    return final_predictions, filenames\n\n\n# ================== 7. 后处理（完全复用训练代码的目标推导逻辑）==================\ndef postprocess_predictions(preds_direct):\n    \"\"\"后处理：和训练代码一致的派生目标推导逻辑\"\"\"\n    print(\"\\nPost-processing predictions...\")\n    \n    # 提取训练代码定义的3个基础目标\n    pred_total = preds_direct[:, 0]\n    pred_gdm = preds_direct[:, 1]\n    pred_green = preds_direct[:, 2]\n    \n    # 推导派生目标（和训练代码一致，添加非负约束）\n    pred_clover = np.maximum(0, pred_gdm - pred_green)\n    pred_dead = np.maximum(0, pred_total - pred_gdm)\n    \n    # 按训练代码的ALL_TARGET_COLS顺序堆叠\n    preds_all = np.stack([\n        pred_green,    # Dry_Green_g\n        pred_dead,     # Dry_Dead_g\n        pred_clover,   # Dry_Clover_g\n        pred_gdm,      # GDM_g\n        pred_total     # Dry_Total_g\n    ], axis=1)\n    \n    # 打印统计信息（和训练代码一致）\n    print(f\"✓ Post-processing complete\")\n    print(f\"  Output shape: {preds_all.shape}\")\n    print(f\"\\nPrediction statistics:\")\n    for i, col in enumerate(CFG.ALL_TARGET_COLS):\n        print(f\"  {col:15s}: mean={preds_all[:, i].mean():.2f}, \"\n              f\"std={preds_all[:, i].std():.2f}, \"\n              f\"min={preds_all[:, i].min():.2f}, \"\n              f\"max={preds_all[:, i].max():.2f}\")\n    \n    return preds_all\n\n\n# ================== 8. 生成提交文件（适配训练代码的目标列）==================\ndef create_submission(predictions, filenames):\n    \"\"\"生成提交文件：适配训练代码的目标列顺序\"\"\"\n    print(\"\\n\" + \"=\"*70)\n    print(\"CREATING SUBMISSION FILE\")\n    print(\"=\"*70)\n    \n    # 加载test.csv（和训练代码的路径逻辑一致）\n    test_df = pd.read_csv(CFG.TEST_CSV)\n    print(f\"\\nTest CSV loaded: {len(test_df)} rows\")\n    print(f\"Sample image_path from test.csv: {test_df['image_path'].iloc[0]}\")\n    print(f\"Sample filename from predictions: {filenames[0]}\")\n    \n    # 修复路径格式（适配训练代码的路径处理）\n    test_path_example = test_df['image_path'].iloc[0]\n    if '/' in test_path_example:\n        prefix = test_path_example.rsplit('/', 1)[0] + '/'\n        corrected_filenames = [prefix + fn for fn in filenames]\n        print(f\"Corrected path format: {corrected_filenames[0]}\")\n    else:\n        corrected_filenames = filenames\n    \n    # 创建宽格式DataFrame（和训练代码的ALL_TARGET_COLS一致）\n    preds_wide = pd.DataFrame(predictions, columns=CFG.ALL_TARGET_COLS)\n    preds_wide.insert(0, 'image_path', corrected_filenames)\n    \n    print(f\"\\nWide format predictions:\")\n    print(preds_wide.head())\n    \n    # 转换为长格式（比赛要求）\n    preds_long = preds_wide.melt(\n        id_vars=['image_path'],\n        value_vars=CFG.ALL_TARGET_COLS,\n        var_name='target_name',\n        value_name='target'\n    )\n    \n    # 匹配sample_id\n    submission = pd.merge(\n        test_df[['sample_id', 'image_path', 'target_name']],\n        preds_long,\n        on=['image_path', 'target_name'],\n        how='left'\n    )\n    \n    # 处理缺失值+排序\n    submission['target'] = submission['target'].fillna(0.0)\n    submission = submission[['sample_id', 'target']].sort_values('sample_id').reset_index(drop=True)\n    \n    # 保存提交文件\n    output_path = os.path.join(CFG.MODEL_DIR, f'submission_{CFG.EXP_ID}.csv')\n    submission.to_csv(\"/kaggle/working/submission.csv\", index=False)\n    \n    # 验证（和训练代码的日志风格一致）\n    print(f\"\\n✓ Submission file saved: {output_path}\")\n    print(f\"  Total rows: {len(submission)}\")\n    print(f\"  Expected rows: {len(test_df)}\")\n    print(f\"  Match: {len(submission) == len(test_df)}\")\n    print(f\"  All sample_ids unique: {submission['sample_id'].is_unique}\")\n    \n    print(f\"\\nFirst 10 rows of submission:\")\n    print(submission.head(10))\n    \n    return submission\n\n\n# ================== 9. 主程序（适配训练代码的执行逻辑）==================\nif __name__ == '__main__':\n    if CFG.CREATE_SUBMISSION:\n        # 执行推理（适配训练代码的流程）\n        predictions_direct, test_filenames = run_inference()\n        # 后处理（和训练代码一致）\n        predictions_all = postprocess_predictions(predictions_direct)\n        # 生成提交文件\n        submission_df = create_submission(predictions_all, test_filenames)\n    \n    print(\"\\n✅ Inference completed successfully!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}