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__) ...@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
D2GO_DATASETS_BASE_MODULE = "d2go.datasets" D2GO_DATASETS_BASE_MODULE = "d2go.datasets"
IM_DIR = "image_directory" IM_DIR = "image_directory"
ANN_FN = "annotation_file" ANN_FN = "annotation_file"
LOAD_KWARGS = "load_kwargs"
COCO_REGISTER_FUNCTION_REGISTRY = Registry("COCO_REGISTER_FUNCTION_REGISTRY") COCO_REGISTER_FUNCTION_REGISTRY = Registry("COCO_REGISTER_FUNCTION_REGISTRY")
COCO_REGISTER_FUNCTION_REGISTRY.__doc__ = "Registry - coco register function" COCO_REGISTER_FUNCTION_REGISTRY.__doc__ = "Registry - coco register function"
...@@ -47,6 +48,7 @@ def _import_dataset(module_name): ...@@ -47,6 +48,7 @@ def _import_dataset(module_name):
def _register_extended_coco(dataset_name, split_dict): def _register_extended_coco(dataset_name, split_dict):
json_file = split_dict[ANN_FN] json_file = split_dict[ANN_FN]
image_root = split_dict[IM_DIR] image_root = split_dict[IM_DIR]
load_kwargs = split_dict.get(LOAD_KWARGS, {})
# 1. register a function which returns dicts # 1. register a function which returns dicts
load_coco_json_func = functools.partial( load_coco_json_func = functools.partial(
...@@ -54,6 +56,7 @@ def _register_extended_coco(dataset_name, split_dict): ...@@ -54,6 +56,7 @@ def _register_extended_coco(dataset_name, split_dict):
json_file=json_file, json_file=json_file,
image_root=image_root, image_root=image_root,
dataset_name=dataset_name, dataset_name=dataset_name,
**load_kwargs,
) )
DatasetCatalog.register(dataset_name, load_coco_json_func) 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