coco数据集格式转yolo数据格式
一、coco数据集是什么？
COCO（Common Objects in Context）是一个广泛使用的目标检测和分割数据集，而YOLO（You Only Look Once）是一种流行的实时目标检测算法。
首先，导入了必要的库，包括json和os。然后，定义了一个名为convert_coco_to_yolo的函数，用于将COCO数据转换为YOLO格式。该函数接受一个COCO数据的字典作为输入，并返回转换后的YOLO数据和类别字典。
在函数内部，提取了COCO数据中的图像列表、注释和类别信息。然后，创建了一个空的YOLO数据列表和一个用于存储类别名和序号映射关系的字典。
接下来，遍历图像列表，并获取图像的ID、文件名、宽度和高度。然后，根据图像ID筛选出与该图像相关的注释，并将它们存储在image_annotations列表中。
然后，遍历image_annotations列表，并提取注释的类别ID。通过在类别列表中查找与类别ID匹配的类别，可以获取类别的名称。如果找不到匹配的类别，则跳过该注释。
接下来，提取注释的边界框信息，并计算出边界框的中心坐标、归一化后的坐标和大小。将这些信息存储在yolo_annotations列表中。
如果存在YOLO注释，则按照类别序号对注释进行排序，并将图像的相关信息和注释存储在yolo_data列表中。
最后，返回转换后的YOLO数据和类别字典。
接下来，定义了几个变量，包括COCO数据文件的路径、文件名和保存目录的路径。然后，使用os.path.join函数构建了完整的文件路径。
通过检查文件是否存在，可以确保文件路径是有效的。如果文件存在，打开文件并加载其中的JSON数据。然后，调用convert_coco_to_yolo函数将COCO数据转换为YOLO数据。
使用os.makedirs函数创建保存目录，如果目录已经存在，则不进行任何操作。
接下来，生成了一个class.txt文件，用于存储类别的名称。遍历类别字典，并按照类别序号对字典进行排序。然后，将类别的名称写入文件中。
最后，遍历YOLO数据列表，并根据图像的文件名创建一个对应的.txt文件。然后，遍历注释列表，并将注释的类别和其他信息写入文件中。
最后，打印转换完成的消息，并显示保存目录的路径。如果文件不存在，则打印文件不存在的消息。
通过运行这个程序，可以将COCO数据集格式转换为YOLO格式，并将转换后的数据保存到指定的目录中。
这篇博客介绍了一个将COCO数据集格式转换为YOLO格式的Python程序。通过这个程序，可以轻松地处理COCO数据，并将其转换为适用于YOLO算法的格式。希望这个程序对你有所帮助！
二、完整代码
import json
import os
def convert_coco_to_yolo(coco_data):
image_list = coco_data['images']
annotations = coco_data['annotations']
categories = coco_data['categories']
yolo_data = []
category_dict = {} # 用于存储类别名和对应的序号
for i, category in enumerate(categories): # 构建类别名和序号的映射关系
category_dict[category['name']] = i
for image in image_list:
image_id = image['id']
file_name = image['file_name']
width = image['width']
height = image['height']
image_annotations = [ann for ann in annotations if ann['image_id'] == image_id]
yolo_annotations = []
for ann in image_annotations:
category_id = ann['category_id']
category = next((cat for cat in categories if cat['id'] == category_id), None)
if category is None:
continue
bbox = ann['bbox']
x, y, w, h = bbox
x_center = x + w / 2
y_center = y + h / 2
normalized_x_center = x_center / width
normalized_y_center = y_center / height
normalized_width = w / width
normalized_height = h / height
yolo_annotations.append({
'category': category_dict[category['name']], # 使用类别序号
'x_center': normalized_x_center,
'y_center': normalized_y_center,
'width': normalized_width,
'height': normalized_height
})
if yolo_annotations:
yolo_annotations.sort(key=lambda x: x['category']) # 按类别序号排序
yolo_data.append({
'file_name': file_name,
'width': width,
'height': height,
'annotations': yolo_annotations
})
return yolo_data, category_dict
path = '/codeyard/yolov5_6.0/data_MS_1001labels' # 修改为包含 via_export_coco.json 文件的目录路径
file_name = 'coco_filtered.json' # 文件名
save_dir = '/codeyard/yolov5_6.0/data_MS_1001labels/labels' # 保存目录
file_path = os.path.join(path, file_name) # 完整文件路径
if os.path.isfile(file_path): # 检查文件是否存在
with open(file_path, 'r', encoding='utf-8') as load_f:
load_dict = json.load(load_f)
yolo_data, category_dict = convert_coco_to_yolo(load_dict)
os.makedirs(save_dir, exist_ok=True) # 创建保存目录
# 生成 class.txt 文件
class_file_path = os.path.join(save_dir, 'classes.txt')
with open(class_file_path, 'w', encoding='utf-8') as class_f:
for category_name, category_index in sorted(category_dict.items(), key=lambda x: x[1]):
class_f.write(f"{category_name}\n")
for data in yolo_data:
file_name = os.path.basename(data['file_name']) # 提取文件名部分
width = data['width']
height = data['height']
annotations = data['annotations']
txt_file_path = os.path.join(save_dir, os.path.splitext(file_name)[0] + '.txt')
with open(txt_file_path, 'w', encoding='utf-8') as save_f:
for annotation in annotations:
category = annotation['category']
x_center = annotation['x_center']
y_center = annotation['y_center']
box_width = annotation['width']
box_height = annotation['height']
line = f"{category} {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f}\n"
save_f.write(line)
print("转换完成，保存到:", save_dir)
else:
print("文件不存在:", file_path)