diff --git a/pytorch3d/implicitron/dataset/json_index_dataset.py b/pytorch3d/implicitron/dataset/json_index_dataset.py index 64e22b959..6dfcccccc 100644 --- a/pytorch3d/implicitron/dataset/json_index_dataset.py +++ b/pytorch3d/implicitron/dataset/json_index_dataset.py @@ -414,15 +414,26 @@ def __getitem__(self, index) -> FrameData: ) if self.load_point_clouds and point_cloud is not None: - frame_data.sequence_point_cloud_path = pcl_path = os.path.join( - self.dataset_root, point_cloud.path - ) + pcl_path = self._fix_point_cloud_path(point_cloud.path) frame_data.sequence_point_cloud = _load_pointcloud( self._local_path(pcl_path), max_points=self.max_points ) + frame_data.sequence_point_cloud_path = pcl_path return frame_data + def _fix_point_cloud_path(self, path: str) -> str: + """ + Fix up a point cloud path from the dataset. + Some files in Co3Dv2 have an accidental absolute path stored. + """ + unwanted_prefix = ( + "/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/" + ) + if path.startswith(unwanted_prefix): + path = path[len(unwanted_prefix) :] + return os.path.join(self.dataset_root, path) + def _load_crop_fg_probability( self, entry: types.FrameAnnotation ) -> Tuple[ diff --git a/tests/implicitron/common_resources.py b/tests/implicitron/common_resources.py index 935f3ee14..93f904b9e 100644 --- a/tests/implicitron/common_resources.py +++ b/tests/implicitron/common_resources.py @@ -17,6 +17,9 @@ CO3D_MANIFOLD_PATH: str = "manifold://co3d/tree/extracted" +CO3DV2_MANIFOLD_PATH: str = "manifold://co3d/tree/v2/extracted" + +INSIDE_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER", False) def get_path_manager(silence_logs: bool = False) -> PathManager: @@ -30,7 +33,7 @@ def get_path_manager(silence_logs: bool = False) -> PathManager: logging.getLogger("iopath.fb.manifold").setLevel(logging.CRITICAL) logging.getLogger("iopath.common.file_io").setLevel(logging.CRITICAL) - if os.environ.get("INSIDE_RE_WORKER", False): + if INSIDE_RE_WORKER: raise ValueError("Cannot get to manifold from RE") path_manager = PathManager() @@ -70,7 +73,7 @@ def get_skateboard_data( raise unittest.SkipTest("Unknown environment. Data not available.") yield "/datasets01/co3d/081922", PathManager() - elif avoid_manifold or os.environ.get("INSIDE_RE_WORKER", False): + elif avoid_manifold or INSIDE_RE_WORKER: from libfb.py.parutil import get_file_path par_path = "skateboard_first_5" @@ -120,7 +123,7 @@ def _provide_torchvision_weights(par_path: str, filename: str) -> None: # (It can't copy straight to a nested location, see # https://fb.workplace.com/groups/askbuck/posts/2644615728920359/) # Here we symlink it to the new cache location. - if os.environ.get("INSIDE_RE_WORKER") is not None: + if INSIDE_RE_WORKER: from libfb.py.parutil import get_file_path os.environ["FVCORE_CACHE"] = "iopath_cache" diff --git a/tests/implicitron/test_json_index_dataset_provider_v2.py b/tests/implicitron/test_json_index_dataset_provider_v2.py index abec3b9c0..04ea2eaa1 100644 --- a/tests/implicitron/test_json_index_dataset_provider_v2.py +++ b/tests/implicitron/test_json_index_dataset_provider_v2.py @@ -29,6 +29,9 @@ SequenceAnnotation, ) from pytorch3d.implicitron.tools.config import expand_args_fields +from tests.common_testing import interactive_testing_requested + +from .common_resources import CO3DV2_MANIFOLD_PATH class TestJsonIndexDatasetProviderV2(unittest.TestCase): @@ -175,3 +178,16 @@ def _make_random_json_dataset_map_provider_v2_data( with open(os.path.join(root, "category_to_subset_name_list.json"), "w") as f: json.dump(category_to_subset_list, f) + + +class TestCo3dv2(unittest.TestCase): + def test_simple(self): + if not interactive_testing_requested(): + return + dataset_provider = JsonIndexDatasetMapProviderV2( + category="apple", + subset_name="manyview_dev_0", + dataset_root=CO3DV2_MANIFOLD_PATH, + dataset_JsonIndexDataset_args={"load_point_clouds": True}, + ) + dataset_provider.get_dataset_map().train[0]