##############################################################################
#
# Below code is inspired on
# https://github.com/facebookresearch/detectron2/blob/master/detectron2/data/datasets/pascal_voc.py
# --------------------------------------------------------
# Detectron2
# Licensed under the Apache 2.0 license.
# --------------------------------------------------------

from fvcore.common.file_io import PathManager
import os
import numpy as np
import xml.etree.ElementTree as ET

from detectron2.structures import BoxMode
from detectron2.data import DatasetCatalog, MetadataCatalog


__all__ = ["register_licenseplates_voc"]


CLASS_NAMES = [
    "license_plate",
]


def load_voc_instances(dirname: str, split: str):
    """
    Load licenseplates VOC detection annotations to Detectron2 format.

    Args:
        dirname: Contain "annotations", "images"
        split (str): one of "train", "test"
    """
    with PathManager.open(os.path.join(dirname, split + ".txt")) as f:
        fileids = np.loadtxt(f, dtype=np.str)

    dicts = []
    for fileid in fileids:
        anno_file = os.path.join(dirname, "annotations", fileid + ".xml")
        jpeg_file = os.path.join(dirname, "images", fileid + ".jpg")

        tree = ET.parse(anno_file)

        r = {
            "file_name": jpeg_file,
            "image_id": fileid,
            "height": int(tree.findall("./size/height")[0].text),
            "width": int(tree.findall("./size/width")[0].text),
        }
        instances = []

        for obj in tree.findall("object"):
            cls = obj.find("name").text
            bbox = obj.find("bndbox")
            bbox = [float(bbox.find(x).text) for x in ["xmin", "ymin", "xmax", "ymax"]]
            instances.append(
                {"category_id": CLASS_NAMES.index(cls), "bbox": bbox, "bbox_mode": BoxMode.XYXY_ABS}
            )
        r["annotations"] = instances
        dicts.append(r)
    return dicts


def register_licenseplates_voc(name, dirname, split):
    DatasetCatalog.register(name,
                            lambda: load_voc_instances(dirname, split))
    MetadataCatalog.get(name).set(thing_classes=CLASS_NAMES,
                                  dirname=dirname,
                                  split=split)


if __name__ == "__main__":
    import random
    import cv2
    from detectron2.utils.visualizer import Visualizer
    import argparse

    # Parse command line arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", default="train")
    ap.add_argument("--samples", type=int, default=10)
    ap.add_argument("--scale", type=float, default=1.0)
    args = ap.parse_args()

    dataset_name = f"licenseplates_{args.split}"
    register_licenseplates_voc(dataset_name, "datasets/licenseplates", args.split)
    dataset_dicts = DatasetCatalog.get(dataset_name)
    for d in random.sample(dataset_dicts, args.samples):
        img = cv2.imread(d["file_name"])
        visualizer = Visualizer(img[:, :, ::-1],
                                metadata=MetadataCatalog.get(dataset_name),
                                scale=args.scale)
        vis = visualizer.draw_dataset_dict(d)
        cv2.imshow(dataset_name, vis.get_image()[:, :, ::-1])

        # Exit? Press ESC
        if cv2.waitKey(0) & 0xFF == 27:
            break

    cv2.destroyAllWindows()

[((2443, 2468), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (2466, 2468), False, 'import argparse\n'), ((2808, 2840), 'detectron2.data.DatasetCatalog.get', 'DatasetCatalog.get', (['dataset_name'], {}), '(dataset_name)\n', (2826, 2840), False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((2854, 2896), 'random.sample', 'random.sample', (['dataset_dicts', 'args.samples'], {}), '(dataset_dicts, args.samples)\n', (2867, 2896), False, 'import random\n'), ((3312, 3335), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), '()\n', (3333, 3335), False, 'import cv2\n'), ((1011, 1038), 'numpy.loadtxt', 'np.loadtxt', (['f'], {'dtype': 'np.str'}), '(f, dtype=np.str)\n', (1021, 1038), True, 'import numpy as np\n'), ((1102, 1155), 'os.path.join', 'os.path.join', (['dirname', '"""annotations"""', "(fileid + '.xml')"], {}), "(dirname, 'annotations', fileid + '.xml')\n", (1114, 1155), False, 'import os\n'), ((1176, 1224), 'os.path.join', 'os.path.join', (['dirname', '"""images"""', "(fileid + '.jpg')"], {}), "(dirname, 'images', fileid + '.jpg')\n", (1188, 1224), False, 'import os\n'), ((1241, 1260), 'xml.etree.ElementTree.parse', 'ET.parse', (['anno_file'], {}), '(anno_file)\n', (1249, 1260), True, 'import xml.etree.ElementTree as ET\n'), ((2912, 2938), 'cv2.imread', 'cv2.imread', (["d['file_name']"], {}), "(d['file_name'])\n", (2922, 2938), False, 'import cv2\n'), ((948, 985), 'os.path.join', 'os.path.join', (['dirname', "(split + '.txt')"], {}), "(dirname, split + '.txt')\n", (960, 985), False, 'import os\n'), ((2106, 2131), 'detectron2.data.MetadataCatalog.get', 'MetadataCatalog.get', (['name'], {}), '(name)\n', (2125, 2131), False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((3029, 3062), 'detectron2.data.MetadataCatalog.get', 'MetadataCatalog.get', (['dataset_name'], {}), '(dataset_name)\n', (3048, 3062), False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((3260, 3274), 'cv2.waitKey', 'cv2.waitKey', (['(0)'], {}), '(0)\n', (3271, 3274), False, 'import cv2\n')]