@@ -824,6 +824,32 @@ def get_reference_frames(start_seconds, stop_seconds):
824
824
frames , get_reference_frames (start_seconds , stop_seconds )
825
825
)
826
826
827
+ @pytest .mark .parametrize ("asset" , (NASA_AUDIO , NASA_AUDIO_MP3 ))
828
+ def test_pts (self , asset ):
829
+ decoder = create_from_file (str (asset .path ), seek_mode = "approximate" )
830
+ add_audio_stream (decoder )
831
+
832
+ for frame_index in range (asset .num_frames ):
833
+ frame_info = asset .get_frame_info (idx = frame_index )
834
+ start_seconds = frame_info .pts_seconds
835
+
836
+ frames , pts_seconds = get_frames_by_pts_in_range_audio (
837
+ decoder , start_seconds = start_seconds , stop_seconds = start_seconds + 1e-3
838
+ )
839
+ torch .testing .assert_close (
840
+ frames , asset .get_frame_data_by_index (frame_index )
841
+ )
842
+
843
+ if asset is NASA_AUDIO_MP3 and frame_index == 0 :
844
+ # TODO This is a bug. The 0.138125 is correct while 0.072 is
845
+ # incorrect, even though it comes from the decoded AVFrame's pts
846
+ # field.
847
+ # See https://github.com/pytorch/torchcodec/issues/553
848
+ assert pts_seconds == 0.072
849
+ assert start_seconds == 0.138125
850
+ else :
851
+ assert pts_seconds == start_seconds
852
+
827
853
828
854
if __name__ == "__main__" :
829
855
pytest .main ()
0 commit comments