From 95a2acf7632241d09baff55e635e1e4904ed7505 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@meta.com>
Date: Sun, 9 Oct 2022 05:06:49 -0700
Subject: [PATCH] Co3Dv2 point cloud fix

Summary: Avoid certain hardcoded paths in co3dv2 data

Reviewed By: davnov134

Differential Revision: D40209309

fbshipit-source-id: 0e83a15baa47d5bd07d2d23c6048cb4522c1ccba
---
 .../implicitron/dataset/json_index_dataset.py   | 17 ++++++++++++++---
 tests/implicitron/common_resources.py           |  9 ++++++---
 .../test_json_index_dataset_provider_v2.py      | 16 ++++++++++++++++
 3 files changed, 36 insertions(+), 6 deletions(-)

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]