API for accessing frames in order in Implicitron dataset.

Summary: We often want to iterate over frames in the sequence in temporal order. This diff provides the API to do that. `seq_to_idx` should probably be considered to have `protected` visibility.

Reviewed By: davnov134

Differential Revision: D35012121

fbshipit-source-id: 41896672ec35cd62f3ed4be3aa119efd33adada1
This commit is contained in:
Roman Shapovalov
2022-03-24 05:33:25 -07:00
committed by Facebook GitHub Bot
parent 05f656c01f
commit c0bb49b5f6
3 changed files with 56 additions and 37 deletions

View File

@@ -9,6 +9,7 @@ import unittest
from collections import defaultdict
from dataclasses import dataclass
from pytorch3d.implicitron.dataset.implicitron_dataset import ImplicitronDatasetBase
from pytorch3d.implicitron.dataset.scene_batch_sampler import SceneBatchSampler
@@ -18,7 +19,7 @@ class MockFrameAnnotation:
frame_timestamp: float = 0.0
class MockDataset:
class MockDataset(ImplicitronDatasetBase):
def __init__(self, num_seq, max_frame_gap=1):
"""
Makes a gap of max_frame_gap frame numbers in the middle of each sequence