Skip to content

Commit 4b1582e

Browse files
committed
Add test to seek to negative pts
1 parent 2640fa9 commit 4b1582e

File tree

2 files changed

+23
-2
lines changed

2 files changed

+23
-2
lines changed

src/torchcodec/decoders/_core/VideoDecoder.cpp

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@ double ptsToSeconds(int64_t pts, const AVRational& timeBase) {
3434
return ptsToSeconds(pts, timeBase.den);
3535
}
3636

37+
int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
38+
return static_cast<int64_t>(std::round(seconds * timeBase.den));
39+
}
40+
3741
struct AVInput {
3842
UniqueAVFormatContext formatContext;
3943
std::unique_ptr<AVIOBytesContext> ioBytesContext;
@@ -663,7 +667,7 @@ void VideoDecoder::maybeSeekToBeforeDesiredPts() {
663667
for (int streamIndex : activeStreamIndices_) {
664668
StreamInfo& streamInfo = streams_[streamIndex];
665669
// clang-format off: clang format clashes
666-
streamInfo.discardFramesBeforePts = *maybeDesiredPts_ * streamInfo.timeBase.den;
670+
streamInfo.discardFramesBeforePts = secondsToClosestPts(*maybeDesiredPts_, streamInfo.timeBase);
667671
// clang-format on
668672
}
669673

@@ -686,7 +690,8 @@ void VideoDecoder::maybeSeekToBeforeDesiredPts() {
686690
}
687691
int firstActiveStreamIndex = *activeStreamIndices_.begin();
688692
const auto& firstStreamInfo = streams_[firstActiveStreamIndex];
689-
int64_t desiredPts = *maybeDesiredPts_ * firstStreamInfo.timeBase.den;
693+
int64_t desiredPts =
694+
secondsToClosestPts(*maybeDesiredPts_, firstStreamInfo.timeBase);
690695

691696
// For some encodings like H265, FFMPEG sometimes seeks past the point we
692697
// set as the max_ts. So we use our own index to give it the exact pts of
@@ -696,6 +701,7 @@ void VideoDecoder::maybeSeekToBeforeDesiredPts() {
696701
if (!firstStreamInfo.keyFrames.empty()) {
697702
int desiredKeyFrameIndex =
698703
getKeyFrameIndexForPts(firstStreamInfo, desiredPts);
704+
desiredKeyFrameIndex = std::max(desiredKeyFrameIndex, 0);
699705
desiredPts = firstStreamInfo.keyFrames[desiredKeyFrameIndex].pts;
700706
}
701707

test/decoders/test_video_decoder_ops.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,21 @@ def test_seek_and_next(self, device):
8484
)
8585
frame_compare_function(frame_time6, reference_frame_time6.to(device))
8686

87+
@pytest.mark.parametrize("device", cpu_and_cuda())
88+
def test_seek_to_negative_pts(self, device):
89+
device = "cpu"
90+
decoder = create_from_file(str(NASA_VIDEO.path))
91+
scan_all_streams_to_update_metadata(decoder)
92+
add_video_stream(decoder, device=device)
93+
frame_compare_function = get_frame_compare_function(device)
94+
frame0, _, _ = get_next_frame(decoder)
95+
reference_frame0 = NASA_VIDEO.get_frame_data_by_index(0)
96+
frame_compare_function(frame0, reference_frame0.to(device))
97+
98+
seek_to_pts(decoder, -1e-4)
99+
frame0, _, _ = get_next_frame(decoder)
100+
frame_compare_function(frame0, reference_frame0.to(device))
101+
87102
@pytest.mark.parametrize("device", cpu_and_cuda())
88103
def test_get_frame_at_pts(self, device):
89104
decoder = create_from_file(str(NASA_VIDEO.path))

0 commit comments

Comments
 (0)