{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":29876,"databundleVersionId":2468249,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --upgrade pip --no-cache-dir\n!pip install ultralytics --no-cache-dir\n\nimport os\nimport pandas as pd\nimport json\nimport shutil\nimport yaml\nfrom ultralytics import YOLO\nfrom sklearn.model_selection import train_test_split\n\n# 配置参数 - 修改为Kaggle路径\nBASE_DATA_DIR = '/kaggle/input/cowboyoutfits'  # 输入数据集根目录\nTEST_CSV_PATH = os.path.join(BASE_DATA_DIR, 'test.csv')  # 测试集CSV\nIMAGE_SRC_DIR = os.path.join(BASE_DATA_DIR, 'images')  # 图片目录\nOUTPUT_DIR = '/kaggle/working'  # 输出目录\n\n# 目标检测类别ID映射\nCATEGORY_IDS = [87, 1034, 131, 318, 588]  # belt, sunglasses, boot, cowboy_hat, jacket\n# 新增：类别名称到ID的映射字典（用于生成JSON）\nCATEGORY_NAME_TO_ID = {\n    'belt': 87,\n    'sunglasses': 1034,\n    'boot': 131,\n    'cowboy_hat': 318,\n    'jacket': 588\n}\nclass_names = ['belt', 'sunglasses', 'boot', 'cowboy_hat', 'jacket']  # 类别名称列表\n\n\n# 数据准备函数\ndef prepare_data():\n    # 加载训练标注\n    train_json_path = os.path.join(BASE_DATA_DIR, 'train.json')\n    with open(train_json_path, 'r') as f:\n        train_info = json.load(f)\n    temp_an = train_info[\"annotations\"]\n    all_images_info = train_info[\"images\"]\n    # 统计信息计算\n    total_annotations = len(temp_an)\n    totle_id = set(each[\"image_id\"] for each in temp_an)\n    print(f\"[统计] 总标注数量: {total_annotations}\")\n    print(f\"[统计] 包含目标类别的图片数: {len(totle_id)}\")\n    # 按类别统计图片数\n    category = {i: [] for i in CATEGORY_IDS}\n    for ann in temp_an:\n        for cid in CATEGORY_IDS:\n            if ann[\"category_id\"] == cid:\n                category[cid].append(ann[\"image_id\"])\n                break\n    for cid in CATEGORY_IDS:\n        print(f\"[统计] 类别 {cid} 的样本数: {len(category[cid])}\")\n    # 找出有标注的图片文件名\n    annotated_image_names = {img['file_name'] for img in all_images_info if img['id'] in totle_id}\n    # 划分训练集和验证集\n    train_image_names, valid_image_names = train_test_split(\n        list(annotated_image_names), test_size=0.2, random_state=42\n    )\n\n    # 复制图片函数\n    def copy_images(img_names, source_dir, target_dir):\n        os.makedirs(target_dir, exist_ok=True)\n        copied = 0\n        missing = 0\n        for img_name in img_names:\n            src = os.path.join(source_dir, img_name)\n            dst = os.path.join(target_dir, img_name)\n            if os.path.exists(src):\n                shutil.copy(src, dst)\n                copied += 1\n            else:\n                print(f\"[警告] 图片不存在: {src}\")\n                missing += 1\n        print(f\"[复制统计] 成功: {copied}, 缺失: {missing}\")\n        return copied\n\n    # 复制训练集和验证集图片\n    train_copied = copy_images(\n        train_image_names, IMAGE_SRC_DIR, os.path.join(OUTPUT_DIR, \"data/train/images\")\n    )\n    val_copied = copy_images(\n        valid_image_names, IMAGE_SRC_DIR, os.path.join(OUTPUT_DIR, \"data/val/images\")\n    )\n    print(f\"[划分] 训练集图片数: {train_copied}, 验证集图片数: {val_copied}\")\n\n    # 生成标签函数\n    def generate_labels(annotations, img_names, img_info_list, label_dir):\n        os.makedirs(label_dir, exist_ok=True)\n        trans = {cid: idx for idx, cid in enumerate(CATEGORY_IDS)}\n        name_to_info = {img['file_name']: img for img in img_info_list}\n        generated_count = 0\n        for img_name in img_names:\n            if img_name in name_to_info:\n                img_info = name_to_info[img_name]\n                img_id = img_info['id']\n                width = img_info['width']\n                height = img_info['height']\n                label_path = os.path.join(label_dir, f\"{os.path.splitext(img_name)[0]}.txt\")\n                with open(label_path, 'w') as f:\n                    for ann in annotations:\n                        if ann['image_id'] == img_id and ann['category_id'] in CATEGORY_IDS:\n                            yolo_class_id = trans[ann['category_id']]\n                            x_center = (ann['bbox'][0] + ann['bbox'][2] / 2) / width\n                            y_center = (ann['bbox'][1] + ann['bbox'][3] / 2) / height\n                            w = ann['bbox'][2] / width\n                            h = ann['bbox'][3] / height\n                            f.write(f\"{yolo_class_id} {x_center:.4f} {y_center:.4f} {w:.4f} {h:.4f}\\n\")\n                generated_count += 1\n        print(f\"[标签生成] 生成了 {generated_count} 个标签文件\")\n\n    # 生成训练集和验证集标签\n    generate_labels(temp_an, train_image_names, all_images_info,\n                    os.path.join(OUTPUT_DIR, 'data/train/labels'))\n    generate_labels(temp_an, valid_image_names, all_images_info,\n                    os.path.join(OUTPUT_DIR, 'data/val/labels'))\n    return len(train_image_names), len(valid_image_names)\n\n\n# 主程序\nif __name__ == '__main__':\n    # 数据准备与模型训练（保留原有逻辑）\n    train_count, val_count = prepare_data()\n    # 生成data.yaml\n    data_yaml = {\n        'train': os.path.abspath(os.path.join(OUTPUT_DIR, 'data/train/images')),\n        'val': os.path.abspath(os.path.join(OUTPUT_DIR, 'data/val/images')),\n        'nc': len(CATEGORY_IDS),\n        'names': class_names\n    }\n    with open(os.path.join(OUTPUT_DIR, 'data.yaml'), 'w') as f:\n        yaml.dump(data_yaml, f, default_flow_style=False)\n    print(f\"[YAML] 数据集配置文件已生成\")\n    # 模型训练\n    os.environ[\"YOLO_DOWNLOAD_FONT\"] = \"0\"\n    model = YOLO('yolov8n.yaml').load('yolov8n.pt')  # 需替换为实际训练好的模型\n    model.train(\n        data=os.path.join(OUTPUT_DIR, 'data.yaml'),\n        epochs=10, imgsz=768, batch=12, lr0=0.004, lrf=0.01, device=0\n    )\n\n    # 重点：测试集处理部分（生成JSON）\n    if os.path.exists(TEST_CSV_PATH):\n        test_df = pd.read_csv(TEST_CSV_PATH)\n        test_images = [os.path.join(IMAGE_SRC_DIR, fname) for fname in test_df['file_name']]\n        test_images_dir = os.path.join(OUTPUT_DIR, 'data/test/images')\n        os.makedirs(test_images_dir, exist_ok=True)\n\n        # 复制测试集图片\n        for src in test_images:\n            if os.path.exists(src):\n                shutil.copy(src, test_images_dir)\n\n        # 模型预测\n        results = model.predict(source=test_images_dir, save=False, save_txt=False)\n\n        # 生成JSON结果\n        kaggle_results = []\n        # 建立文件名到image_id的映射\n        filename_to_id = {row['file_name']: int(row['id']) for _, row in test_df.iterrows()}\n\n        for result in results:\n            img_path = result.path\n            img_name = os.path.basename(img_path)\n            image_id = filename_to_id.get(img_name)\n            if not image_id:\n                continue  # 跳过无image_id的图片\n\n            boxes = result.boxes\n            if boxes is None or len(boxes) == 0:\n                continue  # 跳过无检测结果的图片\n\n            # 处理每个检测框\n            for cls, conf, box in zip(boxes.cls, boxes.conf, boxes.xyxy):\n                class_idx = int(cls)\n                if 0 <= class_idx < len(class_names):\n                    category_name = class_names[class_idx]\n                    category_id = CATEGORY_NAME_TO_ID[category_name]\n                    confidence = float(conf)\n\n                    # 转换边界框格式：[x1, y1, x2, y2] -> [x1, y1, width, height]\n                    x1, y1, x2, y2 = box.tolist()\n                    bbox = [x1, y1, x2 - x1, y2 - y1]\n\n                    # 构造JSON记录（保留6位小数）\n                    prediction = {\n                        \"image_id\": image_id,\n                        \"category_id\": category_id,\n                        \"bbox\": [float(f\"{coord:.6f}\") for coord in bbox],\n                        \"score\": float(f\"{confidence:.6f}\")\n                    }\n                    kaggle_results.append(prediction)\n\n        # 保存JSON文件\n        json_path = os.path.join(OUTPUT_DIR, 'predictions.json')\n        with open(json_path, 'w') as f:\n            json.dump(kaggle_results, f)\n\n        print(f\"[Kaggle] 预测结果已保存至: {json_path}\")\n        print(f\"[统计] 共生成 {len(kaggle_results)} 条预测记录\")\n    else:\n        print(f\"[警告] 测试集CSV不存在: {TEST_CSV_PATH}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}