域迁移相关数据集生成脚本

数据转化

# ----------------------------------------------
# Created by Wei-Jie Huang
# A collection of annotation conversion scripts
# We provide this for reference only
# ----------------------------------------------


import json
from pathlib import Path
from tqdm import tqdm


r"""
coco_anno_dict is like {
    "images": list of image_info's
    "annotations": list of annotation_info's
    "categories": list of categorie_info's
}
where img_info is like: {
    "id": ...,                      # 0-indexed
    "width": ...,
    "height": ...,
    "file_name": ...,
}, annotation_info is like: {
    "id": ...,                      # 0-indexed
    "image_id": ...,
    "category_id": ...,
    "segmentation": ...,
    "iscrowd": ...,
    "area": ...,
    "bbox": ...,                    # (x, y, w, h)
}, and category_info is like: {
    "id": ...,                      # 1-indexed
    "name": ...,
}
"""


def sim10k_to_coco(
    src_path: str = "VOC2012/Annotations",
    des_path: str = "annotations/sim10k_caronly.json",
    categories: tuple = ("car",)
    ) -> None:

    r""" Convert Sim10k (in VOC format) into COCO format.
    Args:
        src_path: path of the directory containing VOC-format annotations
        des_path: destination of the converted COCO-fomat annotation
        categories: only category ``car`` is considered by default
    """

    from xml.etree import ElementTree

    src_path = Path(src_path)
    des_path = Path(des_path)
    assert src_path.exists(), "Annotation directory does not exist"
    if des_path.exists():
        print(f"{
     
     des_path} exists. Override? (y/n)", end=" ")
        if input() != "y":
            print("Abort")
            return
    else:
        des_path.parent.mkdir(parents=True, exist_ok=True)

    # Initialization
    coco_anno_dict = {
   
   
        "images": [],
        "categories": [],
        "annotations": [],
    }
    num_images = 0
    num_categories = 0
    num_annotations = 0

    # Categories
    category_to_id = {
   
   }
    for category in categories:
        coco_anno_dict["categories"].append({
   
   
            "id": num_categories + 1,
            "name": category
        })
        category_to_id[category] = num_categories + 1
        num_categories += 1

    # Start Conversion
    for anno_file in tqdm(list(src_path.glob("*.xml"))):
        et_root = ElementTree.parse(anno_file).getroot()

        ##### Images #####
        img_info = {
   
   
            "id": num_images,
            "file_name": anno_file.stem + ".jpg",
        }
        num_images += 1

        # Image Size
        size = et_root.find("size")
        img_info["width"] = int(size.find("width").text)
        img_info["height"] = int(size.find("height").text)

        coco_anno_dict["images"].append(img_info)

        ##### Annotations #####
        for anno_object in et_root.findall("object"):
            category = anno_object.find("name").text
            if category not in categories:
                continue
            anno_info = {
   
   
                "id": num_annotations,
                "image_id": img_info["id"],
                "category_id": category_to_id[category],
                "segmentation": [],
                "iscrowd": 0
            }
            num_annotations += 1

            # Bounding box
            bndbox = anno_object.find("bndbox")
            xmin = float(bndbox.find("xmin").text)
            ymin = float(bndbox.find("ymin").text)
            xmax = float(bndbox.find("xmax").text)
            ymax = float(bndbox.find("ymax").text)
            # COCO format expects (x, y, w, h)
            anno_info["bbox"] = [xmin