数据转化
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)
coco_anno_dict = {
"images": [],
"categories": [],
"annotations": [],
}
num_images = 0
num_categories = 0
num_annotations = 0
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
for anno_file in tqdm(list(src_path.glob("*.xml"))):
et_root = ElementTree.parse(anno_file).getroot()
img_info = {
"id": num_images,
"file_name": anno_file.stem + ".jpg",
}
num_images += 1
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)
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
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)
anno_info["bbox"] = [xmin