diff --git a/projects/implicitron_trainer/experiment.py b/projects/implicitron_trainer/experiment.py index 0ce35e5a..a888e225 100755 --- a/projects/implicitron_trainer/experiment.py +++ b/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: """ diff --git a/projects/implicitron_trainer/visualize_reconstruction.py b/projects/implicitron_trainer/visualize_reconstruction.py index 773a321d..f8a9aa1f 100644 --- a/projects/implicitron_trainer/visualize_reconstruction.py +++ b/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, diff --git a/pytorch3d/implicitron/dataset/implicitron_dataset.py b/pytorch3d/implicitron/dataset/json_index_dataset.py similarity index 99% rename from pytorch3d/implicitron/dataset/implicitron_dataset.py rename to pytorch3d/implicitron/dataset/json_index_dataset.py index 78e8c00f..bd8d1711 100644 --- a/pytorch3d/implicitron/dataset/implicitron_dataset.py +++ b/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) diff --git a/pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py b/pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py index 1932697f..461cb8b5 100644 --- a/pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py +++ b/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, diff --git a/pytorch3d/implicitron/dataset/visualize.py b/pytorch3d/implicitron/dataset/visualize.py index ebd957c5..331ee892 100644 --- a/pytorch3d/implicitron/dataset/visualize.py +++ b/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, diff --git a/pytorch3d/implicitron/eval_demo.py b/pytorch3d/implicitron/eval_demo.py index 2be1c0df..54267c42 100644 --- a/pytorch3d/implicitron/eval_demo.py +++ b/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") diff --git a/tests/implicitron/test_dataset_visualize.py b/tests/implicitron/test_dataset_visualize.py index 380e14ce..2eb4d5c5 100644 --- a/tests/implicitron/test_dataset_visualize.py +++ b/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, diff --git a/tests/implicitron/test_evaluation.py b/tests/implicitron/test_evaluation.py index 3731d036..f8a7ddfd 100644 --- a/tests/implicitron/test_evaluation.py +++ b/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,