diff --git a/pytorch3d/implicitron/dataset/frame_data.py b/pytorch3d/implicitron/dataset/frame_data.py index 88455319..9fff7195 100644 --- a/pytorch3d/implicitron/dataset/frame_data.py +++ b/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: diff --git a/pytorch3d/implicitron/dataset/json_index_dataset.py b/pytorch3d/implicitron/dataset/json_index_dataset.py index caa016d2..8caf581d 100644 --- a/pytorch3d/implicitron/dataset/json_index_dataset.py +++ b/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)) diff --git a/pytorch3d/implicitron/dataset/sql_dataset.py b/pytorch3d/implicitron/dataset/sql_dataset.py index 605e8b52..1eb26964 100644 --- a/pytorch3d/implicitron/dataset/sql_dataset.py +++ b/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}")