Skip to content

Commit 3a0089a

Browse files
kongchen1992meta-codesync[bot]
authored andcommitted
{Feature} Core - Add random-access P-frame decoding support in ImageSensorPlayer
Summary: Adds random-access support for P-frame video streams in VrsDataProvider. Previously, getImageDataByIndex() only worked correctly for sequential reads or I-frames. Reading a mid-GOP P-frame out of order returned incorrect pixels because the decoder's DPB (Decoded Picture Buffer) lacked the required reference frames. This adds a recordReadComplete() override in ImageSensorPlayer that detects isMissingFrames() and walks back to the preceding keyframe via readMissingFrames(), then replays through the target frame. The user callback fires exactly once per getImageDataByIndex() call, even when the replay decodes multiple intermediate frames. Enables random-access video frame retrieval for training data pipelines, video editors, and any workflow that needs non-sequential frame access. allow-large-files Reviewed By: robertl0 Differential Revision: D101928539 fbshipit-source-id: 9af8cb30a11aacc465e4c51451e430b5c63e1cfc
1 parent 16a3468 commit 3a0089a

4 files changed

Lines changed: 198 additions & 3 deletions

File tree

core/data_provider/players/ImageSensorPlayer.cpp

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include <vrs/RecordFormat.h>
2525

2626
namespace projectaria::tools::data_provider {
27+
2728
std::optional<projectaria::tools::image::ImageVariant> ImageData::imageVariant() const {
2829
if (pixelFrame->getSpec().getImageFormat() == vrs::ImageFormat::JPG) {
2930
std::shared_ptr<vrs::utils::PixelFrame> normalizedFrame;
@@ -126,11 +127,16 @@ bool ImageSensorPlayer::onImageRead(
126127
success = handleNormalImageProcessing(r, cb, imageSpec);
127128
}
128129

129-
if (success) {
130-
invokeCallbackAndCache();
131-
} else {
130+
if (!success) {
132131
return false;
133132
}
133+
// During a P-frame replay walk (triggered from recordReadComplete), let data_/dataRecord_
134+
// update naturally so the final replayed record (the user-visible target frame) leaves them
135+
// in the correct state, but suppress the user callback and dedup cache update — those fire
136+
// exactly once for the target from recordReadComplete().
137+
if (!whileReadingMissingFrames()) {
138+
invokeCallbackAndCache();
139+
}
134140

135141
if (verbose_) {
136142
fmt::print(
@@ -207,4 +213,21 @@ void ImageSensorPlayer::invokeCallbackAndCache() {
207213
cachedCaptureTimestampNs_ = dataRecord_.captureTimestampNs;
208214
}
209215

216+
int ImageSensorPlayer::recordReadComplete(
217+
vrs::RecordFileReader& fileReader,
218+
const vrs::IndexRecord::RecordInfo& recordInfo) {
219+
// Guard against recursive calls during readMissingFrames replay.
220+
if (whileReadingMissingFrames()) {
221+
return 0;
222+
}
223+
if (!getVideoFrameHandler(streamId_).isMissingFrames()) {
224+
return 0;
225+
}
226+
int result = readMissingFrames(fileReader, recordInfo, /*exactFrame=*/true);
227+
if (result == 0 && data_.isValid()) {
228+
invokeCallbackAndCache();
229+
}
230+
return result;
231+
}
232+
210233
} // namespace projectaria::tools::data_provider

core/data_provider/players/ImageSensorPlayer.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,13 @@ class ImageSensorPlayer : public vrs::utils::VideoRecordFormatStreamPlayer {
162162
override;
163163
bool onImageRead(const vrs::CurrentRecord& r, size_t /*idx*/, const vrs::ContentBlock& cb)
164164
override;
165+
// Random-access P-frame fix. When a mid-GOP read sets isMissingFrames(), walk back to the
166+
// keyframe and replay until the target frame is decoded. Reference implementation lives behind
167+
// #if 0 in vrs/utils/VideoRecordFormatStreamPlayer.h:87-95; if VRS ever
168+
// enables it upstream this override becomes a trivial delete.
169+
int recordReadComplete(
170+
vrs::RecordFileReader& fileReader,
171+
const vrs::IndexRecord::RecordInfo& recordInfo) override;
165172

166173
// Helper methods for different image processing modes
167174
bool handleEmptyFrameMode(const vrs::ContentBlock& cb);
Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
"""
17+
Python parity test for random-access P-frame decoding.
18+
19+
Mirrors the C++ ImageSensorPlayer.RandomAccessParity test at the PyBind
20+
boundary. This test fails iff random access from Python returns different
21+
pixels than sequential access.
22+
"""
23+
24+
import os
25+
import random
26+
import unittest
27+
28+
import numpy as np
29+
from projectaria_tools.core import data_provider
30+
from projectaria_tools.core.sensor_data import SensorDataType
31+
32+
TEST_FOLDER = os.getenv("TEST_FOLDER")
33+
TEST_FIXTURE = os.path.join(
34+
os.getenv("TEST_FOLDER_GEN2"), "aria_gen2_unit_test_sequence_with_pframe.vrs"
35+
)
36+
37+
# Per-component LSB drift between independent decodes — covers both lossless
38+
# (where the drift is 0) and YUV/NV12 paths (where decoder rounding can drift
39+
# by 1 LSB). Strictly weaker than the C++ test's split tolerance, but still
40+
# orders of magnitude tighter than the failure mode this test guards against
41+
# (broken P-frame decode returns garbage / undefined buffer content).
42+
PIXEL_TOLERANCE = 1
43+
44+
45+
def sample_indices(num_frames, n=20):
46+
"""Boundary indices + uniformly-spaced interior samples. Mirrors
47+
sampleIndices() in ImageSensorPlayerTest.cpp."""
48+
if num_frames <= 0:
49+
return []
50+
sample = {0}
51+
if num_frames >= 2:
52+
sample.add(1)
53+
sample.add(num_frames - 1)
54+
if num_frames >= 3:
55+
sample.add(num_frames - 2)
56+
if num_frames > 4 and n > 4:
57+
interior = n - 4
58+
for i in range(interior):
59+
idx = int((i + 1) * (num_frames - 2) / (interior + 1))
60+
if 1 < idx < num_frames - 2:
61+
sample.add(idx)
62+
return sorted(sample)
63+
64+
65+
def get_image_array(provider, stream_id, index):
66+
"""Return a deep-copied numpy array for the image at (stream_id, index).
67+
68+
`to_numpy_array()` may return a view into pybind-managed memory whose
69+
lifetime is tied to the ImageData object. We copy so the snapshot survives
70+
the next decode call into the same player.
71+
"""
72+
image_data, _ = provider.get_image_data_by_index(stream_id, index)
73+
if not image_data.is_valid():
74+
return None
75+
return np.array(image_data.to_numpy_array(), copy=True)
76+
77+
78+
class PFrameRandomAccessTest(unittest.TestCase):
79+
def test_random_access_parity(self):
80+
ref_provider = data_provider.create_vrs_data_provider(TEST_FIXTURE)
81+
rand_provider = data_provider.create_vrs_data_provider(TEST_FIXTURE)
82+
self.assertIsNotNone(
83+
ref_provider, f"failed to open ref provider for {TEST_FIXTURE}"
84+
)
85+
self.assertIsNotNone(
86+
rand_provider, f"failed to open rand provider for {TEST_FIXTURE}"
87+
)
88+
89+
stream_ids = ref_provider.get_all_streams()
90+
image_streams_checked = 0
91+
92+
for stream_id in stream_ids:
93+
if ref_provider.get_sensor_data_type(stream_id) != SensorDataType.IMAGE:
94+
continue
95+
image_streams_checked += 1
96+
97+
num_frames = ref_provider.get_num_data(stream_id)
98+
indices = sample_indices(num_frames, n=20)
99+
if not indices:
100+
continue
101+
index_set = set(indices)
102+
103+
print(
104+
f"stream {stream_id}: numFrames={num_frames}, "
105+
f"sampling {len(indices)} indices"
106+
)
107+
108+
# Pass 1 — sequential reference. Walk every frame so the decoder
109+
# state stays clean (no random access on ref_provider). Snapshot
110+
# the sampled indices into owned numpy arrays.
111+
reference = {}
112+
for j in range(num_frames):
113+
arr = get_image_array(ref_provider, stream_id, j)
114+
if j in index_set and arr is not None:
115+
reference[j] = arr
116+
117+
# Pass 2 — random-access reads on a fresh provider in a
118+
# deterministically shuffled order, exercising forward jumps,
119+
# backward jumps, and cross-GOP seeks.
120+
shuffled = list(indices)
121+
rng = random.Random(0xA1B2C3D4)
122+
rng.shuffle(shuffled)
123+
124+
for idx in shuffled:
125+
rand_arr = get_image_array(rand_provider, stream_id, idx)
126+
ref_arr = reference.get(idx)
127+
self.assertIsNotNone(
128+
rand_arr,
129+
f"random-access read returned invalid image at "
130+
f"stream={stream_id} idx={idx}",
131+
)
132+
self.assertIsNotNone(
133+
ref_arr,
134+
f"sequential reference missing for stream={stream_id} idx={idx}",
135+
)
136+
self.assertEqual(
137+
rand_arr.shape,
138+
ref_arr.shape,
139+
f"shape mismatch at stream={stream_id} idx={idx}: "
140+
f"{rand_arr.shape} vs {ref_arr.shape}",
141+
)
142+
self.assertEqual(
143+
rand_arr.dtype,
144+
ref_arr.dtype,
145+
f"dtype mismatch at stream={stream_id} idx={idx}: "
146+
f"{rand_arr.dtype} vs {ref_arr.dtype}",
147+
)
148+
self.assertTrue(
149+
np.allclose(
150+
rand_arr.astype(np.int32),
151+
ref_arr.astype(np.int32),
152+
atol=PIXEL_TOLERANCE,
153+
rtol=0,
154+
),
155+
f"random vs sequential pixel mismatch at "
156+
f"stream={stream_id} idx={idx} "
157+
f"(max abs diff = "
158+
f"{np.max(np.abs(rand_arr.astype(np.int32) - ref_arr.astype(np.int32)))})",
159+
)
160+
161+
self.assertGreater(
162+
image_streams_checked,
163+
0,
164+
"no image streams found in fixture",
165+
)
17.1 MB
Binary file not shown.

0 commit comments

Comments
 (0)