Subsets in dataset iterators

Summary: For the new API, filtering iterators over sequences by subsets is quite helpful. The change is backwards compatible.

Reviewed By: bottler

Differential Revision: D42739669

fbshipit-source-id: d150a404aeaf42fd04a81304c63a4cba203f897d

pytorch3d/implicitron/dataset/dataset_base.py

@@ -237,7 +237,7 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         raise NotImplementedError()
 
     def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int]
+        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
     ) -> List[Tuple[int, float]]:
         """
         If the sequences in the dataset are videos rather than
@@ -251,7 +251,9 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         frames.
 
         Args:
-            idx: frame index in self
+            idxs: frame index in self
+            subset_filter: If given, an index in idxs is ignored if the
+                corresponding frame is not in any of the named subsets.
 
         Returns:
             tuple of
@@ -291,7 +293,7 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         return dict(c2seq)
 
     def sequence_frames_in_order(
-        self, seq_name: str
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
     ) -> Iterator[Tuple[float, int, int]]:
         """Returns an iterator over the frame indices in a given sequence.
         We attempt to first sort by timestamp (if they are available),
@@ -308,7 +310,9 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         """
         # pyre-ignore[16]
         seq_frame_indices = self._seq_to_idx[seq_name]
-        nos_timestamps = self.get_frame_numbers_and_timestamps(seq_frame_indices)
+        nos_timestamps = self.get_frame_numbers_and_timestamps(
+            seq_frame_indices, subset_filter
+        )
 
         yield from sorted(
             [
@@ -317,11 +321,13 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
             ]
         )
 
-    def sequence_indices_in_order(self, seq_name: str) -> Iterator[int]:
+    def sequence_indices_in_order(
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
+    ) -> Iterator[int]:
         """Same as `sequence_frames_in_order` but returns the iterator over
         only dataset indices.
         """
-        for _, _, idx in self.sequence_frames_in_order(seq_name):
+        for _, _, idx in self.sequence_frames_in_order(seq_name, subset_filter):
             yield idx
 
     # frame_data_type is the actual type of frames returned by the dataset.

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -888,10 +888,16 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         return self.path_manager.get_local_path(path)
 
     def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int]
+        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
     ) -> List[Tuple[int, float]]:
         out: List[Tuple[int, float]] = []
         for idx in idxs:
+            if (
+                subset_filter is not None
+                and self.frame_annots[idx]["subset"] not in subset_filter
+            ):
+                continue
+
             # pyre-ignore[16]
             frame_annotation = self.frame_annots[idx]["frame_annotation"]
             out.append(