import io import torch import torchaudio from parameterized import parameterized, parameterized_class from torchaudio.io import StreamReader, StreamWriter from torchaudio_unittest.common_utils import ( disabledInCI, get_asset_path, get_image, get_sinusoid, get_wav_data, nested_params, rgb_to_gray, rgb_to_yuv_ccir, save_image, save_wav, skipIfNoFFmpeg, skipIfNoHWAccel, TempDirMixin, TorchaudioTestCase, ) from torio.io._streaming_media_decoder import ( ChunkTensor, OutputAudioStream, OutputVideoStream, SourceAudioStream, SourceStream, SourceVideoStream, ) @skipIfNoFFmpeg class ChunkTensorTest(TorchaudioTestCase): def test_chunktensor(self): """ChunkTensor serves as a replacement of tensor""" data = torch.randn((256, 2)) pts = 16.0 c = ChunkTensor(data, pts) assert c.pts == pts self.assertEqual(c, data) # method sum_ = c.sum() assert isinstance(sum_, torch.Tensor) self.assertEqual(sum_, data.sum()) # function form min_ = torch.min(c) assert isinstance(min_, torch.Tensor) self.assertEqual(min_, torch.min(data)) # attribute t = c.T assert isinstance(t, torch.Tensor) self.assertEqual(t, data.T) # in-place op c[0] = 0 self.assertEqual(c, data) # pass to other C++ code buffer = io.BytesIO() w = StreamWriter(buffer, format="wav") w.add_audio_stream(8000, 2) with w.open(): w.write_audio_chunk(0, c) w.write_audio_chunk(0, c, c.pts) ################################################################################ # Helper decorator and Mixin to duplicate the tests for fileobj _media_source = parameterized_class( ("test_type",), [("str",), ("fileobj",), ("bytes",)], class_name_func=lambda cls, _, params: f'{cls.__name__}_{params["test_type"]}', ) class _MediaSourceMixin: def setUp(self): super().setUp() self.src = None def get_src(self, path): if self.src is not None: raise ValueError("get_src can be called only once.") if self.test_type == "str": self.src = path elif self.test_type == "fileobj": self.src = open(path, "rb") elif self.test_type == "bytes": with open(path, "rb") as f: self.src = f.read() return self.src def tearDown(self): if self.test_type == "fileobj" and self.src is not None: self.src.close() super().tearDown() ################################################################################ @skipIfNoFFmpeg @_media_source class StreamReaderInterfaceTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase): """Test suite for interface behaviors around StreamReader""" def get_src(self, file="nasa_13013.mp4"): return super().get_src(get_asset_path(file)) def test_streamer_invalid_input(self): """StreamReader constructor does not segfault but raise an exception when the input is invalid""" with self.assertRaises(RuntimeError): StreamReader("foobar") @nested_params( [ ("foo",), ( "foo", "bar", ), ], [{}, {"sample_rate": "16000"}], ) def test_streamer_invalide_option(self, invalid_keys, options): """When invalid options are given, StreamReader raises an exception with these keys""" options.update({k: k for k in invalid_keys}) with self.assertRaises(RuntimeError) as ctx: StreamReader(self.get_src(), option=options) assert all(k in str(ctx.exception) for k in invalid_keys) def test_src_info(self): """`get_src_stream_info` properly fetches information""" s = StreamReader(self.get_src()) assert s.num_src_streams == 6 # Note: # Starting from FFmpeg 4.4, audio/video stream metadata # include "vendor_id" ver = torchaudio.utils.ffmpeg_utils.get_versions()["libavutil"] print(ver) major, minor, _ = ver if major >= 57 or (major == 56 and minor >= 70): base_metadata = {"vendor_id": "[0][0][0][0]"} else: base_metadata = {} expected = [ SourceVideoStream( media_type="video", codec="h264", codec_long_name="H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10", format="yuv420p", bit_rate=71925, num_frames=325, bits_per_sample=8, metadata=dict( base_metadata, handler_name="\x1fMainconcept Video Media Handler", language="eng", ), width=320, height=180, frame_rate=25.0, ), SourceAudioStream( media_type="audio", codec="aac", codec_long_name="AAC (Advanced Audio Coding)", format="fltp", bit_rate=72093, num_frames=103, bits_per_sample=0, metadata=dict( base_metadata, handler_name="#Mainconcept MP4 Sound Media Handler", language="eng", ), sample_rate=8000.0, num_channels=2, ), SourceStream( media_type="subtitle", codec="mov_text", codec_long_name="MOV text", format=None, bit_rate=None, num_frames=None, bits_per_sample=None, metadata={ "handler_name": "SubtitleHandler", "language": "eng", }, ), SourceVideoStream( media_type="video", codec="h264", codec_long_name="H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10", format="yuv420p", bit_rate=128783, num_frames=390, bits_per_sample=8, metadata=dict( base_metadata, handler_name="\x1fMainconcept Video Media Handler", language="eng", ), width=480, height=270, frame_rate=29.97002997002997, ), SourceAudioStream( media_type="audio", codec="aac", codec_long_name="AAC (Advanced Audio Coding)", format="fltp", bit_rate=128837, num_frames=205, bits_per_sample=0, metadata=dict( base_metadata, handler_name="#Mainconcept MP4 Sound Media Handler", language="eng", ), sample_rate=16000.0, num_channels=2, ), SourceStream( media_type="subtitle", codec="mov_text", codec_long_name="MOV text", format=None, bit_rate=None, num_frames=None, bits_per_sample=None, metadata={ "handler_name": "SubtitleHandler", "language": "eng", }, ), ] output = [s.get_src_stream_info(i) for i in range(6)] assert expected == output def test_output_info(self): s = StreamReader(self.get_src()) s.add_audio_stream(-1) s.add_audio_stream(-1, filter_desc="aresample=8000") s.add_audio_stream(-1, filter_desc="aformat=sample_fmts=s16p") s.add_video_stream(-1) s.add_video_stream(-1, filter_desc="fps=10") s.add_video_stream(-1, filter_desc="format=rgb24") s.add_video_stream(-1, filter_desc="scale=w=160:h=90") # Note: # Somehow only FFmpeg 5 reports invalid video frame rate. (24576/0) # FFmpeg 4 and 6 work fine. # Perhaps this is a regression in FFmpeg or it could actually originate # from other libraries. # It consistently fails with FFmpeg installed via conda, so we change # the value based on FFmpeg version. ver = torchaudio.utils.ffmpeg_utils.get_versions()["libavutil"] print(ver) major, minor, _ = ver if major == 57: video_frame_rate = -1 else: video_frame_rate = 30000 / 1001 print(video_frame_rate) expected = [ OutputAudioStream( source_index=4, filter_description="anull", media_type="audio", format="fltp", sample_rate=16000.0, num_channels=2, ), OutputAudioStream( source_index=4, filter_description="aresample=8000", media_type="audio", format="fltp", sample_rate=8000.0, num_channels=2, ), OutputAudioStream( source_index=4, filter_description="aformat=sample_fmts=s16p", media_type="audio", format="s16p", sample_rate=16000.0, num_channels=2, ), OutputVideoStream( source_index=3, filter_description="null", media_type="video", format="yuv420p", width=480, height=270, frame_rate=30000 / 1001, ), OutputVideoStream( source_index=3, filter_description="fps=10", media_type="video", format="yuv420p", width=480, height=270, frame_rate=10, ), OutputVideoStream( source_index=3, filter_description="format=rgb24", media_type="video", format="rgb24", width=480, height=270, frame_rate=30000 / 1001, ), OutputVideoStream( source_index=3, filter_description="scale=w=160:h=90", media_type="video", format="yuv420p", width=160, height=90, frame_rate=30000 / 1001, ), ] output = [s.get_out_stream_info(i) for i in range(s.num_out_streams)] assert expected == output def test_id3tag(self): """get_metadata method can fetch id3tag properly""" s = StreamReader(self.get_src("steam-train-whistle-daniel_simon.mp3")) output = s.get_metadata() expected = { "title": "SoundBible.com Must Credit", "artist": "SoundBible.com Must Credit", "date": "2017", } assert output == expected def test_video_metadata(self): """get_metadata method can fetch video metadata""" s = StreamReader(self.get_src()) output = s.get_metadata() expected = { "compatible_brands": "isomiso2avc1mp41", "encoder": "Lavf58.76.100", "major_brand": "isom", "minor_version": "512", } assert output == expected def test_src_info_invalid_index(self): """`get_src_stream_info` does not segfault but raise an exception when input is invalid""" s = StreamReader(self.get_src()) for i in [-1, 6, 7, 8]: with self.assertRaises(RuntimeError): s.get_src_stream_info(i) def test_default_streams(self): """default stream is not None""" s = StreamReader(self.get_src()) assert s.default_audio_stream is not None assert s.default_video_stream is not None def test_default_audio_stream_none(self): """default audio stream is None for video without audio""" s = StreamReader(self.get_src("nasa_13013_no_audio.mp4")) assert s.default_audio_stream is None def test_default_video_stream_none(self): """default video stream is None for video with only audio""" s = StreamReader(self.get_src("nasa_13013_no_video.mp4")) assert s.default_video_stream is None def test_num_out_stream(self): """num_out_streams gives the correct count of output streams""" s = StreamReader(self.get_src()) n, m = 6, 4 for i in range(n): assert s.num_out_streams == i s.add_audio_stream(frames_per_chunk=-1) for i in range(m): assert s.num_out_streams == n - i s.remove_stream(0) for i in range(m): assert s.num_out_streams == n - m + i s.add_video_stream(frames_per_chunk=-1) for i in range(n): assert s.num_out_streams == n - i s.remove_stream(n - i - 1) assert s.num_out_streams == 0 def test_basic_audio_stream(self): """`add_basic_audio_stream` constructs a correct filter.""" s = StreamReader(self.get_src()) s.add_basic_audio_stream(frames_per_chunk=-1, format=None) s.add_basic_audio_stream(frames_per_chunk=-1, sample_rate=8000) s.add_basic_audio_stream(frames_per_chunk=-1, format="s16p") sinfo = s.get_out_stream_info(0) assert sinfo.source_index == s.default_audio_stream assert sinfo.filter_description == "anull" sinfo = s.get_out_stream_info(1) assert sinfo.source_index == s.default_audio_stream assert "aresample=8000" in sinfo.filter_description sinfo = s.get_out_stream_info(2) assert sinfo.source_index == s.default_audio_stream assert "aformat=sample_fmts=s16" in sinfo.filter_description def test_basic_video_stream(self): """`add_basic_video_stream` constructs a correct filter.""" s = StreamReader(self.get_src()) s.add_basic_video_stream(frames_per_chunk=-1, format=None) s.add_basic_video_stream(frames_per_chunk=-1, width=3, height=5) s.add_basic_video_stream(frames_per_chunk=-1, frame_rate=7) s.add_basic_video_stream(frames_per_chunk=-1, format="bgr24") sinfo = s.get_out_stream_info(0) assert sinfo.source_index == s.default_video_stream assert sinfo.filter_description == "null" sinfo = s.get_out_stream_info(1) assert sinfo.source_index == s.default_video_stream assert "scale=width=3:height=5" in sinfo.filter_description sinfo = s.get_out_stream_info(2) assert sinfo.source_index == s.default_video_stream assert "fps=7" in sinfo.filter_description sinfo = s.get_out_stream_info(3) assert sinfo.source_index == s.default_video_stream assert "format=pix_fmts=bgr24" in sinfo.filter_description def test_remove_streams(self): """`remove_stream` removes the correct output stream""" s = StreamReader(self.get_src()) s.add_basic_audio_stream(frames_per_chunk=-1, sample_rate=24000) s.add_basic_video_stream(frames_per_chunk=-1, width=16, height=16) s.add_basic_audio_stream(frames_per_chunk=-1, sample_rate=8000) sinfo = [s.get_out_stream_info(i) for i in range(3)] s.remove_stream(1) del sinfo[1] assert sinfo == [s.get_out_stream_info(i) for i in range(s.num_out_streams)] s.remove_stream(1) del sinfo[1] assert sinfo == [s.get_out_stream_info(i) for i in range(s.num_out_streams)] s.remove_stream(0) del sinfo[0] assert [] == [s.get_out_stream_info(i) for i in range(s.num_out_streams)] def test_remove_stream_invalid(self): """Attempt to remove invalid output streams raises IndexError""" s = StreamReader(self.get_src()) for i in range(-3, 3): with self.assertRaises(RuntimeError): s.remove_stream(i) s.add_audio_stream(frames_per_chunk=-1) for i in range(-3, 3): if i == 0: continue with self.assertRaises(RuntimeError): s.remove_stream(i) def test_process_packet(self): """`process_packet` method returns 0 while there is a packet in source stream""" s = StreamReader(self.get_src()) # nasa_1013.mp3 contains 1023 packets. for _ in range(1023): code = s.process_packet() assert code == 0 # now all the packets should be processed, so process_packet returns 1. code = s.process_packet() assert code == 1 def test_pop_chunks_no_output_stream(self): """`pop_chunks` method returns empty list when there is no output stream""" s = StreamReader(self.get_src()) assert s.pop_chunks() == [] def test_pop_chunks_empty_buffer(self): """`pop_chunks` method returns None when a buffer is empty""" s = StreamReader(self.get_src()) s.add_basic_audio_stream(frames_per_chunk=-1) s.add_basic_video_stream(frames_per_chunk=-1) assert s.pop_chunks() == [None, None] def test_pop_chunks_exhausted_stream(self): """`pop_chunks` method returns None when the source stream is exhausted""" s = StreamReader(self.get_src()) # video is 16.57 seconds. # audio streams per 10 second chunk # video streams per 20 second chunk # The first `pop_chunk` call should return 2 Tensors (10 second audio and 16.57 second video) # The second call should return 1 Tensor (6.57 second audio) and None. # After that, `pop_chunk` should keep returning None-s. s.add_basic_audio_stream(frames_per_chunk=100, sample_rate=10, buffer_chunk_size=3) s.add_basic_video_stream(frames_per_chunk=200, frame_rate=10, buffer_chunk_size=3) s.process_all_packets() chunks = s.pop_chunks() assert chunks[0] is not None assert chunks[1] is not None assert chunks[0].shape[0] == 100 # audio tensor contains 10 second chunk assert chunks[1].shape[0] < 200 # video tensor contains less than 20 second chunk chunks = s.pop_chunks() assert chunks[0] is not None assert chunks[1] is None assert chunks[0].shape[0] < 100 # audio tensor contains less than 10 second chunk for _ in range(10): chunks = s.pop_chunks() assert chunks[0] is None assert chunks[1] is None def test_stream_empty(self): """`stream` fails when no output stream is configured""" s = StreamReader(self.get_src()) with self.assertRaises(RuntimeError): next(s.stream()) def test_stream_smoke_test(self): """`stream` streams chunks fine""" w, h = 256, 198 s = StreamReader(self.get_src()) s.add_basic_audio_stream(frames_per_chunk=2000, sample_rate=8000) s.add_basic_video_stream(frames_per_chunk=15, frame_rate=60, width=w, height=h) for i, (achunk, vchunk) in enumerate(s.stream()): assert achunk.shape == torch.Size([2000, 2]) assert vchunk.shape == torch.Size([15, 3, h, w]) if i >= 40: break def test_stream_requires_grad_false(self): """Tensors produced by StreamReader are requires_grad=False""" s = StreamReader(self.get_src()) s.add_basic_audio_stream(frames_per_chunk=2000) s.add_basic_video_stream(frames_per_chunk=15) s.fill_buffer() audio, video = s.pop_chunks() assert not audio._elem.requires_grad assert not video._elem.requires_grad @parameterized.expand(["key", "any", "precise"]) def test_seek(self, mode): """Calling `seek` multiple times should not segfault""" s = StreamReader(self.get_src()) for i in range(10): s.seek(i, mode) for _ in range(0): s.seek(0, mode) for i in range(10, 0, -1): s.seek(i, mode) def test_seek_negative(self): """Calling `seek` with negative value should raise an exception""" s = StreamReader(self.get_src()) with self.assertRaises(RuntimeError): s.seek(-1.0) def test_seek_invalid_mode(self): """Calling `seek` with an invalid model should raise an exception""" s = StreamReader(self.get_src()) with self.assertRaises(ValueError): s.seek(10, "magic_seek") @parameterized.expand( [ # Test keyframe seek # The source mp4 video has two key frames the first frame and 203rd frame at 8.08 second. # If the seek time stamp is smaller than 8.08, it will seek into the first frame at 0.0 second. ("nasa_13013.mp4", "key", 0.2, (0, slice(None))), ("nasa_13013.mp4", "key", 8.04, (0, slice(None))), ("nasa_13013.mp4", "key", 8.08, (0, slice(202, None))), ("nasa_13013.mp4", "key", 8.12, (0, slice(202, None))), # The source avi video has one keyframe every twelve frames 0, 12, 24,.. or every 0.4004 seconds. # if we seek to a time stamp smaller than 0.4004 it will seek into the first frame at 0.0 second. ("nasa_13013.avi", "key", 0.2, (0, slice(None))), ("nasa_13013.avi", "key", 1.01, (0, slice(24, None))), ("nasa_13013.avi", "key", 7.37, (0, slice(216, None))), ("nasa_13013.avi", "key", 7.7, (0, slice(216, None))), # Test precise seek ("nasa_13013.mp4", "precise", 0.0, (0, slice(None))), ("nasa_13013.mp4", "precise", 0.2, (0, slice(5, None))), ("nasa_13013.mp4", "precise", 8.04, (0, slice(201, None))), ("nasa_13013.mp4", "precise", 8.08, (0, slice(202, None))), ("nasa_13013.mp4", "precise", 8.12, (0, slice(203, None))), ("nasa_13013.avi", "precise", 0.0, (0, slice(None))), ("nasa_13013.avi", "precise", 0.2, (0, slice(1, None))), ("nasa_13013.avi", "precise", 8.1, (0, slice(238, None))), ("nasa_13013.avi", "precise", 8.14, (0, slice(239, None))), ("nasa_13013.avi", "precise", 8.17, (0, slice(240, None))), # Test precise seek on video with missing PTS ("RATRACE_wave_f_nm_np1_fr_goo_37.avi", "precise", 0.0, (0, slice(None))), ("RATRACE_wave_f_nm_np1_fr_goo_37.avi", "precise", 0.2, (0, slice(4, None))), ("RATRACE_wave_f_nm_np1_fr_goo_37.avi", "precise", 0.3, (0, slice(7, None))), # Test any seek # The source avi video has one keyframe every twelve frames 0, 12, 24,.. or every 0.4004 seconds. ("nasa_13013.avi", "any", 0.0, (0, slice(None))), ("nasa_13013.avi", "any", 0.56, (0, slice(12, None))), ("nasa_13013.avi", "any", 7.77, (0, slice(228, None))), ("nasa_13013.avi", "any", 0.2002, (11, slice(12, None))), ("nasa_13013.avi", "any", 0.233567, (10, slice(12, None))), ("nasa_13013.avi", "any", 0.266933, (9, slice(12, None))), ] ) def test_seek_modes(self, src, mode, seek_time, ref_indices): """We expect the following behaviour from the diferent kinds of seek: - `key`: the reader will seek to the first keyframe from the timestamp given - `precise`: the reader will seek to the first keyframe from the timestamp given and start decoding from that position until the given timestmap (discarding all frames in between) - `any`: the reader will seek to the colsest frame to the timestamp given but if this is not a keyframe, the content will be the delta from other frames To thest this behaviour we can parameterize the test with the tupple ref_indices. ref_indices[0] is the expected index on the frames list decoded after seek and ref_indices[1] is exepected index for the list of all frames decoded from the begining (reference frames). This test checks if the reference frame at index ref_indices[1] is the same as ref_indices[0]. Plese note that with `any` and `key` seek we only compare keyframes, but with `precise` seek we can compare any frame content. """ # Using the first video stream (which is not default video stream) stream_index = 0 # Decode all frames for reference src_bin = self.get_src(src) s = StreamReader(src_bin) s.add_basic_video_stream(-1, stream_index=stream_index) s.process_all_packets() (ref_frames,) = s.pop_chunks() s.seek(seek_time, mode=mode) s.process_all_packets() (frame,) = s.pop_chunks() hyp_index, ref_index = ref_indices hyp, ref = frame[hyp_index:], ref_frames[ref_index] print(hyp.shape, ref.shape) self.assertEqual(hyp, ref) @parameterized.expand( [ ("nasa_13013.mp4", [195, 3, 270, 480]), # RATRACE does not have valid PTS metadata. ("RATRACE_wave_f_nm_np1_fr_goo_37.avi", [36, 3, 240, 560]), ] ) def test_change_fps(self, src, shape): """Can change the FPS of videos""" tgt_frame_rate = 15 s = StreamReader(self.get_src(src)) info = s.get_src_stream_info(s.default_video_stream) assert info.frame_rate != tgt_frame_rate s.add_basic_video_stream(frames_per_chunk=-1, frame_rate=tgt_frame_rate) s.process_all_packets() (chunk,) = s.pop_chunks() assert chunk.shape == torch.Size(shape) def test_invalid_chunk_option(self): """Passing invalid `frames_per_chunk` and `buffer_chunk_size` raises error""" s = StreamReader(self.get_src()) for fpc, bcs in ((0, 3), (3, 0), (-2, 3), (3, -2)): with self.assertRaises(RuntimeError): s.add_audio_stream(frames_per_chunk=fpc, buffer_chunk_size=bcs) with self.assertRaises(RuntimeError): s.add_video_stream(frames_per_chunk=fpc, buffer_chunk_size=bcs) def test_unchunked_stream(self): """`frames_per_chunk=-1` disable chunking. When chunking is disabled, frames contained in one AVFrame become one chunk. For video, that is always one frame, but for audio, it depends. """ s = StreamReader(self.get_src()) s.add_video_stream(frames_per_chunk=-1, buffer_chunk_size=10000) s.add_audio_stream(frames_per_chunk=-1, buffer_chunk_size=10000) s.process_all_packets() video, audio = s.pop_chunks() assert video.shape == torch.Size([390, 3, 270, 480]) assert audio.shape == torch.Size([208896, 2]) @parameterized.expand([(1,), (3,), (5,), (10,)]) def test_frames_per_chunk(self, fpc): """Changing frames_per_chunk does not change the returned content""" src = self.get_src() s = StreamReader(src) s.add_video_stream(frames_per_chunk=-1, buffer_chunk_size=-1) s.add_audio_stream(frames_per_chunk=-1, buffer_chunk_size=-1) s.process_all_packets() ref_video, ref_audio = s.pop_chunks() if self.test_type == "fileobj": src.seek(0) s = StreamReader(src) s.add_video_stream(frames_per_chunk=fpc, buffer_chunk_size=-1) s.add_audio_stream(frames_per_chunk=fpc, buffer_chunk_size=-1) chunks = list(s.stream()) video_chunks = torch.cat([c[0] for c in chunks if c[0] is not None]) audio_chunks = torch.cat([c[1] for c in chunks if c[1] is not None]) self.assertEqual(ref_video, video_chunks) self.assertEqual(ref_audio, audio_chunks) def test_buffer_chunk_size(self): """`buffer_chunk_size=-1` does not drop frames.""" src = self.get_src() s = StreamReader(src) s.add_video_stream(frames_per_chunk=30, buffer_chunk_size=-1) s.add_audio_stream(frames_per_chunk=16000, buffer_chunk_size=-1) s.process_all_packets() for _ in range(13): video, audio = s.pop_chunks() assert video.shape == torch.Size([30, 3, 270, 480]) assert audio.shape == torch.Size([16000, 2]) video, audio = s.pop_chunks() assert video is None assert audio.shape == torch.Size([896, 2]) if self.test_type == "fileobj": src.seek(0) s = StreamReader(src) s.add_video_stream(frames_per_chunk=30, buffer_chunk_size=3) s.add_audio_stream(frames_per_chunk=16000, buffer_chunk_size=3) s.process_all_packets() for _ in range(2): video, audio = s.pop_chunks() assert video.shape == torch.Size([30, 3, 270, 480]) assert audio.shape == torch.Size([16000, 2]) video, audio = s.pop_chunks() assert video.shape == torch.Size([30, 3, 270, 480]) assert audio.shape == torch.Size([896, 2]) @parameterized.expand([(1,), (3,), (5,), (10,)]) def test_video_pts(self, fpc): """PTS values of the first frame are reported in .pts attribute""" rate, num_frames = 30000 / 1001, 390 ref_pts = [i / rate for i in range(0, num_frames, fpc)] s = StreamReader(self.get_src()) s.add_video_stream(fpc) pts = [video.pts for video, in s.stream()] self.assertEqual(pts, ref_pts) @parameterized.expand([(256,), (512,), (1024,), (4086,)]) def test_audio_pts(self, fpc): """PTS values of the first frame are reported in .pts attribute""" rate, num_frames = 16000, 208896 ref_pts = [i / rate for i in range(0, num_frames, fpc)] s = StreamReader(self.get_src()) s.add_audio_stream(fpc, buffer_chunk_size=-1) pts = [audio.pts for audio, in s.stream()] self.assertEqual(pts, ref_pts) def test_pts_unchunked_process_all(self): """PTS is zero when loading the entire media with unchunked buffer""" s = StreamReader(self.get_src()) s.add_audio_stream(-1, buffer_chunk_size=-1) s.add_video_stream(-1, buffer_chunk_size=-1) s.process_all_packets() audio, video = s.pop_chunks() assert audio.pts == 0.0 assert video.pts == 0.0 assert audio.size(0) == 208896 assert video.size(0) == 390 def test_pts_unchunked(self): """PTS grows proportionally to the number of frames decoded""" s = StreamReader(self.get_src()) s.add_audio_stream(-1, buffer_chunk_size=-1) s.add_video_stream(-1, buffer_chunk_size=-1) num_audio_frames, num_video_frames = 0, 0 while num_audio_frames < 208896 and num_video_frames < 390: s.process_packet() audio, video = s.pop_chunks() if audio is None and video is None: continue if audio is not None: assert audio.pts == num_audio_frames / 16000 num_audio_frames += audio.size(0) if video is not None: assert video.pts == num_video_frames * 1001 / 30000 num_video_frames += video.size(0) def _to_fltp(original): """Convert Tensor to float32 with value range [-1, 1]""" denom = { torch.uint8: 2**7, torch.int16: 2**15, torch.int32: 2**31, }[original.dtype] fltp = original.to(torch.float32) if original.dtype == torch.uint8: fltp -= 128 fltp /= denom return fltp @skipIfNoFFmpeg @_media_source class StreamReaderAudioTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase): """Test suite for audio streaming""" def _get_reference_wav(self, sample_rate, channels_first=False, **kwargs): data = get_wav_data(**kwargs, normalize=False, channels_first=channels_first) path = self.get_temp_path("ref.wav") save_wav(path, data, sample_rate, channels_first=channels_first) return path, data def get_src(self, *args, **kwargs): path, data = self._get_reference_wav(*args, **kwargs) src = super().get_src(path) return src, data def _test_wav(self, src, original, fmt): s = StreamReader(src) s.add_basic_audio_stream(frames_per_chunk=-1, format=fmt) s.process_all_packets() (output,) = s.pop_chunks() self.assertEqual(original, output) @nested_params( ["int16", "uint8", "int32"], # "float", "double", "int64"] [1, 2, 4, 8], ) def test_basic_audio_stream(self, dtype, num_channels): """`basic_audio_stream` can load WAV file properly.""" src, original = self.get_src(8000, dtype=dtype, num_channels=num_channels) fmt = { "uint8": "u8p", "int16": "s16p", "int32": "s32p", }[dtype] # provide the matching dtype self._test_wav(src, original, fmt=fmt) # use the internal dtype ffmpeg picks if self.test_type == "fileobj": src.seek(0) self._test_wav(src, original, fmt=None) def test_audio_stream_format(self): "`format` argument properly changes the sample format of decoded audio" num_channels = 2 src, s32 = self.get_src(8000, dtype="int32", num_channels=num_channels) args = { "num_channels": num_channels, "normalize": False, "channels_first": False, "num_frames": 1 << 16, } u8 = get_wav_data("uint8", **args) s16 = get_wav_data("int16", **args) s64 = s32.to(torch.int64) * (1 << 32) f32 = get_wav_data("float32", **args) f64 = get_wav_data("float64", **args) s = StreamReader(src) s.add_basic_audio_stream(frames_per_chunk=-1, format="u8") s.add_basic_audio_stream(frames_per_chunk=-1, format="u8p") s.add_basic_audio_stream(frames_per_chunk=-1, format="s16") s.add_basic_audio_stream(frames_per_chunk=-1, format="s16p") s.add_basic_audio_stream(frames_per_chunk=-1, format="s32") s.add_basic_audio_stream(frames_per_chunk=-1, format="s32p") s.add_basic_audio_stream(frames_per_chunk=-1, format="s64") s.add_basic_audio_stream(frames_per_chunk=-1, format="s64p") s.add_basic_audio_stream(frames_per_chunk=-1, format="flt") s.add_basic_audio_stream(frames_per_chunk=-1, format="fltp") s.add_basic_audio_stream(frames_per_chunk=-1, format="dbl") s.add_basic_audio_stream(frames_per_chunk=-1, format="dblp") s.process_all_packets() chunks = s.pop_chunks() self.assertEqual(chunks[0], u8, atol=1, rtol=0) self.assertEqual(chunks[1], u8, atol=1, rtol=0) self.assertEqual(chunks[2], s16) self.assertEqual(chunks[3], s16) self.assertEqual(chunks[4], s32) self.assertEqual(chunks[5], s32) self.assertEqual(chunks[6], s64) self.assertEqual(chunks[7], s64) self.assertEqual(chunks[8], f32) self.assertEqual(chunks[9], f32) self.assertEqual(chunks[10], f64) self.assertEqual(chunks[11], f64) @nested_params([4000, 16000]) def test_basic_audio_stream_sample_rate(self, sr): """`sample_rate` argument changes the sample_rate of decoded audio""" src_num_channels, src_sr = 2, 8000 data = get_sinusoid(sample_rate=src_sr, n_channels=src_num_channels, channels_first=False) path = self.get_temp_path("ref.wav") save_wav(path, data, src_sr, channels_first=False) s = StreamReader(path) s.add_basic_audio_stream(frames_per_chunk=-1, format="flt", sample_rate=sr) self.assertEqual(s.get_src_stream_info(0).sample_rate, src_sr) self.assertEqual(s.get_out_stream_info(0).sample_rate, sr) s.process_all_packets() (chunks,) = s.pop_chunks() self.assertEqual(chunks.shape, [sr, src_num_channels]) @nested_params([1, 2, 3, 8, 16]) def test_basic_audio_stream_num_channels(self, num_channels): """`sample_rate` argument changes the number of channels of decoded audio""" src_num_channels, sr = 2, 8000 data = get_sinusoid(sample_rate=sr, n_channels=src_num_channels, channels_first=False) path = self.get_temp_path("ref.wav") save_wav(path, data, sr, channels_first=False) s = StreamReader(path) s.add_basic_audio_stream(frames_per_chunk=-1, format="flt", num_channels=num_channels) self.assertEqual(s.get_src_stream_info(0).num_channels, src_num_channels) self.assertEqual(s.get_out_stream_info(0).num_channels, num_channels) s.process_all_packets() (chunks,) = s.pop_chunks() self.assertEqual(chunks.shape, [sr, num_channels]) @nested_params( ["int16", "uint8", "int32"], # "float", "double", "int64"] [1, 2, 4, 8], ) def test_audio_stream(self, dtype, num_channels): """`add_audio_stream` can apply filter""" src, original = self.get_src(8000, dtype=dtype, num_channels=num_channels) expected = torch.flip(original, dims=(0,)) s = StreamReader(src) s.add_audio_stream(frames_per_chunk=-1, filter_desc="areverse") s.process_all_packets() (output,) = s.pop_chunks() self.assertEqual(expected, output) @nested_params( ["int16", "uint8", "int32"], # "float", "double", "int64"] [1, 2, 4, 8], ) def test_audio_seek(self, dtype, num_channels): """`seek` changes the position properly""" src, original = self.get_src(1, dtype=dtype, num_channels=num_channels, num_frames=30) for t in range(10, 20): expected = original[t:, :] if self.test_type == "fileobj": src.seek(0) s = StreamReader(src) s.add_audio_stream(frames_per_chunk=-1) s.seek(float(t)) s.process_all_packets() (output,) = s.pop_chunks() self.assertEqual(expected, output) def test_audio_seek_multiple(self): """Calling `seek` after streaming is started should change the position properly""" src, original = self.get_src(1, dtype="int16", num_channels=2, num_frames=30) s = StreamReader(src) s.add_audio_stream(frames_per_chunk=-1) ts = list(range(20)) + list(range(20, 0, -1)) + list(range(20)) for t in ts: s.seek(float(t)) s.process_all_packets() (output,) = s.pop_chunks() expected = original[t:, :] self.assertEqual(expected, output) @nested_params( [ (18, 6, 3), # num_frames is divisible by frames_per_chunk (18, 5, 4), # num_frames is not divisible by frames_per_chunk (18, 32, 1), # num_frames is shorter than frames_per_chunk ], [1, 2, 4, 8], ) def test_audio_frames_per_chunk(self, frame_param, num_channels): """Different chunk parameter covers the source media properly""" num_frames, frames_per_chunk, buffer_chunk_size = frame_param src, original = self.get_src( 8000, dtype="int16", num_channels=num_channels, num_frames=num_frames, channels_first=False ) s = StreamReader(src) s.add_audio_stream(frames_per_chunk=frames_per_chunk, buffer_chunk_size=buffer_chunk_size) i, outputs = 0, [] for (output,) in s.stream(): expected = original[frames_per_chunk * i : frames_per_chunk * (i + 1), :] outputs.append(output) self.assertEqual(expected, output) i += 1 assert i == num_frames // frames_per_chunk + (1 if num_frames % frames_per_chunk else 0) self.assertEqual(torch.cat(outputs, 0), original) @skipIfNoFFmpeg @_media_source class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase): def _get_reference_png(self, width: int, height: int, grayscale: bool): original = get_image(width, height, grayscale=grayscale) path = self.get_temp_path("ref.png") save_image(path, original, mode="L" if grayscale else "RGB") return path, original def get_src(self, *args, **kwargs): path, data = self._get_reference_png(*args, **kwargs) src = super().get_src(path) return src, data def _test_png(self, path, original, format=None): s = StreamReader(path) s.add_basic_video_stream(frames_per_chunk=-1, format=format) s.process_all_packets() (output,) = s.pop_chunks() self.assertEqual(original, output) @nested_params([True, False]) def test_png(self, grayscale): # TODO: # Add test with alpha channel (RGBA, ARGB, BGRA, ABGR) w, h = 32, 18 src, original = self.get_src(w, h, grayscale=grayscale) expected = original[None, ...] self._test_png(src, expected) @parameterized.expand( [ ("hflip", 2), ("vflip", 1), ] ) def test_png_effect(self, filter_desc, index): h, w = 111, 250 src, original = self.get_src(w, h, grayscale=False) expected = torch.flip(original, dims=(index,))[None, ...] s = StreamReader(src) s.add_video_stream(frames_per_chunk=-1, filter_desc=filter_desc) s.process_all_packets() output = s.pop_chunks()[0] print("expected", expected) print("output", output) self.assertEqual(expected, output) def test_png_yuv_read_out(self): """Providing format prpoerly change the color space""" rgb = torch.empty(1, 3, 256, 256, dtype=torch.uint8) rgb[0, 0] = torch.arange(256, dtype=torch.uint8).reshape([1, -1]) rgb[0, 1] = torch.arange(256, dtype=torch.uint8).reshape([-1, 1]) alpha = torch.full((1, 1, 256, 256), 255, dtype=torch.uint8) for i in range(256): rgb[0, 2] = i path = self.get_temp_path(f"ref_{i}.png") save_image(path, rgb[0], mode="RGB") rgb16 = ((rgb.to(torch.int32) - 128) << 8).to(torch.int16) yuv = rgb_to_yuv_ccir(rgb) yuv16 = yuv.to(torch.int16) * 4 bgr = rgb[:, [2, 1, 0], :, :] gray = rgb_to_gray(rgb) argb = torch.cat([alpha, rgb], dim=1) rgba = torch.cat([rgb, alpha], dim=1) abgr = torch.cat([alpha, bgr], dim=1) bgra = torch.cat([bgr, alpha], dim=1) s = StreamReader(path) s.add_basic_video_stream(frames_per_chunk=-1, format="yuv444p") s.add_basic_video_stream(frames_per_chunk=-1, format="yuv420p") s.add_basic_video_stream(frames_per_chunk=-1, format="nv12") s.add_basic_video_stream(frames_per_chunk=-1, format="rgb24") s.add_basic_video_stream(frames_per_chunk=-1, format="bgr24") s.add_basic_video_stream(frames_per_chunk=-1, format="gray8") s.add_basic_video_stream(frames_per_chunk=-1, format="rgb48le") s.add_basic_video_stream(frames_per_chunk=-1, format="argb") s.add_basic_video_stream(frames_per_chunk=-1, format="rgba") s.add_basic_video_stream(frames_per_chunk=-1, format="abgr") s.add_basic_video_stream(frames_per_chunk=-1, format="bgra") s.add_basic_video_stream(frames_per_chunk=-1, format="yuv420p10le") s.process_all_packets() chunks = s.pop_chunks() self.assertEqual(chunks[0], yuv, atol=1, rtol=0) self.assertEqual(chunks[1], yuv, atol=1, rtol=0) self.assertEqual(chunks[2], yuv, atol=1, rtol=0) self.assertEqual(chunks[3], rgb, atol=0, rtol=0) self.assertEqual(chunks[4], bgr, atol=0, rtol=0) self.assertEqual(chunks[5], gray, atol=1, rtol=0) self.assertEqual(chunks[6], rgb16, atol=256, rtol=0) self.assertEqual(chunks[7], argb, atol=0, rtol=0) self.assertEqual(chunks[8], rgba, atol=0, rtol=0) self.assertEqual(chunks[9], abgr, atol=0, rtol=0) self.assertEqual(chunks[10], bgra, atol=0, rtol=0) self.assertEqual(chunks[11], yuv16, atol=4, rtol=0) @skipIfNoHWAccel("h264_cuvid") class CuvidHWAccelInterfaceTest(TorchaudioTestCase): def test_dup_hw_acel(self): """Specifying the same source stream with and without HW accel should fail (instead of segfault later)""" src = get_asset_path("nasa_13013.mp4") r = StreamReader(src) r.add_video_stream(-1, decoder="h264_cuvid") with self.assertRaises(RuntimeError): r.add_video_stream(-1, decoder="h264_cuvid", hw_accel="cuda") r = StreamReader(src) r.add_video_stream(-1, decoder="h264_cuvid", hw_accel="cuda") with self.assertRaises(RuntimeError): r.add_video_stream(-1, decoder="h264_cuvid") @_media_source class CudaDecoderTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase): def _test_decode( self, decoder: str, src_path: str, height: int, width: int, ref_num_frames: int, hw_accel=None, decoder_option=None, dtype: torch.dtype = torch.uint8, ): src = self.get_src(get_asset_path(src_path)) r = StreamReader(src) r.add_video_stream(10, decoder=decoder, decoder_option=decoder_option, hw_accel=hw_accel) num_frames = 0 for (chunk,) in r.stream(): self.assertEqual(chunk.device, torch.device(hw_accel or "cpu")) self.assertEqual(chunk.dtype, dtype) self.assertEqual(chunk.shape, torch.Size([10, 3, height, width])) num_frames += chunk.size(0) assert num_frames == ref_num_frames @skipIfNoHWAccel("h264_cuvid") def test_h264_cuvid(self): """GPU decoder works for H264""" self._test_decode("h264_cuvid", "nasa_13013.mp4", 270, 480, 390) @skipIfNoHWAccel("h264_cuvid") def test_h264_cuvid_hw_accel(self): """GPU decoder works for H264 with HW acceleration, and put the frames on CUDA tensor""" self._test_decode("h264_cuvid", "nasa_13013.mp4", 270, 480, 390, hw_accel="cuda:0") @skipIfNoHWAccel("h264_cuvid") def test_h264_cuvid_hw_accel_resize(self): """GPU decoder works for H264 with HW acceleration and resize option""" w, h = 240, 136 self._test_decode( "h264_cuvid", "nasa_13013.mp4", h, w, 390, hw_accel="cuda:0", decoder_option={"resize": f"{w}x{h}"} ) @skipIfNoHWAccel("h264_cuvid") def test_h264_cuvid_hw_accel_crop(self): """GPU decoder works for H264 with HW acceleration and crop option""" top, bottom, left, right = 3, 5, 7, 9 self._test_decode( "h264_cuvid", "nasa_13013.mp4", 262, 464, 390, hw_accel="cuda:0", decoder_option={"crop": f"{top}x{bottom}x{left}x{right}"}, ) @skipIfNoHWAccel("hevc_cuvid") def test_hevc_cuvid(self): """GPU decoder works for H265/HEVC""" self._test_decode("hevc_cuvid", "testsrc.hevc", 144, 256, 300) @skipIfNoHWAccel("hevc_cuvid") def test_hevc_cuvid_hw_accel(self): """GPU decoder works for H265/HEVC with HW acceleration, and put the frames on CUDA tensor""" self._test_decode("hevc_cuvid", "testsrc.hevc", 144, 256, 300, hw_accel="cuda:0", dtype=torch.int16) @skipIfNoHWAccel("hevc_cuvid") def test_hevc_cuvid_hw_accel_resize(self): """GPU decoder works for H265/HEVC with HW acceleration and resize option""" w, h = 128, 64 self._test_decode( "hevc_cuvid", "testsrc.hevc", h, w, 300, hw_accel="cuda:0", dtype=torch.int16, decoder_option={"resize": f"{w}x{h}"}, ) @skipIfNoHWAccel("hevc_cuvid") def test_hevc_cuvid_hw_accel_crop(self): """GPU decoder works for H265/HEVC with HW acceleration and crop option""" top, bottom, left, right = 3, 5, 7, 9 self._test_decode( "hevc_cuvid", "testsrc.hevc", 136, 240, 300, hw_accel="cuda:0", dtype=torch.int16, decoder_option={"crop": f"{top}x{bottom}x{left}x{right}"}, ) @skipIfNoHWAccel("h264_cuvid") # Disabled in CI: https://github.com/pytorch/audio/issues/3376 @disabledInCI class FilterGraphWithCudaAccel(TorchaudioTestCase): def test_sclae_cuda_change_size(self): """scale_cuda filter can be used when HW accel is on""" src = get_asset_path("nasa_13013.mp4") r = StreamReader(src) r.add_video_stream(10, decoder="h264_cuvid", hw_accel="cuda", filter_desc="scale_cuda=iw/2:ih/2") num_frames = 0 for (chunk,) in r.stream(): self.assertEqual(chunk.device, torch.device("cuda:0")) self.assertEqual(chunk.dtype, torch.uint8) self.assertEqual(chunk.shape, torch.Size([10, 3, 135, 240])) num_frames += chunk.size(0) assert num_frames == 390 def test_scale_cuda_format(self): """yuv444p format conversion should work""" src = get_asset_path("nasa_13013.mp4") r = StreamReader(src) r.add_video_stream(10, decoder="h264_cuvid", hw_accel="cuda", filter_desc="scale_cuda=format=yuv444p") num_frames = 0 for (chunk,) in r.stream(): self.assertEqual(chunk.device, torch.device("cuda:0")) self.assertEqual(chunk.dtype, torch.uint8) self.assertEqual(chunk.shape, torch.Size([10, 3, 270, 480])) num_frames += chunk.size(0) assert num_frames == 390