{"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":"none","dataSources":[{"sourceId":19989,"databundleVersionId":1160143,"sourceType":"competition"},{"sourceId":1132110,"sourceType":"datasetVersion","datasetId":637207},{"sourceId":13986085,"sourceType":"datasetVersion","datasetId":8914930}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 导入所有必要的库\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport os\nimport re\nimport torch\nimport torchvision\nfrom torchvision import transforms \nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import fasterrcnn_resnet50_fpn\nfrom torch.utils.data import DataLoader, Dataset\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\nfrom pathlib import Path\nfrom tqdm import tqdm\nimport matplotlib.font_manager as fm\n\n# 设置中文字体\nfont_path = \"/kaggle/input/chinese-font/STKAITI.TTF\"\nfm.fontManager.addfont(font_path)\nprop = fm.FontProperties(fname=font_path)\nfont_name = prop.get_name()\nplt.rcParams['font.sans-serif'] = [font_name]\nplt.rcParams['axes.unicode_minus'] = False\n\n# ==================== 数据准备部分 ====================\nprint(\"=\"*50)\nprint(\"数据准备\")\nprint(\"=\"*50)\n\n# 读取数据\ntrain_df = pd.read_csv(\"/kaggle/input/global-wheat-detection/train.csv\")\nsubmit = pd.read_csv(\"/kaggle/input/global-wheat-detection/sample_submission.csv\")\n\n# 数据预处理\ntrain_df = train_df.drop(columns=['width', 'height', 'source'])\n\n# 解析边界框坐标\ntrain_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x))) \ntrain_df['x'] = train_df['x'].astype(np.float16)\ntrain_df['y'] = train_df['y'].astype(np.float16)\ntrain_df['w'] = train_df['w'].astype(np.float16)\ntrain_df['h'] = train_df['h'].astype(np.float16)\n\n# 分割训练集和验证集\nimage_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\n\nprint(f\"训练集大小: {len(train_ids)}\")\nprint(f\"验证集大小: {len(valid_ids)}\")\n\n# ==================== 数据集类 ====================\nclass WheatDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transforms=None, train=True):\n        super().__init__()\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n        self.train = train\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        \n        if self.transforms is not None:  \n            image = self.transforms(image)\n        \n        if not self.train:  \n            return image, image_id\n        \n        # 训练模式：返回图像和标注\n        records = self.df[self.df['image_id'] == image_id]   \n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float16)\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int32)\n        \n        target = {\n            'boxes': boxes,\n            'labels': labels,\n            'image_id': torch.tensor([index]),\n            'area': area,\n            'iscrowd': iscrowd\n        }\n\n        return image, target, image_id\n\n# ==================== 数据加载器 ====================\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# 数据路径\ntrain_dir = '/kaggle/input/global-wheat-detection/train'\ntest_dir = '/kaggle/input/global-wheat-detection/test'\n\n# 数据转换\ntrans = transforms.Compose([transforms.ToTensor()])\n\n# 创建数据集\ntrain_dataset = WheatDataset(train_df, train_dir, trans, True)\nvalid_dataset = WheatDataset(valid_df, train_dir, trans, True)\ntest_dataset = WheatDataset(submit, test_dir, trans, False)\n\n# 创建数据加载器\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=2,\n    shuffle=False,\n    num_workers=1,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=1,\n    shuffle=False,\n    num_workers=1,\n    collate_fn=collate_fn\n)\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=1,\n    shuffle=False,\n    num_workers=1,\n    collate_fn=collate_fn\n)\n\n# ==================== 模型训练部分 ====================\nprint(\"\\n\" + \"=\"*50)\nprint(\"模型训练\")\nprint(\"=\"*50)\n\n# 设备设置\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"使用设备: {device}\")\n\n# 创建模型\nWEIGHTS_FILE = '/kaggle/input/fasterrcnn/fasterrcnn_resnet50_fpn_best.pth'\nmodel = fasterrcnn_resnet50_fpn(weights=None, weights_backbone=None)\nnum_classes = 2\n\n# 修改分类头\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n# 加载预训练权重（如果有的话）\nif os.path.exists(WEIGHTS_FILE):\n    print(f\"加载预训练权重: {WEIGHTS_FILE}\")\n    state_dict = torch.load(WEIGHTS_FILE, map_location=device)\n    model.load_state_dict(state_dict)\nelse:\n    print(\"未找到预训练权重，从头开始训练\")\n\nmodel.to(device)\n\n# 定义平均损失计算器\nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n# 训练一个epoch（示例）\ndef train_one_epoch(model, data_loader, optimizer, device):\n    model.train()\n    loss_hist = Averager()\n    \n    for images, targets, image_ids in data_loader:\n        images = list(image.to(device) for image in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        loss_dict = model(images, targets)\n        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n        \n        loss_hist.send(loss_value)\n        \n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n    \n    return loss_hist.value\n\n# 设置优化器\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=0.00001)\n\n# 训练一个epoch（如果需要的话）\nprint(\"训练一个epoch...\")\ntrain_loss = train_one_epoch(model, train_data_loader, optimizer, device)\nprint(f\"训练损失: {train_loss:.4f}\")\n\n# ==================== 测试集可视化类 ====================\nclass WheatTestVisualizer:\n    def __init__(self, model, device, detection_threshold=0.45):\n        \"\"\"\n        小麦检测测试集可视化器\n        \"\"\"\n        self.model = model\n        self.device = device\n        self.detection_threshold = detection_threshold\n        self.model.eval()\n        \n    def process_single_image(self, image_path, output_path=None, show=True):\n        \"\"\"\n        处理单张测试图像\n        \"\"\"\n        # 读取图像\n        image = cv2.imread(str(image_path))\n        if image is None:\n            print(f\"无法读取图像: {image_path}\")\n            return None, None, None\n            \n        # 转换颜色空间\n        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # 预处理图像\n        image_tensor = torch.from_numpy(image_rgb.transpose(2, 0, 1)).float() / 255.0\n        image_tensor = image_tensor.unsqueeze(0).to(self.device)\n        \n        # 模型推理\n        with torch.no_grad():\n            outputs = self.model(image_tensor)\n        \n        # 提取预测结果\n        boxes = outputs[0]['boxes'].data.cpu().numpy()\n        scores = outputs[0]['scores'].data.cpu().numpy()\n        labels = outputs[0]['labels'].data.cpu().numpy()\n        \n        # 过滤低置信度框\n        keep_indices = scores >= self.detection_threshold\n        boxes = boxes[keep_indices]\n        scores = scores[keep_indices]\n        labels = labels[keep_indices]\n        \n        # 绘制预测框\n        image_with_boxes = self.draw_boxes(image_rgb.copy(), boxes, scores)\n        \n        # 显示结果\n        if show:\n            self.display_result(image_with_boxes, image_path.name, boxes, scores)\n        \n        # 保存结果\n        if output_path:\n            cv2.imwrite(str(output_path), cv2.cvtColor(image_with_boxes, cv2.COLOR_RGB2BGR))\n            print(f\"结果已保存至: {output_path}\")\n        \n        return image_with_boxes, boxes, scores\n    \n    def draw_boxes(self, image, boxes, scores):\n        \"\"\"\n        在图像上绘制边界框\n        \"\"\"\n        image_with_boxes = image.copy()\n        \n        for i, (box, score) in enumerate(zip(boxes, scores)):\n            # 转换坐标为整数\n            x1, y1, x2, y2 = map(int, box)\n            \n            # 计算框的颜色（基于置信度）\n            color_intensity = int(score * 255)\n            color = (0, 255 - color_intensity, color_intensity)  # 从蓝色到绿色\n            \n            # 绘制边界框\n            cv2.rectangle(image_with_boxes, (x1, y1), (x2, y2), color, 3)\n            \n            # 准备标签文本\n            label_text = f'小麦: {score:.2f}'\n            \n            # 计算文本尺寸\n            (text_width, text_height), baseline = cv2.getTextSize(\n                label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2\n            )\n            \n            # 绘制文本背景\n            cv2.rectangle(\n                image_with_boxes,\n                (x1, y1 - text_height - 10),\n                (x1 + text_width, y1),\n                color,\n                -1\n            )\n            \n            # 绘制文本\n            cv2.putText(\n                image_with_boxes,\n                label_text,\n                (x1, y1 - 5),\n                cv2.FONT_HERSHEY_SIMPLEX,\n                0.8,\n                (255, 255, 255),\n                2,\n                cv2.LINE_AA\n            )\n        \n        return image_with_boxes\n    \n    def display_result(self, image, image_name, boxes, scores):\n        \"\"\"\n        显示结果图像\n        \"\"\"\n        fig, axes = plt.subplots(1, 2, figsize=(20, 10))\n        \n        # 左侧：带预测框的图像\n        axes[0].imshow(image)\n        axes[0].set_title(f'预测结果 - {image_name}', fontsize=16)\n        axes[0].axis('off')\n        \n        # 右侧：统计信息\n        axes[1].axis('off')\n        \n        # 创建统计信息文本\n        stats_text = f\"\"\"\n        图像名称: {image_name}\n        \n        检测统计:\n        --------------------\n        总检测框数: {len(boxes)}\n        平均置信度: {scores.mean():.3f if len(scores) > 0 else 0}\n        最高置信度: {scores.max() if len(scores) > 0 else 0:.3f}\n        最低置信度: {scores.min() if len(scores) > 0 else 0:.3f}\n        \n        置信度分布:\n        --------------------\n        \"\"\"\n        \n        # 添加置信度分布\n        if len(scores) > 0:\n            bins = [0.45, 0.6, 0.7, 0.8, 0.9, 1.0]\n            hist, _ = np.histogram(scores, bins=bins)\n            for i, (low, high) in enumerate(zip(bins[:-1], bins[1:])):\n                stats_text += f\"{low:.1f}-{high:.1f}: {hist[i]}个框\\n\"\n        \n        axes[1].text(0.1, 0.9, stats_text, fontsize=14, \n                    verticalalignment='top',\n                    transform=axes[1].transAxes)\n        \n        plt.tight_layout()\n        plt.show()\n    \n    def process_test_directory(self, test_dir, output_dir=None, max_images=None):\n        \"\"\"\n        批量处理测试目录中的所有图像\n        \"\"\"\n        # 获取所有测试图像\n        test_path = Path(test_dir)\n        image_paths = list(test_path.glob('*.jpg')) + \\\n                     list(test_path.glob('*.png')) + \\\n                     list(test_path.glob('*.jpeg'))\n        \n        if not image_paths:\n            print(f\"在目录 {test_dir} 中未找到图像\")\n            return None\n        \n        # 限制处理数量\n        if max_images:\n            image_paths = image_paths[:max_images]\n        \n        print(f\"开始处理 {len(image_paths)} 张测试图像...\")\n        \n        # 创建输出目录\n        if output_dir:\n            output_path = Path(output_dir)\n            output_path.mkdir(parents=True, exist_ok=True)\n        \n        # 存储所有结果\n        all_results = []\n        \n        # 处理每张图像\n        for img_path in tqdm(image_paths, desc=\"处理测试图像\"):\n            try:\n                # 设置输出路径\n                save_path = None\n                if output_dir:\n                    save_path = output_path / f\"pred_{img_path.name}\"\n                \n                # 处理单张图像\n                result_img, boxes, scores = self.process_single_image(\n                    img_path, \n                    output_path=save_path, \n                    show=False\n                )\n                \n                if result_img is not None:\n                    # 收集结果信息\n                    result_info = {\n                        'image_id': img_path.stem,\n                        'num_detections': len(boxes),\n                        'avg_score': scores.mean() if len(scores) > 0 else 0,\n                        'max_score': scores.max() if len(scores) > 0 else 0,\n                        'min_score': scores.min() if len(scores) > 0 else 0,\n                        'boxes': boxes.tolist() if len(boxes) > 0 else [],\n                        'scores': scores.tolist() if len(scores) > 0 else []\n                    }\n                    all_results.append(result_info)\n                    \n                    # 打印进度\n                    print(f\"{img_path.name}: 检测到 {len(boxes)} 个小麦穗\")\n                    \n            except Exception as e:\n                print(f\"处理图像 {img_path} 时出错: {str(e)}\")\n                continue\n        \n        # 创建结果DataFrame\n        results_df = pd.DataFrame(all_results)\n        \n        # 打印总体统计\n        self.print_summary_statistics(results_df)\n        \n        # 保存结果到CSV\n        if output_dir:\n            csv_path = output_path / \"detection_results.csv\"\n            results_df.to_csv(csv_path, index=False)\n            print(f\"详细结果已保存至: {csv_path}\")\n        \n        return results_df\n    \n    def print_summary_statistics(self, results_df):\n        \"\"\"\n        打印总体统计信息\n        \"\"\"\n        if results_df.empty:\n            print(\"没有可用的检测结果\")\n            return\n        \n        print(\"\\n\" + \"=\"*50)\n        print(\"测试集检测结果汇总\")\n        print(\"=\"*50)\n        \n        total_images = len(results_df)\n        total_detections = results_df['num_detections'].sum()\n        avg_detections = results_df['num_detections'].mean()\n        \n        print(f\"处理图像总数: {total_images}\")\n        print(f\"总检测框数: {total_detections}\")\n        print(f\"平均每张图像检测数: {avg_detections:.2f}\")\n        print(f\"平均置信度: {results_df['avg_score'].mean():.3f}\")\n        \n        # 检测数量分布\n        print(\"\\n检测数量分布:\")\n        det_counts = results_df['num_detections'].value_counts().sort_index()\n        for count, freq in det_counts.items():\n            print(f\"  检测到 {count} 个小麦穗: {freq} 张图像\")\n        \n        # 置信度分布\n        print(\"\\n置信度分布:\")\n        all_scores = []\n        for scores_list in results_df['scores']:\n            all_scores.extend(scores_list)\n        \n        if all_scores:\n            all_scores = np.array(all_scores)\n            print(f\"  最高置信度: {all_scores.max():.3f}\")\n            print(f\"  最低置信度: {all_scores.min():.3f}\")\n            print(f\"  平均置信度: {all_scores.mean():.3f}\")\n            \n            # 显示置信度直方图\n            plt.figure(figsize=(10, 6))\n            plt.hist(all_scores, bins=20, alpha=0.7, color='blue', edgecolor='black')\n            plt.xlabel('置信度分数')\n            plt.ylabel('频率')\n            plt.title('所有检测框的置信度分布')\n            plt.grid(True, alpha=0.3)\n            plt.show()\n    \n    def visualize_sample_predictions(self, test_dir, num_samples=6):\n        \"\"\"\n        可视化样本预测结果\n        \"\"\"\n        test_path = Path(test_dir)\n        image_paths = list(test_path.glob('*.jpg'))[:num_samples]\n        \n        if not image_paths:\n            print(f\"在目录 {test_dir} 中未找到图像\")\n            return\n        \n        # 创建子图\n        fig, axes = plt.subplots(2, 3, figsize=(20, 12))\n        axes = axes.flatten()\n        \n        for idx, (ax, img_path) in enumerate(zip(axes, image_paths)):\n            try:\n                # 读取图像\n                image = cv2.imread(str(img_path))\n                if image is None:\n                    continue\n                    \n                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n                image_tensor = torch.from_numpy(image_rgb.transpose(2, 0, 1)).float() / 255.0\n                image_tensor = image_tensor.unsqueeze(0).to(self.device)\n                \n                # 推理\n                with torch.no_grad():\n                    outputs = self.model(image_tensor)\n                \n                # 提取结果\n                boxes = outputs[0]['boxes'].data.cpu().numpy()\n                scores = outputs[0]['scores'].data.cpu().numpy()\n                \n                # 过滤\n                keep_indices = scores >= self.detection_threshold\n                boxes = boxes[keep_indices]\n                scores = scores[keep_indices]\n                \n                # 绘制到子图\n                for box, score in zip(boxes, scores):\n                    x1, y1, x2, y2 = map(int, box)\n                    rect = Rectangle((x1, y1), x2-x1, y2-y1, \n                                    linewidth=2, edgecolor='red', facecolor='none')\n                    ax.add_patch(rect)\n                    \n                    # 添加标签\n                    ax.text(x1, y1-5, f'{score:.2f}', \n                           fontsize=10, color='red', \n                           bbox=dict(facecolor='white', alpha=0.7))\n                \n                ax.imshow(image_rgb)\n                ax.set_title(f'{img_path.stem}\\n检测数: {len(boxes)}', fontsize=12)\n                ax.axis('off')\n                \n            except Exception as e:\n                print(f\"处理图像 {img_path} 时出错: {str(e)}\")\n                ax.axis('off')\n                continue\n        \n        # 隐藏多余的子图\n        for idx in range(len(image_paths), len(axes)):\n            axes[idx].axis('off')\n        \n        plt.suptitle('小麦检测测试集样本预测结果', fontsize=16)\n        plt.tight_layout()\n        plt.show()\n    \n    def process_with_data_loader(self, data_loader, output_dir=None):\n        \"\"\"\n        使用DataLoader处理测试集\n        \"\"\"\n        print(\"使用DataLoader处理测试集...\")\n        \n        model.eval()\n        detection_threshold = self.detection_threshold\n        \n        all_results = []\n        \n        for images, image_ids in tqdm(data_loader, desc=\"处理测试批次\"):\n            images = list(image.to(device) for image in images)\n            \n            with torch.no_grad():\n                outputs = model(images)\n            \n            for i, image in enumerate(images):\n                boxes = outputs[i]['boxes'].data.cpu().numpy()\n                scores = outputs[i]['scores'].data.cpu().numpy()\n                \n                # 过滤低置信度框\n                keep_indices = scores >= detection_threshold\n                boxes = boxes[keep_indices]\n                scores = scores[keep_indices]\n                \n                # 保存结果\n                result_info = {\n                    'image_id': image_ids[i],\n                    'num_detections': len(boxes),\n                    'avg_score': scores.mean() if len(scores) > 0 else 0,\n                    'boxes': boxes.tolist(),\n                    'scores': scores.tolist()\n                }\n                all_results.append(result_info)\n                \n                print(f\"{image_ids[i]}: 检测到 {len(boxes)} 个目标\")\n        \n        # 创建DataFrame\n        results_df = pd.DataFrame(all_results)\n        \n        # 打印统计信息\n        self.print_summary_statistics(results_df)\n        \n        # 保存结果\n        if output_dir:\n            output_path = Path(output_dir)\n            output_path.mkdir(parents=True, exist_ok=True)\n            \n            csv_path = output_path / \"loader_results.csv\"\n            results_df.to_csv(csv_path, index=False)\n            print(f\"结果已保存至: {csv_path}\")\n        \n        return results_df\n\n# ==================== 主执行函数 ====================\nprint(\"\\n\" + \"=\"*50)\nprint(\"测试集可视化\")\nprint(\"=\"*50)\n\n# 创建可视化器\nvisualizer = WheatTestVisualizer(model, device, detection_threshold=0.45)\n\n# 方法1: 使用DataLoader处理测试集\nprint(\"\\n1. 使用DataLoader处理测试集...\")\nloader_results = visualizer.process_with_data_loader(\n    test_data_loader,\n    output_dir='/kaggle/working/loader_predictions'\n)\n\n# 方法2: 可视化样本预测\nprint(\"\\n2. 可视化样本预测结果...\")\nvisualizer.visualize_sample_predictions(\n    test_dir='/kaggle/input/global-wheat-detection/test',\n    num_samples=6\n)\n\n# 方法3: 直接处理测试目录\nprint(\"\\n3. 直接处理测试目录...\")\ntest_results = visualizer.process_test_directory(\n    test_dir='/kaggle/input/global-wheat-detection/test',\n    output_dir='/kaggle/working/direct_predictions',\n    max_images=10  # 限制处理数量\n)\n\n# 方法4: 处理单个图像示例\nprint(\"\\n4. 处理单个图像示例...\")\ntest_images = list(Path('/kaggle/input/global-wheat-detection/test').glob('*.jpg'))\nif test_images:\n    sample_image = test_images[0]\n    print(f\"处理示例图像: {sample_image.name}\")\n    \n    result_img, boxes, scores = visualizer.process_single_image(\n        sample_image,\n        output_path='/kaggle/working/sample_prediction.jpg',\n        show=True\n    )\n    \n    if result_img is not None:\n        print(f\"检测到 {len(boxes)} 个小麦穗\")\n        if len(scores) > 0:\n            print(f\"平均置信度: {scores.mean():.3f}\")\n            print(f\"最高置信度: {scores.max():.3f}\")\n            print(f\"最低置信度: {scores.min():.3f}\")\n\n# ==================== 生成提交文件 ====================\nprint(\"\\n\" + \"=\"*50)\nprint(\"生成提交文件\")\nprint(\"=\"*50)\n\ndef format_prediction_string(boxes, scores):\n    \"\"\"定义存储预测结果的格式\"\"\"\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(\n            j[0], j[1][0], j[1][1], j[1][2], j[1][3]\n        ))\n    return \" \".join(pred_strings)\n\n# 预测并生成提交文件\nresults = []\nmodel.eval()\n\nprint(\"生成提交预测...\")\nfor images, image_ids in test_data_loader:\n    images = list(image.to(device) for image in images)\n    \n    with torch.no_grad():\n        outputs = model(images)\n    \n    for i, image in enumerate(images):\n        boxes = outputs[i]['boxes'].data.cpu().numpy()\n        scores = outputs[i]['scores'].data.cpu().numpy()\n        \n        boxes = boxes[scores >= 0.45].astype(np.int32)\n        scores = scores[scores >= 0.45]\n        \n        image_id = image_ids[i]\n        \n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores) if len(boxes) > 0 else \"\"\n        }\n        results.append(result)\n\n# 创建提交DataFrame\ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('/kaggle/working/submission.csv', index=False)\nprint(f\"提交文件已保存至: /kaggle/working/submission.csv\")\nprint(f\"提交文件大小: {len(test_df)} 行\")\n\n# 显示前几行\nprint(\"\\n提交文件前5行:\")\nprint(test_df.head())\n\nprint(\"\\n\" + \"=\"*50)\nprint(\"任务完成！\")\nprint(\"=\"*50)\nprint(f\"输出文件:\")\nprint(f\"1. 提交文件: /kaggle/working/submission.csv\")\nprint(f\"2. DataLoader预测结果: /kaggle/working/loader_predictions/\")\nprint(f\"3. 直接处理预测结果: /kaggle/working/direct_predictions/\")\nprint(f\"4. 样本预测图像: /kaggle/working/sample_prediction.jpg\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-05T03:24:07.704346Z","iopub.execute_input":"2025-12-05T03:24:07.704647Z"}},"outputs":[],"execution_count":null}]}