【数据集映射】（含完整可行Python代码）Yolo格式数据（txt）转RCNN格式数据（coco）

1：需求背景

目标检测研究中，通常不同的模型需要不同的数据格式，如yolo需要txt、rcnn需要coco等，因此就需要对标注的数据格式进行转换。

通常的数据格式有txt、coco(json)、Pascal VOC、xml等。

本文的需求是txt转coco，下面提供一种可行的代码（可完全复制）

2：格式转换代码

需要修改的内容是：

1：图像和txt数据集的路径。这个是自己配制的，一般在dataset里面。将图像路径赋给images_dir，txt路径赋给annotations_dir即可。

# 路径配置（根据自己的路径修改）
images_dir = ""
annotations_dir = ""

2：categories内的映射字典。详情可以打开自己的txt文件，查看每一类对应的class_number，例如我们这里的yolo格式中的GD图像分类编号对应0，则映射字典出设置为【0: "xxx",】，其他内容同理。

# 类别映射字典（根据自己的内容修改）
categories = {
    0: "xxx",
}

# 后面可以加1、2、3...

3：数据集图像的大小参数。在以下位置处的【image_width】和【image_height】替换为自己图像的分辨率。

# TODO: 如果每张图片的大小都一样，你可以在这里指定
        # 如果大小不一样，你需要读取图片文件来获取实际尺寸
        image_width = 960
        image_height = 960

4：json文件的保存位置。在以下位置处的路径【'xxx.json'】修改为自己的想要存放的绝对路径/相对路径即可。

# 将COCO数据结构写入JSON文件（保存的json路径自己修改）
with open('xxx.json', 'w') as json_file:
    json.dump(coco_output, json_file, indent=4)

完整代码如下： 

import json
import os

# 路径配置（根据自己的路径修改）
images_dir = '/home/ubuntu/datasets'
annotations_dir = '/home/ubuntu/datasets'

# 类别映射字典（根据自己的内容修改）
categories = {
    0: "xxx",
}

# COCO格式的基本结构
coco_output = {
    "info": {},
    "licenses": [],
    "images": [],
    "annotations": [],
    "categories": [{"id": k, "name": v, "supercategory": ""} for k, v in categories.items()]
}

# 图片和标注的ID计数器
image_id = 1
annotation_id = 1

# 遍历annotations目录下的所有TXT文件
for filename in os.listdir(annotations_dir):
    if filename.endswith('.txt'):
        # 假设文件名与图片文件名一致（不包含扩展名）
        image_filename = filename.replace('.txt', '.jpg')
        image_path = os.path.join(images_dir, image_filename)
        
        # TODO: 如果每张图片的大小都一样，你可以在这里指定
        # 如果大小不一样，你需要读取图片文件来获取实际尺寸
        image_width = 960
        image_height = 960
        
        # 添加图片信息到COCO数据结构
        coco_output['images'].append({
            "id": image_id,
            "file_name": image_filename,
            "width": image_width,
            "height": image_height
        })
        
        # 读取每个TXT文件并添加标注信息
        txt_file_path = os.path.join(annotations_dir, filename)
        with open(txt_file_path, 'r') as file:
            for line in file:
                class_id, x_center, y_center, width, height = map(float, line.strip().split())
                
                # COCO要求bbox是[x_min, y_min, width, height]，而不是中心点坐标
                x_min = (x_center - width / 2) * image_width
                y_min = (y_center - height / 2) * image_height
                width = width * image_width
                height = height * image_height
                
                # 添加标注信息
                coco_output['annotations'].append({
                    "id": annotation_id,
                    "image_id": image_id,
                    "category_id": class_id,
                    "bbox": [x_min, y_min, width, height],
                    "area": width * height,
                    "segmentation": [],  # 如果你有分割信息可以在这里添加
                    "iscrowd": 0
                })
                
                annotation_id += 1
        
        # 更新图片ID
        image_id += 1

# 将COCO数据结构写入JSON文件（保存的json路径自己修改）
with open('xxx.json', 'w') as json_file:
    json.dump(coco_output, json_file, indent=4)

print("Conversion completed!")