PASCAL VOC目标检测数据集格式转化为COCO目标检测数据集格式

VOC2012目标检测数据集格式转化为COCO目标检测数据集格式

你只需要传入一个VOC数据集根路径可以直接运行，简直不要太方便好吧

需要材料：

• VOC2012原始数据集
• 一个pycharm项目文件夹（任意）

达到效果：

• 自动生成当前项目下的 data/coco 文件夹

• VOC2012数据集标注格式转为标准的COCO2017数据集格式

• 图片按照训练集和验证集划分好

voc2012目标检测数据集格式

主要关注框起来的三个文件夹：

COCO2017数据集标注格式

annotation里面就是标注，而train2017和val2017是放的图片

转换代码如下：

• 文件复制到项目里面任何一个地方都能运行

运行代码：
python demo.py --voc_root={YOUR VOC} # 这里的 {YOUR VOC}就是你自己的voc路径

import os
import json
import xml.etree.ElementTree as ET
import argparse
import shutil


def transform_voc2coco(args):
    voc_root = args.voc_root

    # 获取工程的根绝对路径
    project_root = os.path.realpath(os.curdir)
    while True:
        if ".idea" in os.listdir(project_root):
            break
        else:
            project_root = os.path.join(project_root, "..")

    # 构建COCO完整目录
    coco_root = os.path.join(project_root, "data/coco")
    coco_train = os.path.join(coco_root, "train2017")
    coco_val = os.path.join(coco_root, "val2017")
    coco_anno = os.path.join(coco_root, "annotations")
    coco_train_anno = os.path.join(coco_anno, "instances_train2017.json")
    coco_val_anno = os.path.join(coco_anno, "instances_val2017.json")
    if not os.path.exists(coco_root):
        os.makedirs(coco_root)
        os.mkdir(coco_train)
        os.mkdir(coco_val)
        os.mkdir(coco_anno)

    coco_train_num = len(os.listdir(coco_train))
    coco_val_num = len(os.listdir(coco_val))
    print(f"train_2017 number: {coco_train_num}")
    print(f"val_2017 number: {coco_val_num}")

    # voc数据集目录
    voc_anno_dir = os.path.join(voc_root, "Annotations")
    voc_images_dir = os.path.join(voc_root, "JPEGImages")
    voc_train_txt = os.path.join(voc_root, "ImageSets/Main/train.txt")
    voc_val_txt = os.path.join(voc_root, "ImageSets/Main/val.txt")

    overwrite_images = True
    # 复制voc图片到coco
    if overwrite_images:
        # 复制训练集图片
        with open(voc_train_txt, "r") as f:
            lines = f.readlines()
            for line in lines:
                image_name = line.strip() + ".jpg"
                image_path = os.path.join(voc_images_dir, image_name)
                shutil.copy(image_path, os.path.join(coco_train, image_name))

        # 复制验证集图片
        with open(voc_val_txt, "r") as f:
            lines = f.readlines()
            for line in lines:
                image_name = line.strip() + ".jpg"
                image_path = os.path.join(voc_images_dir, image_name)
                shutil.copy(image_path, os.path.join(coco_val, image_name))

    def _extract_anno(fp, mode: str = "train"):
        txt_file = voc_train_txt if mode == "train" else voc_val_txt

        # 预定义VOC检测的20个类别以及超类
        supercategorys = ["vehicles", "household", "animals", "person"]
        vehicles = ["car", "bus", "bicycle", "motorbike", "aeroplane", "boat", "train"]
        household = ["chair", "sofa", "diningtable", "tvmonitor", "bottle", "pottedplant"]
        animals = ["cat", "dog", "cow", "horse", "sheep", "bird"]
        person = ["person"]
        classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
                   "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

        # 预构建coco格式的json文件
        json_file = {"info": [], "license": [], "images": [], "annotations": [],
                     "categories": [{"id": i, "name": class_,
                                     "supercategory": ""} for i, class_ in enumerate(classes)]}
        for i, class_ in enumerate(classes):
            if class_ in vehicles:
                json_file["categories"][i]["supercategory"] = supercategorys[0]
            elif class_ in household:
                json_file["categories"][i]["supercategory"] = supercategorys[1]
            elif class_ in animals:
                json_file["categories"][i]["supercategory"] = supercategorys[2]
            elif class_ in person:
                json_file["categories"][i]["supercategory"] = supercategorys[3]
            else:
                raise "unsupported class"

        # 写入json文件
		with open(txt_file, "r") as f_:
            lines = f_.readlines()
            for line in lines:
                image_xml = line.strip() + ".xml"
                image_xml_path = os.path.join(voc_anno_dir, image_xml)
                xml_obj = ET.parse(image_xml_path)
                root = xml_obj.getroot()
                img_dir = {"file_name": "", "width": 0, "height": 0, "objects": []}
                i = 0
                for eles in root:
                    if eles.tag == "filename":
                        img_dir["file_name"] = eles.text
                    elif eles.tag == "size":
                        for ele in eles:
                            img_dir["width"] = int(ele.text) if ele.tag == "width" else int(img_dir["width"])
                            img_dir["height"] = int(ele.text) if ele.tag == "height" else int(img_dir["height"])
                    elif eles.tag == "object":
                        obj_dir = {"name": "", "bndbox": [], "image_id": int(img_dir["file_name"].split(".")[0]),
                                   "id": i}
                        i = i + 1
                        for ele in eles:
                            obj_dir["name"] = ele.text if ele.tag == "name" else obj_dir["name"]
                            if ele.tag == "bndbox":
                                for pos in ele:
                                    if pos.tag == "xmin":
                                        xmin = int(pos.text)
                                    elif pos.tag == "xmax":
                                        xmax = int(pos.text)
                                    elif pos.tag == "ymin":
                                        ymin = int(pos.text)
                                    elif pos.tag == "ymax":
                                        ymax = int(pos.text)
                                    else:
                                        raise "unsupported pose"
                                obj_dir["bndbox"] = [xmin, ymin, xmax-xmin, ymax-ymin]
                        img_dir["objects"].append(obj_dir)
                        json_file["annotations"].append({"id": obj_dir["id"],
                                                         "image_id": obj_dir["image_id"],
                                                         "category_id": classes.index(obj_dir["name"]),
                                                         "segmentation": [],
                                                         "area": float(obj_dir["bndbox"][2]*obj_dir["bndbox"][3]),
                                                         "bbox": obj_dir["bndbox"],
                                                         "iscrowd": 0})
                    else:
                        continue
                json_file["images"].append({"file_name": img_dir["file_name"],
                                            "height": img_dir["height"],
                                            "width": img_dir["width"],
                                            "id": int(img_dir["file_name"].split(".")[0])})
            json.dump(json_file, fp)
            
    # 生成coco的annotation标注文件
    override_anno = True
    # 有一个文件不存在或者需要覆盖掉之前的标注
    if not os.path.exists(coco_train_anno) or not os.path.exists(coco_val_anno) or override_anno:
        with open(coco_train_anno, "w") as f:
            _extract_anno(f, "train")

        with open(coco_val_anno, "w") as f:
            _extract_anno(f, "val")

    print("-" * 30 + "finish" + "-" * 30)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--voc_root', type=str, default="/project/datasets/VOC2012")
    args = parser.parse_args()
    transform_voc2coco(args)