✅ YOLO获取COCO指标(1):YOLO 转 COCO 标签格式 (txt→json) | 发论文必看!| Ultralytics | 小白友好教程
文章目录
YOLO到COCO格式转换工具(修复版)
这个脚本用于将YOLO格式的标注文件转换为COCO格式的JSON文件,解决了原始转换工具中的一些问题。
1. 主要特点
- 类别ID从1开始:修复了原始脚本中类别ID从0开始的问题,现在类别ID从1开始,符合COCO格式规范。
- 图像ID使用文件名:使用图像文件名(不含扩展名)作为image_id,而不是数字索引。
- 支持子集转换:可以通过指定包含图像路径列表的文本文件(如val.txt)来只转换数据集的一个子集。
- 改进的类别处理:确保不同的类别(如Spur和Spurious_copper)有不同的类别ID。
2. 依赖项
- Python 3.6+
- PIL/Pillow (
pip install pillow
) - tqdm (
pip install tqdm
)
3. 使用方法
3.1 基本用法
python yolo2coco_fixed.py --dataset_path . --output instances_val2017.json
3.2 使用子集文件
python yolo2coco_fixed.py --dataset_path . --output instances_val2017.json --subset val.txt
4. 参数说明
--dataset_path
: YOLO数据集的根目录,默认为当前目录(.
)--output
: 输出的COCO格式JSON文件路径,默认为instances_val2017.json
--subset
: 可选,指定包含要处理的图像列表的文本文件路径,默认为val.txt
5. 数据集目录结构
--dataset_path
指定的目录应包含以下文件和子目录:
dataset_path/
├── classes.txt # 类别文件,包含所有类别名称,每行一个
├── images/ # 图像目录
│ ├── Missing_hole/ # 按类别组织的图像子目录(可选)
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ │ └── ...
│ ├── Mouse_bite/
│ │ └── ...
│ └── ... # 其他类别子目录
├── labels/ # 标签目录
│ ├── Missing_hole/ # 按类别组织的标签子目录(可选)
│ │ ├── image1.txt
│ │ ├── image2.txt
│ │ └── ...
│ ├── Mouse_bite/
│ │ └── ...
│ └── ... # 其他类别子目录
└── val.txt # 可选,包含要处理的图像路径列表
如果图像和标签不是按类别组织在子目录中,也可以直接放在 images/
和 labels/
目录下:
dataset_path/
├── classes.txt # 类别文件
├── images/ # 图像目录
│ ├── image1.jpg
│ ├── image2.jpg
│ └── ...
├── labels/ # 标签目录
│ ├── image1.txt
│ ├── image2.txt
│ └── ...
└── val.txt # 可选,包含要处理的图像路径列表
6. 子集文件格式
子集文件(如val.txt)应包含要处理的图像的完整路径,每行一个路径,例如:
E:\project\YOLOv12\dataset\PCB_DATASET\images\Missing_hole\08_missing_hole_08.jpg
E:\project\YOLOv12\dataset\PCB_DATASET\images\Missing_hole\07_missing_hole_05.jpg
...
7. 注意事项
- 脚本会自动查找类别文件(classes.txt),首先在数据集根目录查找,然后在labels目录及其子目录中查找。
- 对于每个图像,脚本会尝试根据图像路径推断对应的标签文件路径。
- 如果提供了子集文件,脚本只会处理子集文件中列出的图像。
- 类别ID从1开始,而不是从0开始,这与原始的YOLO格式不同。
8. 输出格式
输出的JSON文件符合COCO数据格式,包含以下主要部分:
images
: 所有图像的信息,包括ID(使用文件名)、文件名、宽度和高度。annotations
: 所有标注信息,包括ID、图像ID(使用文件名)、类别ID(从1开始)、边界框等。categories
: 所有类别信息,包括ID(从1开始)、名称和超类别。
9. 修复的问题
- 原始脚本中,类别ID从0开始,现在从1开始。
- 原始脚本中,图像ID使用数字索引,现在使用文件名。
- 原始脚本可能将不同的类别(如Spur和Spurious_copper)转换为相同的类别ID,现在确保每个类别有唯一的ID。
- 添加了对子集转换的支持,可以只处理指定的图像。
10. 完整脚本代码
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import json
import glob
import cv2
from tqdm import tqdm
from PIL import Image
import argparse
import numpy as np
from datetime import datetime
def yolo2coco(yolo_dataset_path, output_json_path, subset_file=None):
"""
将YOLO格式的数据集转换为COCO格式
Args:
yolo_dataset_path: YOLO数据集的根目录
output_json_path: 输出的COCO格式JSON文件路径
subset_file: 可选,指定包含要处理的图像列表的文本文件路径
"""
# 获取类别信息
classes_file = os.path.join(yolo_dataset_path, "classes.txt")
if not os.path.exists(classes_file):
# 尝试在其他目录中查找classes.txt
classes_file = os.path.join(yolo_dataset_path, "labels", "classes.txt")
if not os.path.exists(classes_file):
for root, dirs, files in os.walk(os.path.join(yolo_dataset_path, "labels")):
if "classes.txt" in files:
classes_file = os.path.join(root, "classes.txt")
break
if not os.path.exists(classes_file):
raise FileNotFoundError(f"未找到classes.txt文件")
# 读取类别信息
with open(classes_file, 'r', encoding='utf-8') as f:
classes = [line.strip() for line in f.readlines() if line.strip()]
# 创建COCO格式的数据结构
coco_data = {
"info": {
"description": "PCB_DATASET",
"url": "",
"version": "1.0",
"year": datetime.now().year,
"contributor": "",
"date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
},
"licenses": [
{
"id": 1,
"name": "Unknown",
"url": ""
}
],
"categories": [],
"images": [],
"annotations": []
}
# 添加类别信息 - 从1开始
for i, cls_name in enumerate(classes):
coco_data["categories"].append({
"id": i + 1, # 修改点1:类别ID从1开始
"name": cls_name,
"supercategory": "PCB"
})
# 创建类别名称到ID的映射
class_name_to_id = {
cls_name: i + 1 for i, cls_name in enumerate(classes)}
# 获取所有图像文件
image_files = []
# 如果提供了子集文件,则从中读取图像列表
if subset_file and os.path.exists(subset_file):
print(f"使用子集文件: {
subset_file}")
with open(subset_file, 'r', encoding='utf-8') as f:
file_paths = [line.strip() for line in f.readlines() if line.strip()]
# 确保所有路径都存在
image_files = [path for path in file_paths if os.path.exists(path)]
if len(image_files) != len(file_paths):
print(f"警告: 子集文件中的 {
len(file_paths) - len(image_files)} 个图像路径不存在")
else:
# 如果没有提供子集文件,则按原来的方式获取所有图像
image_dir = os.path.join(yolo_dataset_path, "images")
for cls_name in classes:
cls_image_dir = os.path.join(image_dir, cls_name)
if os.path.exists(cls_image_dir):
cls_images = glob.glob(os.path.join(cls_image_dir, "*.jpg")) + \
glob.glob(os.path.join(cls_image_dir, "*.jpeg")) + \
glob.glob(os.path.join(cls_image_dir, "*.png"))
image_files.extend(cls_images)
# 如果没有在类别子目录中找到图像,尝试直接在images目录下查找
if not image_files:
image_files = glob.glob(os.path.join(image_dir, "*.jpg")) + \
glob.glob(os.path.join(image_dir, "*.jpeg")) + \
glob.glob(os.path.join(image_dir, "*.png"))
if not image_files:
raise FileNotFoundError(f"未找到任何图像文件")
print(f"找到 {
len(image_files)} 个图像文件")
# 处理每个图像和对应的标签
ann_id = 0
for img_id, img_path in enumerate(tqdm(image_files, desc="处理图像")):
# 获取图像信息
img_filename = os.path.basename(img_path)
img_name_without_ext = os.path.splitext(img_filename)[0]
# 获取图像尺寸
try:
img = Image.open(img_path)
img_width, img_height = img.size
except Exception as e:
print(f"无法打开图像 {
img_path}: {
e}")
continue
# 添加图像信息到COCO数据 - 使用文件名作为image_id
coco_data["images"].append({
"id": img_name_without_ext, # 修改点2:使用文件名作为image_id
"file_name": img_filename,
"width": img_width,
"height": img_height,
"license": 1,
"date_captured": ""
})
# 确定标签文件路径
# 首先尝试根据图像路径推断标签路径
img_dir = os.path.dirname(img_path)
cls_name = os.path.basename(img_dir)
# 尝试在类别子目录中查找
label_found = False
label_path = os.path.join(yolo_dataset_path, "labels", cls_name, f"{
img_name_without_ext}.txt")
if os.path.exists(label_path):
label_found = True
# 如果在类别子目录中没找到,尝试直接在labels目录下查找
if not label_found:
label_path = os.path.join(yolo_dataset_path, "labels", f"{
img_name_without_ext}.txt")
if not os.path.exists(label_path):
print(f"未找到图像 {
img_filename} 对应的标签文件")
continue
# 读取标签信息
try:
with open(label_path, 'r') as f:
label_lines = [line.strip() for line in f.readlines() if line.strip()]
# 处理每个标签
for line in label_lines:
parts = line.split()
if len(parts) != 5:
print(f"标签格式错误: {
line}")
continue
# 解析YOLO格式的标签
cls_id = int(parts[0])
if cls_id >= len(classes):
print(f"警告: 标签中的类别ID {
cls_id} 超出了类别列表范围 (0-{
len(classes)-1})")
continue
# 获取类别名称和新的类别ID
cls_name = classes[cls_id]
new_cls_id = class_name_to_id[cls_name] # 修改点3:使用从1开始的类别ID
x_center = float(parts[1])
y_center = float(parts[2])
width = float(parts[3])
height = float(parts[4])
# 转换为COCO格式(左上角坐标和宽高)
x = (x_center - width / 2) * img_width
y = (y_center - height / 2) * img_height
w = width * img_width
h = height * img_height
# 添加标注信息到COCO数据
coco_data["annotations"].append({
"id": ann_id,
"image_id": img_name_without_ext, # 修改点4:使用文件名作为image_id
"category_id": new_cls_id, # 修改点5:使用从1开始的类别ID
"bbox": [x, y, w, h],
"area": w * h,
"segmentation": [],
"iscrowd": 0
})
ann_id += 1
except Exception as e:
print(f"处理标签文件 {
label_path} 时出错: {
e}")
# 保存COCO格式的数据到JSON文件
with open(output_json_path, 'w', encoding='utf-8') as f:
json.dump(coco_data, f, ensure_ascii=False, indent=2)
print(f"转换完成,共处理 {
len(coco_data['images'])} 张图像,{
len(coco_data['annotations'])} 个标注")
print(f"COCO格式数据已保存到: {
output_json_path}")
def main():
parser = argparse.ArgumentParser(description="将YOLO格式的数据集转换为COCO格式")
parser.add_argument("--dataset_path", type=str, default=".", help="YOLO数据集的根目录")
parser.add_argument("--output", type=str, default="instances_val2017.json", help="输出的COCO格式JSON文件路径")
parser.add_argument("--subset", type=str, default="val.txt", help="可选,指定包含要处理的图像列表的文本文件路径")
args = parser.parse_args()
yolo2coco(args.dataset_path, args.output, args.subset)
if __name__ == "__main__":
main()