From 3d886c32d513c14ae6f178e1daa79f5952334881 Mon Sep 17 00:00:00 2001 From: Roman Shapovalov Date: Wed, 14 Jun 2023 08:48:14 -0700 Subject: [PATCH] Fix: subset filter in DatasetBase implementation Summary: In D42739669, I forgot to update the API of existing implementations of DatasetBase to take `subset_filter`. Looks like only one was missing. Reviewed By: bottler Differential Revision: D46724488 fbshipit-source-id: 13ab7a457f853278cf06955aad0cc2bab5fbcce6 --- pytorch3d/implicitron/dataset/single_sequence_dataset.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/pytorch3d/implicitron/dataset/single_sequence_dataset.py b/pytorch3d/implicitron/dataset/single_sequence_dataset.py index 16972c6c..1090faa1 100644 --- a/pytorch3d/implicitron/dataset/single_sequence_dataset.py +++ b/pytorch3d/implicitron/dataset/single_sequence_dataset.py @@ -9,7 +9,7 @@ # provide data for a single scene. from dataclasses import field -from typing import Iterable, Iterator, List, Optional, Tuple +from typing import Iterable, Iterator, List, Optional, Sequence, Tuple import numpy as np import torch @@ -47,13 +47,12 @@ class SingleSceneDataset(DatasetBase, Configurable): def __len__(self) -> int: return len(self.poses) - # pyre-fixme[14]: `sequence_frames_in_order` overrides method defined in - # `DatasetBase` inconsistently. def sequence_frames_in_order( - self, seq_name: str + self, seq_name: str, subset_filter: Optional[Sequence[str]] = None ) -> Iterator[Tuple[float, int, int]]: for i in range(len(self)): - yield (0.0, i, i) + if subset_filter is None or self.frame_types[i] in subset_filter: + yield 0.0, i, i def __getitem__(self, index) -> FrameData: if index >= len(self):