add extra_annotation_keys parameter to lvis json loader

Summary:
Thanks for your contribution!

If you're sending a large PR (e.g., >100 lines),
please open an issue first about the feature / bug, and indicate how you want to contribute.

We do not always accept features.
See https://detectron2.readthedocs.io/notes/contributing.html#pull-requests about how we handle PRs.

Before submitting a PR, please run `dev/linter.sh` to lint the code.

Pull Request resolved: fairinternal/detectron2#556

Reviewed By: ppwwyyxx

Differential Revision: D29925800

Pulled By: alexander-kirillov

fbshipit-source-id: 8e8b7452da98d73f2fd6a7f936658cfcb9ed98f4

Author: alexander-kirillov

detectron2/data/datasets/lvis.py

@@ -37,7 +37,7 @@ def register_lvis_instances(name, metadata, json_file, image_root):
     )
 
 
-def load_lvis_json(json_file, image_root, dataset_name=None):
+def load_lvis_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None):
     """
     Load a json file in LVIS's annotation format.
 
@@ -47,6 +47,9 @@ def load_lvis_json(json_file, image_root, dataset_name=None):
         dataset_name (str): the name of the dataset (e.g., "lvis_v0.5_train").
             If provided, this function will put "thing_classes" into the metadata
             associated with this dataset.
+        extra_annotation_keys (list[str]): list of per-annotation keys that should also be
+            loaded into the dataset dict (besides "bbox", "bbox_mode", "category_id",
+            "segmentation"). The values for these keys will be returned as-is.
 
     Returns:
         list[dict]: a list of dicts in Detectron2 standard format. (See
@@ -106,6 +109,13 @@ def load_lvis_json(json_file, image_root, dataset_name=None):
 
     logger.info("Loaded {} images in the LVIS format from {}".format(len(imgs_anns), json_file))
 
+    if extra_annotation_keys:
+        logger.info(
+            "The following extra annotation keys will be loaded: {} ".format(extra_annotation_keys)
+        )
+    else:
+        extra_annotation_keys = []
+
     def get_file_name(img_root, img_dict):
         # Determine the path including the split folder ("train2017", "val2017", "test2017") from
         # the coco_url field. Example:
@@ -145,6 +155,8 @@ def get_file_name(img_root, img_dict):
             ), "Annotation contains an invalid polygon with < 3 points"
             assert len(segm) > 0
             obj["segmentation"] = segm
+            for extra_ann_key in extra_annotation_keys:
+                obj[extra_ann_key] = anno[extra_ann_key]
             objs.append(obj)
         record["annotations"] = objs
         dataset_dicts.append(record)