From 4fddaa85d490dfc5d9e35ed7bcd641ce195096f6 Mon Sep 17 00:00:00 2001 From: Zhicheng Yan Date: Thu, 22 Jun 2023 13:29:13 -0700 Subject: [PATCH] update INJECTED_COCO_DATASETS_LUT when registering AdhocCOCODataset Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/584 When registering AdhocCOCODataset, INJECTED_COCO_DATASETS_LUT needs to be updated as well. For example, if a dataset uses custom registering function, it can be only retrieved from INJECTED_COCO_DATASETS_LUT. Otherwise, it uses the default registering function as in branch `register_dataset_split`. Reviewed By: antonrigner Differential Revision: D46826507 fbshipit-source-id: 31bc5c79e3b82a20cb798dec37b18d9271040cdf --- d2go/data/utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/d2go/data/utils.py b/d2go/data/utils.py index e6779051..d8305ff9 100644 --- a/d2go/data/utils.py +++ b/d2go/data/utils.py @@ -20,6 +20,7 @@ ANN_FN, IM_DIR, INJECTED_COCO_DATASETS_LUT, + InjectedCocoEntry, register_dataset_split, ) from detectron2.data import DatasetCatalog, MetadataCatalog @@ -168,9 +169,10 @@ def register_catalog(self): DatasetCatalog.register(self.new_ds_name, new_func) elif self.src_ds_name in INJECTED_COCO_DATASETS_LUT: _src_func, _src_dict = INJECTED_COCO_DATASETS_LUT[self.src_ds_name] - _src_func( - self.new_ds_name, - split_dict={**_src_dict, ANN_FN: tmp_file, IM_DIR: metadata.image_root}, + split_dict = {**_src_dict, ANN_FN: tmp_file, IM_DIR: metadata.image_root} + _src_func(self.new_ds_name, split_dict=split_dict) + INJECTED_COCO_DATASETS_LUT[self.new_ds_name] = InjectedCocoEntry( + func=_src_func, split_dict=split_dict ) else: # NOTE: only supports COCODataset as DS_TYPE since we cannot reconstruct