Commit 705f5476 authored by Bo (YANG, Bo) Yang's avatar Bo (YANG, Bo) Yang Committed by Facebook GitHub Bot
Browse files

Allow passing additional load parameters to extended_coco_load

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/251

As title

Reviewed By: wenliangzhao2018

Differential Revision: D36360786

fbshipit-source-id: c52fa9a108a908fb96d16b2affd18bed19efff5e
parent 091213f0
......@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
D2GO_DATASETS_BASE_MODULE = "d2go.datasets"
IM_DIR = "image_directory"
ANN_FN = "annotation_file"
LOAD_KWARGS = "load_kwargs"
COCO_REGISTER_FUNCTION_REGISTRY = Registry("COCO_REGISTER_FUNCTION_REGISTRY")
COCO_REGISTER_FUNCTION_REGISTRY.__doc__ = "Registry - coco register function"
......@@ -47,6 +48,7 @@ def _import_dataset(module_name):
def _register_extended_coco(dataset_name, split_dict):
json_file = split_dict[ANN_FN]
image_root = split_dict[IM_DIR]
load_kwargs = split_dict.get(LOAD_KWARGS, {})
# 1. register a function which returns dicts
load_coco_json_func = functools.partial(
......@@ -54,6 +56,7 @@ def _register_extended_coco(dataset_name, split_dict):
json_file=json_file,
image_root=image_root,
dataset_name=dataset_name,
**load_kwargs,
)
DatasetCatalog.register(dataset_name, load_coco_json_func)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment