VOC转COCO格式数据集

VOC转COCO格式数据集

此代码用于将VOC格式转换为COCO格式的数据集
在main里选择你的 txt 来筛选图片（train / txt / val）

import os.path as osp
import xml.etree.ElementTree as ET
import mmcv
from glob import glob
from tqdm import tqdm
from PIL import Image

def coco_classes():#在这里更改你的类别
    name_dict = { '1': 'pedestrian', '2': 'people',
             '3': 'bicycle', '4': 'car', '5': 'van', '6': 'truck',
             '7': 'tricycle', '8': 'awning-tricycle', '9': 'bus',
             '10': 'motor', '11': 'others'}
    return name_dict.values()

label_ids = {name: i for i, name in enumerate(coco_classes())}

def parse_xml(xml_path, img_id, anno_id):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    annotation = []
    for obj in root.findall('object'):
        name = obj.find('name').text
        category_id = label_ids[name]
        bnd_box = obj.find('bndbox')
        xmin = int(bnd_box.find('xmin').text)
        ymin = int(bnd_box.find('ymin').text)
        xmax = int(bnd_box.find('xmax').text)
        ymax = int(bnd_box.find('ymax').text)
        w = xmax - xmin + 1
        h = ymax - ymin + 1
        area = w*h
        annotation.append({
                        "area": area,
                        "iscrowd": 0,
                        "image_id": img_id,
                        "bbox": [xmin, ymin, w, h],
                        "category_id": category_id,
                        "id": anno_id,
                        "ignore": 0})
        anno_id += 1
    return annotation, anno_id


def cvt_annotations(img_path, xml_path, out_file, path):
    images = []
    annotations = []
    img_id = 1
    anno_id = 1
    for img_path in tqdm(glob(img_path + '/*.jpg')):
        if img_path.split("/")[-1].split(".")[0] in path:
            w, h = Image.open(img_path).size
            img_name = osp.basename(img_path)
            img = {"file_name": img_name, "height": int(h), "width": int(w), "id": img_id}
            images.append(img)

            xml_file_name = img_name.split('.')[0] + '.xml'
            xml_file_path = osp.join(xml_path, xml_file_name)
            annos, anno_id = parse_xml(xml_file_path, img_id, anno_id)
            annotations.extend(annos)
            img_id += 1

    categories = []
    for k,v in label_ids.items():
        if (k == 'others') or (k =='ignored regions'):
            continue
        categories.append({"name": k, "id": v})
    final_result = {"images": images, "annotations": annotations, "categories": categories}
    mmcv.dump(final_result, out_file)
    return annotations


def main():
    # train = "/home/zsy/workspace/UAV/datauav/source/train.txt"
    # val = "/home/zsy/workspace/UAV/datauav/source/val.txt"
    # test = "/home/zsy/workspace/UAV/datauav/source/test.txt"
    # with open(train) as t:
    #     t = t.read().splitlines()
    # with open(val) as v:
    #     v = v.read().splitlines()
    # with open(test) as te:
    #     te = te.read().splitlines()
    # train = [i + ".xml" for i in t]
    # val = [i + ".xml" for i in v]
    # test = [i + ".xml" for i in te]
    xml_path = '/media/zsy/FEC0939E1BEF7D06/xml_val/'
    img_path = '/media/zsy/FEC0939E1BEF7D06/images/'
    print('processing {} ...'.format("xml format annotations"))
    cvt_annotations(img_path, xml_path, '/media/zsy/val.json',val)
    print('Done!')


if __name__ == '__main__':
    main()