##############################################################################
#
# 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()

[((80, 9, 80, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((88, 20, 88, 52), 'detectron2.data.DatasetCatalog.get', 'DatasetCatalog.get', ({(88, 39, 88, 51): 'dataset_name'}, {}), '(dataset_name)', False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((89, 13, 89, 55), 'random.sample', 'random.sample', ({(89, 27, 89, 40): 'dataset_dicts', (89, 42, 89, 54): 'args.samples'}, {}), '(dataset_dicts, args.samples)', False, 'import random\n'), ((101, 4, 101, 27), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ({}, {}), '()', False, 'import cv2\n'), ((36, 18, 36, 45), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((40, 20, 40, 73), 'os.path.join', 'os.path.join', ({(40, 33, 40, 40): 'dirname', (40, 42, 40, 55): '"""annotations"""', (40, 57, 40, 72): "fileid + '.xml'"}, {}), "(dirname, 'annotations', fileid + '.xml')", False, 'import os\n'), ((41, 20, 41, 68), 'os.path.join', 'os.path.join', ({(41, 33, 41, 40): 'dirname', (41, 42, 41, 50): '"""images"""', (41, 52, 41, 67): "fileid + '.jpg'"}, {}), "(dirname, 'images', fileid + '.jpg')", False, 'import os\n'), ((43, 15, 43, 34), 'xml.etree.ElementTree.parse', 'ET.parse', ({(43, 24, 43, 33): 'anno_file'}, {}), '(anno_file)', True, 'import xml.etree.ElementTree as ET\n'), ((90, 14, 90, 40), 'cv2.imread', 'cv2.imread', ({(90, 25, 90, 39): "d['file_name']"}, {}), "(d['file_name'])", False, 'import cv2\n'), ((35, 26, 35, 63), 'os.path.join', 'os.path.join', ({(35, 39, 35, 46): 'dirname', (35, 48, 35, 62): "(split + '.txt')"}, {}), "(dirname, split + '.txt')", False, 'import os\n'), ((68, 4, 68, 29), 'detectron2.data.MetadataCatalog.get', 'MetadataCatalog.get', ({(68, 24, 68, 28): 'name'}, {}), '(name)', False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((92, 41, 92, 74), 'detectron2.data.MetadataCatalog.get', 'MetadataCatalog.get', ({(92, 61, 92, 73): 'dataset_name'}, {}), '(dataset_name)', False, 'from detectron2.data import DatasetCatalog, MetadataCatalog\n'), ((98, 11, 98, 25), 'cv2.waitKey', 'cv2.waitKey', ({(98, 23, 98, 24): '(0)'}, {}), '(0)', False, 'import cv2\n')]