{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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":14165966,"sourceType":"datasetVersion","datasetId":9029583},{"sourceId":14177546,"sourceType":"datasetVersion","datasetId":9037625},{"sourceId":14308733,"sourceType":"datasetVersion","datasetId":9134440},{"sourceId":14309590,"sourceType":"datasetVersion","datasetId":9135029},{"sourceId":14329629,"sourceType":"datasetVersion","datasetId":9148401}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nimport timm\nfrom tqdm import tqdm\n\n# 屏蔽无关警告\nwarnings.filterwarnings('ignore', category=UserWarning)\nwarnings.filterwarnings('ignore', category=FutureWarning)\n\n# ==========================================\n# 配置（已适配你的权重文件名）\n# ==========================================\nclass Config:\n    model_name = 'vit_base_patch14_dinov2.lvd142m'\n    img_size = 224\n    batch_size = 16\n    num_workers = 4\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # 你的权重文件夹路径（正确）\n    MODEL_DIR = '/kaggle/input/metal-logmetal-log'\n    # 权重文件关键词（适配带空格的文件名）\n    WEIGHT_KEYWORD = 'fold'  # 只要文件名包含\"fold\"就匹配\n    \n    META_DIM = 15\n    N_FOLD = 5\n\n# ==========================================\n# 模型定义（不变）\n# ==========================================\nclass BiomassModelMeta(nn.Module):\n    def __init__(self, model_name, pretrained=False, meta_dim=15):\n        super().__init__()\n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=pretrained, \n            num_classes=0, \n            img_size=Config.img_size,\n            global_pool='avg'\n        )\n        visual_dim = self.backbone.num_features * 2\n        \n        self.meta_mlp = nn.Sequential(\n            nn.Linear(meta_dim, 64),\n            nn.BatchNorm1d(64),\n            nn.ReLU(),\n            nn.Linear(64, 64),\n            nn.ReLU()\n        )\n        \n        in_features = visual_dim + 64\n        \n        def make_head(out_dim=1):\n            return nn.Sequential(\n                nn.Linear(in_features, 512),\n                nn.LayerNorm(512),\n                nn.GELU(),\n                nn.Dropout(0.3),\n                nn.Linear(512, out_dim)\n            )\n        \n        self.head_green = make_head()\n        self.head_clover = make_head()\n        self.head_dead = make_head()\n\n    def forward(self, left_x, right_x, meta_x):\n        left_feat = self.backbone(left_x)\n        right_feat = self.backbone(right_x)\n        visual_feat = torch.cat([left_feat, right_feat], dim=1)\n        meta_feat = self.meta_mlp(meta_x)\n        combined = torch.cat([visual_feat, meta_feat], dim=1)\n        green = self.head_green(combined)\n        clover = self.head_clover(combined)\n        dead = self.head_dead(combined)\n        return torch.cat([green, clover, dead], dim=1)\n\n# ==========================================\n# 数据处理（不变）\n# ==========================================\ndef get_transforms(img_size):\n    interp = transforms.InterpolationMode.BICUBIC\n    return transforms.Compose([\n        transforms.Resize((img_size, img_size), interpolation=interp),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ])\n\nclass TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.transform = transform\n        self.default_meta = np.zeros(Config.META_DIM, dtype=np.float32)\n        self.default_meta[14] = 1.0 \n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = row['image_path']\n        try:\n            image = Image.open(img_path).convert('RGB')\n        except:\n            image = Image.new('RGB', (2000, 1000))\n        w, h = image.size\n        mid = w // 2\n        left_image = image.crop((0, 0, mid, h))\n        right_image = image.crop((mid, 0, w, h))\n        if self.transform:\n            left_image = self.transform(left_image)\n            right_image = self.transform(right_image)\n        meta = torch.tensor(self.default_meta, dtype=torch.float32)\n        return left_image, right_image, meta, row['id_prefix']\n\n# ==========================================\n# 推理主逻辑（核心修复：权重加载）\n# ==========================================\ndef run_inference():\n    DATA_ROOT = '/kaggle/input/csiro-biomass'\n    test_csv = os.path.join(DATA_ROOT, 'test.csv')\n    if not os.path.exists(test_csv):\n        print(\"❌ Test CSV not found!\")\n        return\n\n    print(f\"📂 Reading test data from {DATA_ROOT}...\")\n    test_df_raw = pd.read_csv(test_csv)\n    test_df_raw['id_prefix'] = test_df_raw['sample_id'].apply(lambda x: x.split('__')[0])\n    test_df_raw['target_name'] = test_df_raw['sample_id'].apply(lambda x: x.split('__')[1])\n    unique_ids = test_df_raw[['id_prefix']].drop_duplicates().reset_index(drop=True)\n    \n    def find_image_path(prefix):\n        candidate_paths = [\n            f\"{DATA_ROOT}/test/{prefix}.jpg\",\n            f\"{DATA_ROOT}/{prefix}.jpg\",\n            f\"{DATA_ROOT}/test/{prefix}.png\",\n            f\"{DATA_ROOT}/{prefix}.png\"\n        ]\n        for p in candidate_paths:\n            if os.path.exists(p):\n                return p\n        print(f\"⚠️ 未找到图像 {prefix}\")\n        return candidate_paths[0]\n    unique_ids['image_path'] = unique_ids['id_prefix'].apply(find_image_path)\n    \n    dataset = TestDataset(unique_ids, transform=get_transforms(Config.img_size))\n    loader = DataLoader(\n        dataset, \n        batch_size=Config.batch_size, \n        shuffle=False, \n        num_workers=Config.num_workers,\n        pin_memory=True,\n        drop_last=False\n    )\n    \n    # ==========================================\n    # 🔴 核心修复：自动匹配带空格的权重文件\n    # ==========================================\n    models = []\n    print(\"\\n🚀 Loading Models...\")\n    \n    # 遍历权重文件夹，匹配所有包含\"fold\"的文件（自动处理空格）\n    for filename in os.listdir(Config.MODEL_DIR):\n        # 忽略非pth文件\n        if not filename.endswith('.pth'):\n            continue\n        # 只要文件名包含\"fold\"，就尝试加载\n        if Config.WEIGHT_KEYWORD in filename:\n            model_path = os.path.join(Config.MODEL_DIR, filename)\n            try:\n                model = BiomassModelMeta(Config.model_name, pretrained=False, meta_dim=Config.META_DIM)\n                # 加载权重（忽略空格/命名差异）\n                state_dict = torch.load(model_path, map_location=Config.device)\n                model.load_state_dict(state_dict, strict=False)\n                model.to(Config.device)\n                model.eval()\n                models.append(model)\n                print(f\"✅ 成功加载：{filename}\")\n            except Exception as e:\n                print(f\"⚠️ 加载失败 {filename} | 错误：{e}\")\n                continue\n    \n    if not models:\n        print(\"❌ No models loaded!\")\n        return\n    print(f\"✅ 共加载 {len(models)} 个模型\\n\")\n\n    # 推理（不变）\n    results_map = {}\n    print(\"⚡ Starting Inference...\")\n    with torch.no_grad():\n        for left, right, meta, ids in tqdm(loader):\n            left = left.to(Config.device)\n            right = right.to(Config.device)\n            meta = meta.to(Config.device)\n            \n            tta_transforms = [\n                (left, right),\n                (torch.flip(left, dims=[3]), torch.flip(right, dims=[3])),\n                (torch.flip(left, dims=[2]), torch.flip(right, dims=[2])),\n                (torch.flip(left, dims=[2,3]), torch.flip(right, dims=[2,3]))\n            ]\n            \n            batch_preds = torch.zeros((left.size(0), 3)).to(Config.device)\n            for model in models:\n                for l_in, r_in in tta_transforms:\n                    log_preds = model(l_in, r_in, meta)\n                    batch_preds += torch.expm1(log_preds)\n            \n            avg_preds = batch_preds / (len(models)*len(tta_transforms))\n            avg_preds = torch.clamp(avg_preds, min=0.0).cpu().numpy()\n            for i, prefix in enumerate(ids):\n                results_map[prefix] = avg_preds[i]\n\n    # 生成提交文件（不变）\n    print(\"\\n📝 Generating Submission...\")\n    submission_rows = []\n    for _, row in test_df_raw.iterrows():\n        prefix = row['id_prefix']\n        target_name = row['target_name']\n        g, c, d = results_map.get(prefix, [0,0,0])\n        if target_name == 'Dry_Green_g': val = g\n        elif target_name == 'Dry_Clover_g': val = c\n        elif target_name == 'Dry_Dead_g': val = d\n        elif target_name == 'GDM_g': val = g + c\n        elif target_name == 'Dry_Total_g': val = g + c + d\n        else: val = 0\n        submission_rows.append({'sample_id': row['sample_id'], 'target': val})\n    \n    sub_df = pd.DataFrame(submission_rows)\n    sub_df.to_csv('submission.csv', index=False)\n    print(\"✅ Done! submission.csv saved.\")\n    print(sub_df.head(10))\n\nif __name__ == '__main__':\n    run_inference()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}