return types for dataset_zoo, dataloader_zoo

Summary: Stronger typing for these functions

Reviewed By: shapovalov

Differential Revision: D36170489

fbshipit-source-id: a2104b29dbbbcfcf91ae1d076cd6b0e3d2030c0b
This commit is contained in:
Jeremy Reizenstein
2022-05-13 05:38:14 -07:00
committed by Facebook GitHub Bot
parent 90ab219d88
commit 2c1901522a
5 changed files with 176 additions and 117 deletions

View File

@@ -27,6 +27,7 @@ from pytorch3d.implicitron.dataset.dataset_zoo import dataset_zoo
from pytorch3d.implicitron.dataset.implicitron_dataset import (
FrameData,
ImplicitronDataset,
ImplicitronDatasetBase,
)
from pytorch3d.implicitron.dataset.utils import is_train_frame
from pytorch3d.implicitron.models.base_model import EvaluationMode
@@ -342,7 +343,10 @@ def export_scenes(
model.eval()
# Setup the dataset
dataset = dataset_zoo(**config.dataset_args)[split]
datasets = dataset_zoo(**config.dataset_args)
dataset: Optional[ImplicitronDatasetBase] = getattr(datasets, split, None)
if dataset is None:
raise ValueError(f"{split} dataset not provided")
# iterate over the sequences in the dataset
for sequence_name in dataset.sequence_names():