YOLOtxt格式转CoCo数据集格式

在分割任务对比试验需要用Mask RCNN作为对比，需要将yolotxt格式转化coco数据集格式

网上很多版本，多少有点问题，为了方便各位查看和修改我做出了注释和修改意见。

以下是基于网络代码的修改版本，不会有多的报错：

import os
import json
from PIL import Image

# 你的路径定义
coco_format_save_path = 'D:\\Code\\PycharmProjects\\yolov5-seg-master\\dataset\\custom_dataset\\'
yolo_format_annotation_path = 'D:\\Code\\PycharmProjects\\yolov5-seg-master\\dataset\\custom_dataset\\labels\\train\\'
img_pathDir = 'D:\\Code\\PycharmProjects\\yolov5-seg-master\\dataset\\custom_dataset\\images\\train\\'

# 类别映射和其他初始化代码  该代码相对于其他版本用户可以自定义在以下修改类别而不需要额外调用外部文件
categories_mapping = ['square', 'triangle']
categories = [{'id': i + 1, 'name': label, 'supercategory': 'None'} for i, label in enumerate(categories_mapping)]

write_json_context = {
    'info': {'description': '', 'url': '', 'version': '', 'year': 2024, 'contributor': '', 'date_created': '2024-02-16'},
    'licenses': [{'id': 1, 'name': None, 'url': None}],
    'categories': categories,
    'images': [],
    'annotations': []
}

imageFileList = os.listdir(img_pathDir)
for i, imageFile in enumerate(imageFileList):
    imagePath = os.path.join(img_pathDir, imageFile)
    image = Image.open(imagePath)
    W, H = image.size

    img_context = {
        'file_name': imageFile, 'height': H, 'width': W,
        'date_captured': '2021-07-25', 'id': i,
        'license': 1, 'color_url': '', 'flickr_url': ''
    }
    write_json_context['images'].append(img_context)

    txtFile = os.path.splitext(imageFile)[0] + '.txt'  # 修改以正确处理文件名 获取该图片获取的txt文件  # 和其他人写的代码区别是可以保证文件被找到
    with open(os.path.join(yolo_format_annotation_path, txtFile), 'r') as fr:
        lines = fr.readlines()  # 读取txt文件的每一行数据，lines是一个列表，包含了一个图片的所有标注信息

    # 重新引入循环中的enumerate函数
    for j, line in enumerate(lines):  # 这里使用enumerate确保j被正确定义
        parts = line.strip().split(' ')
        if len(parts) >= 5:  # 确保至少有5个部分    # 这里需要注意，yolo格式添加额外的内容容易报错，所以需要你只要前面的主要信息
            class_id, x, y, w, h = map(float, parts[:5])  # 只读取前五个值
            xmin = (x - w / 2) * W  # 坐标转换
            ymin = (y - h / 2) * H
            xmax = (x + w / 2) * W
            ymax = (y + h / 2) * H
            bbox_width, bbox_height = w * W, h * H

            bbox_dict = {
                'id': i * 10000 + j,  # 使用j，它现在被enumerate定义
                'image_id': i,
                'category_id': class_id + 1,  # 注意目标类别要加一
                'iscrowd': 0,
                'area': bbox_width * bbox_height,
                'bbox': [xmin, ymin, bbox_width, bbox_height],
                'segmentation': [[xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax]]
            }
            write_json_context['annotations'].append(bbox_dict)
name = os.path.join(coco_format_save_path, "train.json")
with open(name, 'w') as fw:
    json.dump(write_json_context, fw, indent=2)