【/108533949/】coco制作程序

返回原博客 

import os
import json
import shutil
import numpy as np
import cv2



# 读图片（支持中文路径）->ndarray
def myImread(path):
    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
class MyData2COCO:

    def __init__(self):  # 初始化
        self.images = []  # 存储images键对应的数据
        self.annotations = []  # 存储annotations键对应的数据
        self.categories = []  # 存储categories键对应的数据
        self.img_id = 0  # 统计image的id
        self.ann_id = 0  # 统计annotation的id

    def _categories(self, num_categories):  # 获取categories信息
        for i in range(0, num_categories):
            category = {}
            category['id'] = i
            category['name'] = str(i)  # 可根据实际需要修改
            category['supercategory'] = 'name'  # 可根据实际需要修改
            self.categories.append(category)

    def _image(self, path, h, w):  # 获取images信息
        image = {}
        image['height'] = h
        image['width'] = w
        image['id'] = self.img_id
        image['file_name'] = os.path.basename(path)
        return image

    def _annotation(self, label, bbox):  # 获取annotations信息
        # bbox = list(bbox.values())

        area = bbox[2] * bbox[3]
        points = [[bbox[0], bbox[1]], [bbox[0] + bbox[2], bbox[1]], [bbox[2], bbox[1] + bbox[3]],
                  [bbox[0], bbox[1] + bbox[3]]]
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = label
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['bbox'] = bbox
        annotation['iscrowd'] = 0
        annotation['area'] = area
        return annotation
    # 根据自己的数据集调整
    def to_coco(self, folder_label, folder_image, num_categories=1):  # 转换实现函数
        """
            anno_file: 自己数据的文件路径
            img_dir: 图片文件夹路径（coco分为train和valid）
            num_categories: 对应的总类别数目
        其他：list_bbox[[x1,y1,w,h],....]原数值
             labels[0,1,0,0,... ...] 标签从0开始的数字
        """
     
        self._categories(num_categories)  # 初始化categories基本信息
        ...  ...
            for bbox, label in zip(list_bbox, labels):
                annotation = self._annotation(label, bbox)
                self.annotations.append(annotation)
                self.ann_id += 1

            img_path = os.path.join(folder_image, img_name)

            img = myImread(img_path)
            h, w, c = img.shape
            self.images.append(self._image(img_path, h, w))

            self.img_id += 1
        instance = {}
        instance['info'] = 'WiderPerson_coco'
        instance['license'] = ['none']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    def save_coco_json(self, instance, save_path):
        with open(save_path, 'w') as fp:
            json.dump(instance, fp, indent=1, separators=(',', ': '))
        # 保存文件


if __name__ == '__main__':

    folder_label = r"XXXXXXXXXXX"
    folder_image = r"XXXXXXXXXXXXXXXXXXXXX"
    save_path = r"XXXXXXXXXXXXXXXXXXX"

    if os.path.exists(save_path):
        shutil.rmtree(save_path)
    os.makedirs(save_path)

    fabric2coco_train = MyData2COCO()
    train_instance = fabric2coco_train.to_coco(folder_label, folder_image)
    save_path_name = os.path.join(save_path, 'XXX_coco_trainval.json')
    fabric2coco_train.save_coco_json(train_instance, save_path_name)