Adding SQL dataset classes to ImplicitronDataSource imports

Summary: Making it easier for the clients to use these datasets. Reviewed By: bottler Differential Revision: D46727179 fbshipit-source-id: cf619aee4c4c0222a74b30ea590cf37f08f014cc
facebookresearch · Jun 14, 2023 · cd5db07 · cd5db07
1 parent 3d886c3
commit cd5db07
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 0 deletions.
diff --git a/projects/implicitron_trainer/tests/experiment.yaml b/projects/implicitron_trainer/tests/experiment.yaml
@@ -129,6 +129,19 @@ data_source_ImplicitronDataSource_args:
     dataset_length_train: 0
     dataset_length_val: 0
     dataset_length_test: 0
+  data_loader_map_provider_TrainEvalDataLoaderMapProvider_args:
+    batch_size: 1
+    num_workers: 0
+    dataset_length_train: 0
+    dataset_length_val: 0
+    dataset_length_test: 0
+    train_conditioning_type: SAME
+    val_conditioning_type: SAME
+    test_conditioning_type: KNOWN
+    images_per_seq_options: []
+    sample_consecutive_frames: false
+    consecutive_frames_max_gap: 0
+    consecutive_frames_max_gap_seconds: 0.1
 model_factory_ImplicitronModelFactory_args:
   resume: true
   model_class_type: GenericModel

diff --git a/projects/implicitron_trainer/tests/test_experiment.py b/projects/implicitron_trainer/tests/test_experiment.py
@@ -136,6 +136,8 @@ def test_yaml_contents(self):
         ds_arg = cfg.data_source_ImplicitronDataSource_args
         ds_arg.dataset_map_provider_JsonIndexDatasetMapProvider_args.dataset_root = ""
         ds_arg.dataset_map_provider_JsonIndexDatasetMapProviderV2_args.dataset_root = ""
+        if "dataset_map_provider_SqlIndexDatasetMapProvider_args" in ds_arg:
+            del ds_arg.dataset_map_provider_SqlIndexDatasetMapProvider_args
         cfg.training_loop_ImplicitronTrainingLoop_args.visdom_port = 8097
         yaml = OmegaConf.to_yaml(cfg, sort_keys=False)
         if DEBUG:

diff --git a/pytorch3d/implicitron/dataset/data_source.py b/pytorch3d/implicitron/dataset/data_source.py
@@ -72,6 +72,16 @@ def pre_expand(cls) -> None:
             from .rendered_mesh_dataset_map_provider import (  # noqa: F401
                 RenderedMeshDatasetMapProvider,
             )
+            from .train_eval_data_loader_provider import (  # noqa: F401
+                TrainEvalDataLoaderMapProvider,
+            )
+
+            try:
+                from .sql_dataset_provider import (  # noqa: F401  # pyre-ignore
+                    SqlIndexDatasetMapProvider,
+                )
+            except ModuleNotFoundError:
+                pass  # environment without SQL dataset
         finally:
             pass
 

diff --git a/tests/implicitron/data/data_source.yaml b/tests/implicitron/data/data_source.yaml
@@ -116,3 +116,16 @@ data_loader_map_provider_SimpleDataLoaderMapProvider_args:
   dataset_length_train: 0
   dataset_length_val: 0
   dataset_length_test: 0
+data_loader_map_provider_TrainEvalDataLoaderMapProvider_args:
+  batch_size: 1
+  num_workers: 0
+  dataset_length_train: 0
+  dataset_length_val: 0
+  dataset_length_test: 0
+  train_conditioning_type: SAME
+  val_conditioning_type: SAME
+  test_conditioning_type: KNOWN
+  images_per_seq_options: []
+  sample_consecutive_frames: false
+  consecutive_frames_max_gap: 0
+  consecutive_frames_max_gap_seconds: 0.1
diff --git a/tests/implicitron/test_data_source.py b/tests/implicitron/test_data_source.py
@@ -68,6 +68,9 @@ def test_one(self):
         # making the test invariant to env variables
         cfg.dataset_map_provider_JsonIndexDatasetMapProvider_args.dataset_root = ""
         cfg.dataset_map_provider_JsonIndexDatasetMapProviderV2_args.dataset_root = ""
+        # making the test invariant to the presence of SQL dataset
+        if "dataset_map_provider_SqlIndexDatasetMapProvider_args" in cfg:
+            del cfg.dataset_map_provider_SqlIndexDatasetMapProvider_args
         yaml = OmegaConf.to_yaml(cfg, sort_keys=False)
         if DEBUG:
             (DATA_DIR / "data_source.yaml").write_text(yaml)