COCO数据集提取特定的类，从原有的COCO数据集提取segmenttation，area，bbox的xml格式以及xml转json代码。转换后与原coco的json数据一致

提取coco数据集的数据,##有个缺陷就是当多次提取不同类别的时候，train，val的json文件。类别id可能会不同，请自行手动修改.

提取特定的类：

# -*-coding:utf-8-*-
from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import skimage.io as io
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw

# 需要设置的路径
savepath = "F:/coco-lei-fenbu/coco-donut/"  # "F:/coco-fan/"
img_dir = savepath + 'images/'
anno_dir = savepath + 'annotations/'
datasets_list = ['train2017']

# coco有80类，这里写要提取类的名字，以person为例
classes_names = ['donut']

# 包含所有类别的原coco数据集路径
'''
目录格式如下：
$COCO_PATH
----|annotations
----|train2017
----|val2017
----|test2017
'''
dataDir = 'F:/coco/'

headstr = """\
<annotation>
    <folder>VOC</folder>
    <filename>%s</filename>
    <source>
        <database>My Database</database>
        <annotation>COCO</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>

"""
objstr = """\
    <object>
        <name>%s</name>
        <segmentation>%s</segmentation>
        <area>%s</area>
        <bndbox>
            <cx>%s</cx>
            <cy>%s</cy>
            <w>%s</w>
            <h>%s</h>
        </bndbox>
    </object>
"""
tailstr = '''\
</annotation>
'''


# 检查目录是否存在，如果存在，先删除再创建，否则，直接创建
def mkr(path):
    if not os.path.exists(path):
        os.makedirs(path)  # 可以创建多级目录


def id2name(coco):
    classes = dict()
    for cls in coco.dataset['categories']:
        classes[cls['id']] = cls['name']
    return classes


def write_xml(anno_path, head, objs, tail):
    f = open(anno_path, "w")
    f.write(head)
    for obj in objs:
        f.write(objstr % (obj[0], obj[1], obj[2], obj[3], obj[4], obj[5], obj[6]))
    f.write(tail)


def save_annotations_and_imgs(coco, dataset, filename, objs):
    # 将图片转为xml，例:COCO_train2017_000000196610.jpg-->COCO_train2017_000000196610.xml
    dst_anno_dir = os.path.join(anno_dir, dataset)
    mkr(dst_anno_dir)
    anno_path = dst_anno_dir + '/' + filename[:-3] + 'xml'
    img_path = dataDir + dataset + '/' + filename
    print("img_path: ", img_path)
    dst_img_dir = os.path.join(img_dir, dataset)
    mkr(dst_img_dir)
    dst_imgpath = dst_img_dir + '/' + filename
    print("dst_imgpath: ", dst_imgpath)
    img = cv2.imread(img_path)
    # if (img.shape[2] == 1):
    #    print(filename + " not a RGB image")
    #   return
    shutil.copy(img_path, dst_imgpath)

    head = headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
    tail = tailstr
    write_xml(anno_path, head, objs, tail)


def showimg(coco, dataset, img, classes, cls_id, show=True):
    global dataDir
    I = Image.open('%s/%s/%s' % (dataDir, dataset, img['file_name']))
    # 通过id，得到注释的信息
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
    # print(annIds)
    anns = coco.loadAnns(annIds)
    # print(anns)
    # coco.showAnns(anns)
    objs = []
    for ann in anns:
        class_name = classes[ann['category_id']]
        if class_name in classes_names:
            print(class_name)
            if 'segmentation' in ann:
                segmentation = ann["segmentation"]
            if 'area' in ann:
                area = ann["area"]
            if 'bbox' in ann:
                bbox = ann['bbox']
                xmin = float(bbox[0])
                ymin = float(bbox[1])
                xmax = float(bbox[2])  # bbox[2] + bbox[0]
                ymax = float(bbox[3])  # bbox[3] + bbox[1]
                obj = [class_name, segmentation, area, xmin, ymin, xmax, ymax]
                objs.append(obj)
                draw = ImageDraw.Draw(I)
                draw.rectangle([xmin, ymin, xmax, ymax])
    if show:
        plt.figure()
        plt.axis('off')
        plt.imshow(I)
        plt.show()

    return objs


for dataset in datasets_list:
    # ./COCO/annotations/instances_train2017.json
    annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataset)

    # 使用COCO API用来初始化注释数据
    coco = COCO(annFile)

    # 获取COCO数据集中的所有类别
    classes = id2name(coco)
    print(classes)
    # [1, 2, 3, 4, 6, 8]
    classes_ids = coco.getCatIds(catNms=classes_names)
    print(classes_ids)
    for cls in classes_names:
        # 获取该类的id
        cls_id = coco.getCatIds(catNms=[cls])
        img_ids = coco.getImgIds(catIds=cls_id)
        print(cls, len(img_ids))
        # imgIds=img_ids[0:10]
        for imgId in tqdm(img_ids):
            img = coco.loadImgs(imgId)[0]
            filename = img['file_name']
            # print(filename)
            objs = showimg(coco, dataset, img, classes, classes_ids, show=False)
            print(objs)
            save_annotations_and_imgs(coco, dataset, filename, objs)

xml转json

# -*-coding:utf-8-*-
# -*-coding:utf-8-*-
import xml.etree.ElementTree as ET
import os
import json

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict()
image_set = set()

category_item_id = 0
image_id = 20210000000
annotation_id = 0


def addCatItem(name):
    global category_item_id
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id
    return category_item_id


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox, segmentation, area):
    global annotation_id
    annotation_item = dict()

    annotation_item['segmentation'] = segmentation

    annotation_item['area'] = area
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)


def parseXmlFiles(xml_path):
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue

        bndbox = dict()
        size = dict()

        segmentation = []
        area = []


        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None



        xml_file = os.path.join(xml_path, f)
        print(xml_file)

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == 'folder':
                continue

            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')


            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>


            for subelem in elem:

                bndbox['cx'] = None
                bndbox['cy'] = None
                bndbox['w'] = None
                bndbox['h'] = None

                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':
                    object_name = subelem.text

                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]

                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)

                elif current_parent == 'object' and subelem.tag == 'segmentation':
                    segmentation = eval(subelem.text)

                elif current_parent == 'object' and subelem.tag == 'area':
                    area = float(subelem.text)

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = float(option.text)

                # only after parse the <object> tag
                if bndbox['cx'] is not None:#'xmin'
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['cx'])
                    # y
                    bbox.append(bndbox['cy'])
                    # w
                    bbox.append(bndbox['w'])
                    # h
                    bbox.append(bndbox['h'])
                    print('add annotation with {},{},{},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox,segmentation,area))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox, segmentation, area)


if __name__ == '__main__':
    # 需要自己设定的地址，一个是已生成的是xml文件的父目录；一个是要生成的json文件的目录
    xml_dir = r'F:\coco-lei-fenbu\coco-donut\annotations'
    json_dir = r'F:\coco-lei-fenbu\coco-donut\annotations'
    dataset_lists = ['val2017']
    for dataset in dataset_lists:
        xml_path = os.path.join(xml_dir, dataset)
        json_file = json_dir + '/{}.json'.format(dataset)
        parseXmlFiles(xml_path)
        json.dump(coco, open(json_file, 'w'))