Allow registering pascal voc dataset with custom class names

Summary:
Relate: facebookresearch#991
Pull Request resolved: facebookresearch#1603

Differential Revision: D22118724

Pulled By: ppwwyyxx

fbshipit-source-id: c6ad85e8f250345cf2e2abf512323ae1166051ae

.gitignore

@@ -14,6 +14,7 @@ inference_test_output
 __pycache__
 _ext
 *.pyc
+*.pyd
 *.so
 detectron2.egg-info/
 build/

detectron2/data/datasets/__init__.py

@@ -2,6 +2,7 @@
 from .cityscapes import load_cityscapes_instances
 from .coco import load_coco_json, load_sem_seg
 from .lvis import load_lvis_json, register_lvis_instances, get_lvis_instances_meta
+from .pascal_voc import load_voc_instances, register_pascal_voc
 from .register_coco import register_coco_instances, register_coco_panoptic_separated
 from . import builtin  # ensure the builtin datasets are registered
 

detectron2/data/datasets/pascal_voc.py

@@ -4,30 +4,32 @@
 import numpy as np
 import os
 import xml.etree.ElementTree as ET
+from typing import List, Tuple, Union
 from fvcore.common.file_io import PathManager
 
 from detectron2.data import DatasetCatalog, MetadataCatalog
 from detectron2.structures import BoxMode
 
-__all__ = ["register_pascal_voc"]
+__all__ = ["load_voc_instances", "register_pascal_voc"]
 
 
 # fmt: off
-CLASS_NAMES = [
+CLASS_NAMES = (
     "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat",
     "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person",
-    "pottedplant", "sheep", "sofa", "train", "tvmonitor",
-]
+    "pottedplant", "sheep", "sofa", "train", "tvmonitor"
+)
 # fmt: on
 
 
-def load_voc_instances(dirname: str, split: str):
+def load_voc_instances(dirname: str, split: str, class_names: Union[List[str], Tuple[str, ...]]):
     """
     Load Pascal VOC detection annotations to Detectron2 format.
 
     Args:
         dirname: Contain "Annotations", "ImageSets", "JPEGImages"
         split (str): one of "train", "test", "val", "trainval"
+        class_names: list or tuple of class names
     """
     with PathManager.open(os.path.join(dirname, "ImageSets", "Main", split + ".txt")) as f:
         fileids = np.loadtxt(f, dtype=np.str)
@@ -66,15 +68,15 @@ def load_voc_instances(dirname: str, split: str):
             bbox[0] -= 1.0
             bbox[1] -= 1.0
             instances.append(
-                {"category_id": CLASS_NAMES.index(cls), "bbox": bbox, "bbox_mode": BoxMode.XYXY_ABS}
+                {"category_id": class_names.index(cls), "bbox": bbox, "bbox_mode": BoxMode.XYXY_ABS}
             )
         r["annotations"] = instances
         dicts.append(r)
     return dicts
 
 
-def register_pascal_voc(name, dirname, split, year):
-    DatasetCatalog.register(name, lambda: load_voc_instances(dirname, split))
+def register_pascal_voc(name, dirname, split, year, class_names=CLASS_NAMES):
+    DatasetCatalog.register(name, lambda: load_voc_instances(dirname, split, class_names))
     MetadataCatalog.get(name).set(
-        thing_classes=CLASS_NAMES, dirname=dirname, year=year, split=split
+        thing_classes=list(class_names), dirname=dirname, year=year, split=split
     )