Fix: FrameDataBuilder working with PathManager

Summary: In refactoring, we lost path manager here, which broke manifold storage. Fixing this.

Reviewed By: bottler

Differential Revision: D45574940

fbshipit-source-id: 579349eaa654215a09e057be57b56b46769c986a

pytorch3d/implicitron/dataset/frame_data.py

@@ -450,7 +450,9 @@ class FrameDataBuilderBase(ReplaceableBase, Generic[FrameDataSubtype], ABC):
         self,
         frame_annotation: types.FrameAnnotation,
         sequence_annotation: types.SequenceAnnotation,
+        *,
         load_blobs: bool = True,
+        **kwargs,
     ) -> FrameDataSubtype:
         """An abstract method to build the frame data based on raw frame/sequence
         annotations, load the binary data and adjust them according to the metadata.
@@ -526,7 +528,12 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
                 "Make sure it is set in either FrameDataBuilder or Dataset params."
             )
 
-        if load_any_blob and not os.path.isdir(self.dataset_root):  # pyre-ignore
+        if self.path_manager is None:
+            dataset_root_exists = os.path.isdir(self.dataset_root)  # pyre-ignore
+        else:
+            dataset_root_exists = self.path_manager.isdir(self.dataset_root)
+
+        if load_any_blob and not dataset_root_exists:
             raise ValueError(
                 f"dataset_root is passed but {self.dataset_root} does not exist."
             )
@@ -535,7 +542,9 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
         self,
         frame_annotation: types.FrameAnnotation,
         sequence_annotation: types.SequenceAnnotation,
+        *,
         load_blobs: bool = True,
+        **kwargs,
     ) -> FrameDataSubtype:
         """Builds the frame data based on raw frame/sequence annotations, loads the
         binary data and adjust them according to the metadata. The processing includes:

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -190,6 +190,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             box_crop=self.box_crop,
             box_crop_mask_thr=self.box_crop_mask_thr,
             box_crop_context=self.box_crop_context,
+            path_manager=self.path_manager,
         )
         logger.info(str(self))
 

pytorch3d/implicitron/dataset/sql_dataset.py

@@ -140,6 +140,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):  # pyre-ignore
             ] = self.dataset_root
 
         run_auto_creation(self)
+        self.frame_data_builder.path_manager = self.path_manager
 
         # pyre-ignore
         self._sql_engine = sa.create_engine(f"sqlite:///{self.sqlite_metadata_file}")