rename ImplicitronDataset to JsonIndexDataset

Summary: The ImplicitronDataset class corresponds to JsonIndexDatasetMapProvider

Reviewed By: shapovalov

Differential Revision: D36661396

fbshipit-source-id: 80ca2ff81ef9ecc2e3d1f4e1cd14b6f66a7ec34d

projects/implicitron_trainer/experiment.py

@@ -68,7 +68,7 @@ from pytorch3d.implicitron.dataset.data_loader_map_provider import DataLoaderMap
 from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource, Task
 from pytorch3d.implicitron.dataset.dataset_base import FrameData
 from pytorch3d.implicitron.dataset.dataset_map_provider import DatasetMap
-from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDataset
+from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset
 from pytorch3d.implicitron.evaluation import evaluate_new_view_synthesis as evaluate
 from pytorch3d.implicitron.models.generic_model import EvaluationMode, GenericModel
 from pytorch3d.implicitron.tools import model_io, vis_utils
@@ -641,7 +641,7 @@ def _run_eval(model, all_source_cameras, loader, task: Task, device):
 
 
 def _get_all_source_cameras(
-    dataset: ImplicitronDataset,
+    dataset: JsonIndexDataset,
     num_workers: int = 8,
 ) -> CamerasBase:
     """

projects/implicitron_trainer/visualize_reconstruction.py

@@ -25,7 +25,7 @@ from experiment import init_model
 from omegaconf import OmegaConf
 from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource
 from pytorch3d.implicitron.dataset.dataset_base import FrameData
-from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDataset
+from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset
 from pytorch3d.implicitron.dataset.utils import is_train_frame
 from pytorch3d.implicitron.models.base_model import EvaluationMode
 from pytorch3d.implicitron.tools.configurable import get_default_args
@@ -41,7 +41,7 @@ from tqdm import tqdm
 
 
 def render_sequence(
-    dataset: ImplicitronDataset,
+    dataset: JsonIndexDataset,
     sequence_name: str,
     model: torch.nn.Module,
     video_path,

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -49,9 +49,10 @@ class FrameAnnotsEntry(TypedDict):
 
 
 @dataclass(eq=False)
-class ImplicitronDataset(DatasetBase):
+class JsonIndexDataset(DatasetBase):
     """
-    A class for the Common Objects in 3D (CO3D) dataset.
+    A dataset with annotations in json files like the Common Objects in 3D
+    (CO3D) dataset.
 
     Args:
         frame_annotations_file: A zipped json file containing metadata of the
@@ -148,7 +149,7 @@ class ImplicitronDataset(DatasetBase):
     seq_annots: Dict[str, types.SequenceAnnotation] = field(init=False)
 
     def __post_init__(self) -> None:
-        # pyre-fixme[16]: `ImplicitronDataset` has no attribute `subset_to_image_path`.
+        # pyre-fixme[16]: `JsonIndexDataset` has no attribute `subset_to_image_path`.
         self.subset_to_image_path = None
         self._load_frames()
         self._load_sequences()
@@ -193,7 +194,7 @@ class ImplicitronDataset(DatasetBase):
         return batches_idx
 
     def __str__(self) -> str:
-        return f"ImplicitronDataset #frames={len(self.frame_annots)}"
+        return f"JsonIndexDataset #frames={len(self.frame_annots)}"
 
     def __len__(self) -> int:
         return len(self.frame_annots)

pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py

@@ -14,7 +14,7 @@ from omegaconf import DictConfig
 from pytorch3d.implicitron.tools.config import registry
 
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, Task
-from .implicitron_dataset import ImplicitronDataset
+from .json_index_dataset import JsonIndexDataset
 from .utils import (
     DATASET_TYPE_KNOWN,
     DATASET_TYPE_TEST,
@@ -86,7 +86,7 @@ class JsonIndexDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13]
             are present in all generated datasets.
         only_test_set: Load only the test set.
         aux_dataset_kwargs: Specifies additional arguments to the
-            ImplicitronDataset constructor call.
+            JsonIndexDataset constructor call.
         path_manager: Optional[PathManager] for interpreting paths
     """
 
@@ -196,7 +196,7 @@ class JsonIndexDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13]
 
         train_dataset = None
         if not self.only_test_set:
-            train_dataset = ImplicitronDataset(
+            train_dataset = JsonIndexDataset(
                 n_frames_per_sequence=self.n_frames_per_sequence,
                 subsets=set_names_mapping["train"],
                 pick_sequence=restrict_sequence_name,
@@ -206,13 +206,13 @@ class JsonIndexDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13]
             assert train_dataset is not None
             val_dataset = test_dataset = train_dataset
         else:
-            val_dataset = ImplicitronDataset(
+            val_dataset = JsonIndexDataset(
                 n_frames_per_sequence=-1,
                 subsets=set_names_mapping["val"],
                 pick_sequence=restrict_sequence_name,
                 **common_kwargs,
             )
-            test_dataset = ImplicitronDataset(
+            test_dataset = JsonIndexDataset(
                 n_frames_per_sequence=-1,
                 subsets=set_names_mapping["test"],
                 pick_sequence=restrict_sequence_name,

pytorch3d/implicitron/dataset/visualize.py

@@ -11,11 +11,11 @@ from pytorch3d.implicitron.tools.point_cloud_utils import get_rgbd_point_cloud
 from pytorch3d.structures import Pointclouds
 
 from .dataset_base import FrameData
-from .implicitron_dataset import ImplicitronDataset
+from .json_index_dataset import JsonIndexDataset
 
 
 def get_implicitron_sequence_pointcloud(
-    dataset: ImplicitronDataset,
+    dataset: JsonIndexDataset,
     sequence_name: Optional[str] = None,
     mask_points: bool = True,
     max_frames: int = -1,

pytorch3d/implicitron/eval_demo.py

@@ -14,7 +14,7 @@ import torch
 from iopath.common.file_io import PathManager
 from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource, Task
 from pytorch3d.implicitron.dataset.dataset_base import DatasetBase, FrameData
-from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDataset
+from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset
 from pytorch3d.implicitron.dataset.json_index_dataset_map_provider import (
     CO3D_CATEGORIES,
 )
@@ -126,7 +126,7 @@ def evaluate_dbir_for_category(
     if task == Task.SINGLE_SEQUENCE:
         # all_source_cameras are needed for evaluation of the
         # target camera difficulty
-        # pyre-fixme[16]: `ImplicitronDataset` has no attribute `frame_annots`.
+        # pyre-fixme[16]: `JsonIndexDataset` has no attribute `frame_annots`.
         sequence_name = test_dataset.frame_annots[0]["frame_annotation"].sequence_name
         all_source_cameras = _get_all_source_cameras(
             test_dataset, sequence_name, num_workers=num_workers
@@ -134,7 +134,7 @@ def evaluate_dbir_for_category(
     else:
         all_source_cameras = None
 
-    image_size = cast(ImplicitronDataset, test_dataset).image_width
+    image_size = cast(JsonIndexDataset, test_dataset).image_width
 
     if image_size is None:
         raise ValueError("Image size should be set in the dataset")

tests/implicitron/test_dataset_visualize.py

@@ -12,7 +12,7 @@ import unittest
 
 import torch
 import torchvision
-from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDataset
+from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset
 from pytorch3d.implicitron.dataset.visualize import get_implicitron_sequence_pointcloud
 from pytorch3d.implicitron.tools.point_cloud_utils import render_point_cloud_pytorch3d
 from pytorch3d.vis.plotly_vis import plot_scene
@@ -38,7 +38,7 @@ class TestDatasetVisualize(unittest.TestCase):
         sequence_file = os.path.join(dataset_root, category, "sequence_annotations.jgz")
         self.image_size = 256
         self.datasets = {
-            "simple": ImplicitronDataset(
+            "simple": JsonIndexDataset(
                 frame_annotations_file=frame_file,
                 sequence_annotations_file=sequence_file,
                 dataset_root=dataset_root,
@@ -48,7 +48,7 @@ class TestDatasetVisualize(unittest.TestCase):
                 load_point_clouds=True,
                 path_manager=path_manager,
             ),
-            "nonsquare": ImplicitronDataset(
+            "nonsquare": JsonIndexDataset(
                 frame_annotations_file=frame_file,
                 sequence_annotations_file=sequence_file,
                 dataset_root=dataset_root,
@@ -58,7 +58,7 @@ class TestDatasetVisualize(unittest.TestCase):
                 load_point_clouds=True,
                 path_manager=path_manager,
             ),
-            "nocrop": ImplicitronDataset(
+            "nocrop": JsonIndexDataset(
                 frame_annotations_file=frame_file,
                 sequence_annotations_file=sequence_file,
                 dataset_root=dataset_root,

tests/implicitron/test_evaluation.py

@@ -15,7 +15,7 @@ import unittest
 import lpips
 import torch
 from pytorch3d.implicitron.dataset.dataset_base import FrameData
-from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDataset
+from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset
 from pytorch3d.implicitron.evaluation.evaluate_new_view_synthesis import eval_batch
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
 from pytorch3d.implicitron.models.generic_model import GenericModel  # noqa
@@ -44,7 +44,7 @@ class TestEvaluation(unittest.TestCase):
         frame_file = os.path.join(dataset_root, category, "frame_annotations.jgz")
         sequence_file = os.path.join(dataset_root, category, "sequence_annotations.jgz")
         self.image_size = 64
-        self.dataset = ImplicitronDataset(
+        self.dataset = JsonIndexDataset(
             frame_annotations_file=frame_file,
             sequence_annotations_file=sequence_file,
             dataset_root=dataset_root,