{"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,"sourceType":"competition"},{"sourceId":272104372,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\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\nimport cv2\nfrom tqdm import tqdm\nimport gc\n\n# ===============================================================\n# 1. ⚙️ 配置 \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    \n    # 包含 5 个 .pth 文件的目录\n    MODEL_DIR = '/kaggle/input/csiro/'\n    SUBMISSION_FILE = 'submission.csv'\n    \n    # --- 模型设置 ---\n    MODEL_NAME = 'convnext_tiny'\n    IMG_SIZE = 768              \n    \n    # --- 推理设置 ---\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    BATCH_SIZE = 1 # 推理时可以增大学习率\n    NUM_WORKERS = 1\n    N_FOLDS = 5\n    \n    # --- 目标与损失 (必须匹配) ---\n    # 模型预测的 3 个目标\n    TARGET_COLS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    \n    # 提交所需的 5 个目标\n    ALL_TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n\nprint(f\"使用设备: {CFG.DEVICE}\")\nprint(f\"模型骨干: {CFG.MODEL_NAME}\")\nprint(f\"推理图片尺寸: {CFG.IMG_SIZE}x{CFG.IMG_SIZE}\")\n\n\n# ===============================================================\n# 2. 🏞️ 数据增强\n# ===============================================================\nfrom albumentations import (\n    Compose, \n    Resize, \n    Normalize,\n    HorizontalFlip, \n    VerticalFlip\n)\n\ndef get_tta_transforms():\n    \"\"\"\n    返回一个用于 TTA 的 transform pipeline 列表。\n    每个 pipeline 都是图像的一个不同“视角”。\n    \"\"\"\n    \n    # 这是基本的标准化步骤\n    base_transforms = [\n        Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        ),\n        ToTensorV2()\n    ]\n    \n    # -----------------\n    # 视角 1: 原始图像 \n    # -----------------\n    original_view = Compose([\n        Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        *base_transforms\n    ])\n    \n    # -----------------\n    # 视角 2: 水平翻转 \n    # -----------------\n    hflip_view = Compose([\n        HorizontalFlip(p=1.0), # 始终翻转\n        Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        *base_transforms\n    ])\n    \n    # -----------------\n    # 视角 3: 垂直翻转 \n    # -----------------\n    vflip_view = Compose([\n        VerticalFlip(p=1.0), # 始终翻转\n        Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        *base_transforms\n    ])\n    \n    return [original_view, hflip_view, vflip_view]\n\nprint(\"已定义 get_tta_transforms() 函数。\")\n\n\nclass TestBiomassDataset(Dataset):\n    \"\"\"\n    用于测试图像的自定义 Dataset (双流策略)。\n    修改为接受一个特定的 TTA transform pipeline。\n    \"\"\"\n    def __init__(self, df, transform_pipeline, image_dir):\n        self.df = df\n        self.transforms = transform_pipeline \n        self.image_dir = image_dir\n        self.image_paths = df['image_path'].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        # 1. 获取信息\n        img_path_suffix = self.image_paths[idx]\n        \n        # 2. 读取原始图像 (2000x1000)\n        filename = os.path.basename(img_path_suffix)\n        full_path = os.path.join(self.image_dir, filename)\n        \n        image = cv2.imread(full_path)\n        if image is None:\n            print(f\"警告: 无法读取图像: {full_path}。返回黑色图像。\")\n            image = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        \n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # 3. 裁剪成 2 张图像 (左侧和右侧)\n        height, width, _ = image.shape\n        mid_point = width // 2\n        img_left = image[:, :mid_point]\n        img_right = image[:, mid_point:]\n        \n        # 4. 应用 TTA Transform (对两者使用相同的 Transform)\n        img_left_tensor = self.transforms(image=img_left)['image']\n        img_right_tensor = self.transforms(image=img_right)['image']\n        \n        # 5. 返回\n        return img_left_tensor, img_right_tensor\n\n# ===============================================================\n# 4. 🧠 模型架构\n# ===============================================================\nclass BiomassModel(nn.Module):\n    \"\"\"\n    模型架构 (双流, 三输出)\n    必须与训练文件完全一致。\n    \"\"\"\n    def __init__(self, model_name, pretrained, n_targets=3):\n        super(BiomassModel, self).__init__()\n        \n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=pretrained,\n            num_classes=0,\n            global_pool='avg'\n        )\n        \n        self.n_features = self.backbone.num_features\n        self.n_combined_features = self.n_features * 2\n        \n        # --- Dry_Total_g 的输出头 ---\n        self.head_total = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)\n        )\n        \n        # --- GDM_g 的输出头 ---\n        self.head_gdm = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)\n        )\n        \n        # --- Dry_Green_g 的输出头 ---\n        self.head_green = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)\n        )\n\n    def forward(self, img_left, img_right):\n        features_left = self.backbone(img_left)\n        features_right = self.backbone(img_right)\n        combined = torch.cat([features_left, features_right], dim=1)\n        \n        out_total = self.head_total(combined)\n        out_gdm = self.head_gdm(combined)\n        out_green = self.head_green(combined)\n        \n        return out_total, out_gdm, out_green\n\n\ndef predict_one_view(models_list, test_loader, device):\n    \"\"\"\n    子函数: 为单个 TTA 视角运行 5 折集成预测。\n    \"\"\"\n    view_preds_3 = {'total': [], 'gdm': [], 'green': []}\n    \n    with torch.no_grad():\n        for (img_left, img_right) in tqdm(test_loader, desc=\"  Predicting View\", leave=False):\n            img_left = img_left.to(device)\n            img_right = img_right.to(device)\n            \n            batch_preds_3_folds = {'total': [], 'gdm': [], 'green': []}\n            \n            # 1. 5 折集成循环\n            for model in models_list:\n                pred_total, pred_gdm, pred_green = model(img_left, img_right)\n                batch_preds_3_folds['total'].append(pred_total.cpu())\n                batch_preds_3_folds['gdm'].append(pred_gdm.cpu())\n                batch_preds_3_folds['green'].append(pred_green.cpu())\n            \n            # 2. 计算 5 折的平均值\n            avg_pred_total = torch.mean(torch.stack(batch_preds_3_folds['total']), dim=0)\n            avg_pred_gdm = torch.mean(torch.stack(batch_preds_3_folds['gdm']), dim=0)\n            avg_pred_green = torch.mean(torch.stack(batch_preds_3_folds['green']), dim=0)\n            \n            view_preds_3['total'].append(avg_pred_total.numpy())\n            view_preds_3['gdm'].append(avg_pred_gdm.numpy())\n            view_preds_3['green'].append(avg_pred_green.numpy())\n\n    # 3. 合并此视角的批处理结果\n    preds_np = {\n        'total': np.concatenate(view_preds_3['total']).flatten(),\n        'gdm':   np.concatenate(view_preds_3['gdm']).flatten(),\n        'green': np.concatenate(view_preds_3['green']).flatten()\n    }\n    return preds_np\n\n\ndef run_inference_with_tta():\n    \"\"\"\n    主推理函数, 执行 TTA x 集成。\n    \"\"\"\n    print(f\"\\n{'='*50}\")\n    print(f\"🚀 开始推理 (使用 TTA) 🚀\")\n    print(f\"{'='*50}\")\n\n    # --- 1. 加载测试数据 ---\n    print(f\"正在加载 {CFG.TEST_CSV}...\")\n    try:\n        test_df_long = pd.read_csv(CFG.TEST_CSV)\n        test_df_unique = test_df_long.drop_duplicates(subset=['image_path']).reset_index(drop=True)\n        print(f\"找到 {len(test_df_unique)} 张唯一的测试图像。\")\n    except FileNotFoundError:\n        print(f\"错误: 未找到 {CFG.TEST_CSV}\")\n        return None, None, None\n\n    # --- 2. 加载 5 个模型 (集成) ---\n    print(\"\\n正在加载 5 个已训练的模型...\")\n    models_list = []\n    for fold in range(CFG.N_FOLDS):\n        model_path = os.path.join(CFG.MODEL_DIR, f'best_model_fold{fold}.pth')\n        if not os.path.exists(model_path):\n            print(f\"错误: 未找到模型文件: {model_path}\")\n            return None, None, None\n        model = BiomassModel(CFG.MODEL_NAME, pretrained=False)\n        try:\n            model.load_state_dict(torch.load(model_path, map_location=CFG.DEVICE))\n        except RuntimeError:\n            state_dict = torch.load(model_path, map_location=CFG.DEVICE)\n            from collections import OrderedDict\n            new_state_dict = OrderedDict()\n            for k, v in state_dict.items():\n                name = k.replace('module.', '')\n                new_state_dict[name] = v\n            model.load_state_dict(new_state_dict)\n        model.eval()\n        model.to(CFG.DEVICE)\n        models_list.append(model)\n    print(f\"✓ 成功加载 {len(models_list)} 个模型。\")\n\n    # --- 3. TTA 循环  ---\n    tta_transforms = get_tta_transforms()\n    print(f\"\\n开始使用 {len(tta_transforms)} 个 TTA 视角进行预测...\")\n    \n    all_tta_view_preds = [] \n\n    for i, tta_transform in enumerate(tta_transforms):\n        print(f\"--- 正在运行 TTA 视角 {i+1}/{len(tta_transforms)} ---\")\n        \n        # 为此 TTA 视角创建新的 Dataset/Loader\n        test_dataset = TestBiomassDataset(\n            df=test_df_unique,\n            transform_pipeline=tta_transform,\n            image_dir=CFG.TEST_IMAGE_DIR\n        )\n        test_loader = DataLoader(\n            test_dataset,\n            batch_size=CFG.BATCH_SIZE,\n            shuffle=False,\n            num_workers=CFG.NUM_WORKERS,\n            pin_memory=True\n        )\n        \n        # 为此视角运行 5 折集成\n        view_preds_np = predict_one_view(models_list, test_loader, CFG.DEVICE)\n        all_tta_view_preds.append(view_preds_np)\n        print(f\"✓ 完成 TTA 视角 {i+1}\")\n\n    # --- 4. 集成 (平均) TTA 结果 ---\n    print(\"\\n正在集成 TTA 视角的结果...\")\n    final_ensembled_preds = {\n        'total': np.mean([d['total'] for d in all_tta_view_preds], axis=0),\n        'gdm':   np.mean([d['gdm'] for d in all_tta_view_preds], axis=0),\n        'green': np.mean([d['green'] for d in all_tta_view_preds], axis=0)\n    }\n    \n    print(\"✓ 预测完成。\")\n    \n    del models_list, test_loader, test_dataset\n    gc.collect()\n    torch.cuda.empty_cache()\n    \n    return final_ensembled_preds, test_df_long, test_df_unique\n# ===============================================================\n# 6. ✍️ 创建提交文件的函数\n# ===============================================================\ndef create_submission(preds_np, test_df_long, test_df_unique):\n    \"\"\"\n    此函数接收 3 个集成后的预测结果，\n    计算剩余 2 个预测结果，\n    并格式化提交文件。\n    \"\"\"\n    if preds_np is None:\n        print(\"因上述错误，跳过创建提交文件。\")\n        return\n\n    print(\"\\n正在后处理并创建提交文件...\")\n\n    # 1. 获取 3 个集成后的预测结果\n    pred_total_final = preds_np['total']\n    pred_gdm_final = preds_np['gdm']\n    pred_green_final = preds_np['green']\n\n    # 2. 计算剩余 2 个目标 (后处理)\n    # 使用 np.maximum(0, ...) 确保没有负值\n    pred_clover_final = np.maximum(0, pred_gdm_final - pred_green_final)\n    pred_dead_final = np.maximum(0, pred_total_final - pred_gdm_final)\n\n    # 3. 创建一个包含 5 个预测结果的“宽”DataFrame\n    # (确保 5 个列的顺序与 CFG.ALL_TARGET_COLS 相同)\n    preds_wide_df = pd.DataFrame({\n        'image_path': test_df_unique['image_path'],\n        'Dry_Green_g': pred_green_final,\n        'Dry_Dead_g': pred_dead_final,\n        'Dry_Clover_g': pred_clover_final,\n        'GDM_g': pred_gdm_final,\n        'Dry_Total_g': pred_total_final\n    })\n\n    # 4. “逆透视” DataFrame (转换为“长”格式)\n    # 将其从 5 列转换为“长”格式 (与 sample_submission 类似)\n    preds_long_df = preds_wide_df.melt(\n        id_vars=['image_path'],\n        value_vars=CFG.ALL_TARGET_COLS, # 5 个目标列\n        var_name='target_name',        # 目标名称列\n        value_name='target'            # 预测值列\n    )\n\n    # 5. 与原始 test.csv 文件 (test_df_long) 合并\n    # 这是获取正确 'sample_id' 的关键步骤\n    # (例如: 'ID1001187975__Dry_Clover_g')\n    submission_df = pd.merge(\n        test_df_long[['sample_id', 'image_path', 'target_name']],\n        preds_long_df,\n        on=['image_path', 'target_name'],\n        how='left'\n    )\n\n    # 6. 清理并保存\n    # 仅保留所需的 2 个列\n    submission_df = submission_df[['sample_id', 'target']]\n    \n    # 保存文件\n    submission_df.to_csv(CFG.SUBMISSION_FILE, index=False)\n\n    print(f\"\\n🎉 完成! 提交文件已保存至: {CFG.SUBMISSION_FILE}\")\n    print(\"--- 提交文件的前 5 行 ---\")\n    print(submission_df.head())\n    print(\"\\n--- 提交文件的后 5 行 ---\")\n    print(submission_df.tail())\n    \n# ===============================================================\n# 8. 🏁 运行程序 (已修改)\n# ===============================================================\nif __name__ == \"__main__\":\n    # 1. 运行预测 (已包含 TTA)\n    all_preds_np, df_long, df_unique = run_inference_with_tta()\n    \n    # 2. 创建提交文件 (create_submission 函数保持不变)\n    create_submission(all_preds_np, df_long, df_unique)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-30T13:37:31.585364Z","iopub.execute_input":"2025-10-30T13:37:31.586001Z","iopub.status.idle":"2025-10-30T13:37:37.021319Z","shell.execute_reply.started":"2025-10-30T13:37:31.585975Z","shell.execute_reply":"2025-10-30T13:37:37.020506Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}