pluggable JsonIndexDataset

Summary: Make dataset type and args configurable on JsonIndexDatasetMapProvider.

Reviewed By: davnov134

Differential Revision: D36666705

fbshipit-source-id: 4d0a3781d9a956504f51f1c7134c04edf1eb2846

projects/implicitron_trainer/configs/repro_base.yaml

@@ -24,12 +24,13 @@ data_source_args:
     - 10
   dataset_map_provider_JsonIndexDatasetMapProvider_args:
     dataset_root: ${oc.env:CO3D_DATASET_ROOT}
-    load_point_clouds: false
-    mask_depths: false
-    mask_images: false
     n_frames_per_sequence: -1
     test_on_train: true
     test_restrict_sequence_id: 0
+    dataset_JsonIndexDataset_args:
+      load_point_clouds: false
+      mask_depths: false
+      mask_images: false
 generic_model_args:
   loss_weights:
     loss_mask_bce: 1.0

projects/implicitron_trainer/configs/repro_multiseq_base.yaml

@@ -20,9 +20,6 @@ data_source_args:
   dataset_map_provider_JsonIndexDatasetMapProvider_args:
     assert_single_seq: false
     task_str: multisequence
-    load_point_clouds: false
-    mask_depths: false
-    mask_images: false
     n_frames_per_sequence: -1
     test_on_train: true
     test_restrict_sequence_id: 0

projects/implicitron_trainer/tests/experiment.yaml

@@ -286,24 +286,35 @@ data_source_args:
     category: ???
     task_str: singlesequence
     dataset_root: ''
-    limit_to: -1
-    limit_sequences_to: -1
     n_frames_per_sequence: -1
     test_on_train: false
-    load_point_clouds: false
-    mask_images: false
-    mask_depths: false
     restrict_sequence_name: []
     test_restrict_sequence_id: -1
     assert_single_seq: false
     only_test_set: false
-    aux_dataset_kwargs:
-      box_crop: true
-      box_crop_context: 0.3
-      image_width: 800
-      image_height: 800
-      remove_empty_masks: true
+    dataset_class_type: JsonIndexDataset
     path_manager_factory_class_type: PathManagerFactory
+    dataset_JsonIndexDataset_args:
+      limit_to: 0
+      limit_sequences_to: 0
+      exclude_sequence: []
+      limit_category_to: []
+      load_images: true
+      load_depths: true
+      load_depth_masks: true
+      load_masks: true
+      load_point_clouds: false
+      max_points: 0
+      mask_images: false
+      mask_depths: false
+      image_height: 800
+      image_width: 800
+      box_crop: true
+      box_crop_mask_thr: 0.4
+      box_crop_context: 0.3
+      remove_empty_masks: true
+      seed: 0
+      sort_frames: false
     path_manager_factory_PathManagerFactory_args:
       silence_logs: true
   data_loader_map_provider_SequenceDataLoaderMapProvider_args:

projects/implicitron_trainer/tests/test_experiment.py

@@ -26,8 +26,6 @@ DATA_DIR = Path(__file__).resolve().parent
 DEBUG: bool = False
 
 # TODO:
-# - sort out path_manager config. Here we monkeypatch to avoid
-#    the problem.
 # - add enough files to skateboard_first_5 that this works on RE.
 # - share common code with PyTorch3D tests?
 # - deal with the temporary output files this test creates
@@ -54,7 +52,7 @@ class TestExperiment(unittest.TestCase):
         dataset_args.category = "skateboard"
         dataset_args.test_restrict_sequence_id = 0
         dataset_args.dataset_root = "manifold://co3d/tree/extracted"
-        dataset_args.limit_sequences_to = 5
+        dataset_args.dataset_JsonIndexDataset_args.limit_sequences_to = 5
         dataloader_args.dataset_len = 1
         cfg.solver_args.max_epochs = 2