COCO数据集转YOLO数据集

一、COCO数据格式

# 转换前的数据格式
data_coco/
├── images/
│   ├── train2017/  # 训练集图片
│   └── val2017/    # 验证集图片
└── annotations/    # 原始COCO JSON文件

二、代码

import json
import os
import argparse


class COCO2YOLO:
    def __init__(self):
        self._check_file_and_dir(json_file, output)
        self.labels = json.load(open(json_file, 'r', encoding='utf-8'))
        self.coco_id_name_map = self._categories()
        self.coco_name_list = list(self.coco_id_name_map.values())
        print("total images", len(self.labels['images']))
        print("total categories", len(self.labels['categories']))
        print("total labels", len(self.labels['annotations']))

    @staticmethod
    def _check_file_and_dir(file_path, dir_path):
        """
        文件检测和文件夹检测
        :param file_path:
        :param dir_path:
        :return:
        """
        if not os.path.exists(file_path):
            # 文件路径错误
            raise ValueError("file not found")
        if not os.path.exists(dir_path):
            # 文件夹不存在就创建
            os.makedirs(dir_path)

    def _categories(self):
        # 获取标注文件中的种类
        categories = {}
        for cls in self.labels['categories']:
            categories[cls['id']] = cls['name']
        return categories

    def _load_images_info(self):
        # 获取图片基本信息：文件名、宽、高
        images_info = {}
        for image in self.labels['images']:
            img_id = image['id']
            file_name = image['file_name']
            if file_name.find('\\') > -1:
                file_name = file_name[file_name.index('\\') + 1:]
            w = image['width']
            h = image['height']
            images_info[img_id] = (file_name, w, h)

        return images_info

    @staticmethod
    def _bbox_2_yolo(bbox, img_w, img_h):
        """
        将coco格式的边界框（Bounding Box）坐标转换为Yolo格式的归一化坐标
        :param bbox: COCO格式的边界框，格式为[x_min, y_min, width, height]
        :param img_w: 图片宽度
        :param img_h: 图片的高度
        :return: 归一化后的中心点x、y坐标、边界框的宽、高
        """
        x, y, w, h = bbox[0], bbox[1], bbox[2], bbox[3]
        # 边界框中心点坐标
        center_x = bbox[0] + w / 2
        center_y = bbox[1] + h / 2
        # 归一化因子
        dw = 1 / img_w
        dh = 1 / img_h
        # 归一化
        center_x *= dw
        w *= dw
        center_y *= dh
        h *= dh
        return center_x, center_y, w, h

    def _convert_anno(self, images_info):
        # 格式转换
        anno_dict = dict()
        for anno in self.labels['annotations']:
            bbox = anno['bbox']
            image_id = anno['image_id']
            category_id = anno['category_id']

            image_info = images_info.get(image_id)
            image_name = image_info[0]
            img_w = image_info[1]
            img_h = image_info[2]
            yolo_box = self._bbox_2_yolo(bbox, img_w, img_h)

            anno_info = (image_name, category_id, yolo_box)
            anno_infos = anno_dict.get(image_id)
            if not anno_infos:
                anno_dict[image_id] = [anno_info]
            else:
                anno_infos.append(anno_info)
                anno_dict[image_id] = anno_infos
        return anno_dict

    def coco2yolo(self):
        print("loading image info...")
        images_info = self._load_images_info()
        print("loading done, total images", len(images_info))

        print("start converting...")
        anno_dict = self._convert_anno(images_info)
        print("converting done, total labels", len(anno_dict))

        print("saving txt file...")
        self._save_txt(anno_dict)
        print("saving done")

    def _save_txt(self, anno_dict):
        for k, v in anno_dict.items():
            file_name = os.path.splitext(v[0][0])[0] + ".txt"
            with open(os.path.join(output, file_name), 'w', encoding='utf-8') as f:
                print(k, v)
                for obj in v:
                    cat_name = self.coco_id_name_map.get(obj[1])
                    category_id = self.coco_name_list.index(cat_name)
                    box = ['{:.6f}'.format(x) for x in obj[2]]
                    box = ' '.join(box)
                    line = str(category_id) + ' ' + box
                    f.write(line + '\n')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Test yolo data.')
    parser.add_argument('-i', help='JSON file', dest='json',
                        default='/home/data_coco/annotations/meter_coco_val.json')
    parser.add_argument('-o', help='path to output folder', dest='out',
                        default='/home/data_coco/labels/val2017')

    args = parser.parse_args()

    json_file = args.json
    output = args.out

    COCO2YOLO().coco2yolo()

三、YOLO数据格式

# 转换后的数据格式
data_coco/
├── images/
│   ├── train2017/  # 训练集图片
│   └── val2017/    # 验证集图片
├── labels/
│   ├── train2017/  # 训练集标签（YOLO格式）
│   └── val2017/    # 验证集标签（YOLO格式）
└── annotations/    # 原始COCO JSON文件