.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated_examples/low_level/raw_data.py" .. LINE NUMBERS ARE GIVEN BELOW. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_low_level_raw_data.py: ================================ Raw frames and raw audio samples ================================ .. currentmodule:: torchcodec.decoders .. important:: **The low-level APIs are in beta.** Their signatures and semantics may still change slightly, in response to user feedback. Please `share your feedback `__! In this tutorial, we'll skip the conversion stage of a decoding pipeline and read the decoder's own YUV planes and audio samples directly, at the source's own precision. The last stage of a decoding pipeline - a :class:`ColorConverter` for video, an :class:`AudioConverter` for audio - is optional. If you stop before it, you get what the decoder actually produced: YUV planes in the codec's own pixel format, and audio samples in the codec's own sample type, with no conversion, no normalisation, and no copy. This is worth doing when the conversion you want isn't the one the converters perform. You may want to apply a custom colorspace, train a model directly in YUV space, write a kernel fused with the first layer of your model, decode a 10-bit HDR source without flattening it to `uint8` or `float32`, or scale integer audio samples yourself. This tutorial assumes you are familiar with the three stages described in :ref:`sphx_glr_generated_examples_low_level_basics.py`. .. GENERATED FROM PYTHON SOURCE LINES 41-42 Boilerplate: a test video, and the device we'll run on. .. GENERATED FROM PYTHON SOURCE LINES 42-72 .. code-block:: Python import subprocess import tempfile from pathlib import Path import torch device = "cuda" if torch.cuda.is_available() else "cpu" print(f"{device = }") temp_dir = Path(tempfile.mkdtemp()) video_path = temp_dir / "video.mp4" subprocess.run( [ "ffmpeg", "-y", "-hide_banner", "-loglevel", "error", "-f", "lavfi", "-i", "testsrc2=size=1280x720:rate=30:duration=5", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-g", "30", "-colorspace", "bt709", "-color_primaries", "bt709", "-color_trc", "bt709", str(video_path), ], check=True, ) def decode_raw(demuxer, packet_decoder): for packet in demuxer: yield from packet_decoder.decode(packet) yield from packet_decoder.drain() .. rst-class:: sphx-glr-script-out .. code-block:: none device = 'cuda' .. GENERATED FROM PYTHON SOURCE LINES 74-81 Raw frames ---------- A :class:`VideoPacketDecoder` produces :class:`RawFrame` objects, and a :class:`RawFrame` can hand out the decoder's own planes as tensor views, with no copy and no conversion. It also carries everything you need to interpret them. .. GENERATED FROM PYTHON SOURCE LINES 81-92 .. code-block:: Python from torchcodec.decoders import ColorConverter, Demuxer demuxer = Demuxer(video_path) packet_decoder = demuxer.streams[0].make_decoder(device=device) raw_frame = next(decode_raw(demuxer, packet_decoder)) Y, U, V = raw_frame.planes print(f"{raw_frame.pixel_format = }, {raw_frame.bit_depth = }") print(f"{raw_frame.color_space = }, {raw_frame.color_range = }") print(f"{Y.shape = }, {U.shape = }, {Y.dtype = }, {Y.stride() = }") .. rst-class:: sphx-glr-script-out .. code-block:: none raw_frame.pixel_format = 'nv12', raw_frame.bit_depth = 8 raw_frame.color_space = 'bt709', raw_frame.color_range = 'tv' Y.shape = torch.Size([720, 1280]), U.shape = torch.Size([360, 640]), Y.dtype = torch.uint8, Y.stride() = (1536, 1) .. GENERATED FROM PYTHON SOURCE LINES 93-108 There is one tensor per component of :attr:`RawFrame.pixel_format`, in the order that format describes, and only the luma and alpha ones are full size: the chroma planes are subsampled by whatever the format says - half in both directions for the 4:2:0 formats above. These are views into the frame's memory, so they are not contiguous in general: the row stride is the decoder's own line size, and the chroma planes of a semi-planar format such as an NVDEC ``nv12`` surface are two interleaved views over a single allocation. Writing through them is visible to whatever reads the :class:`RawFrame` next. Being the decoder's own planes, they are also never rotated - a video whose container asks for a rotation gives you the samples as they were encoded, and :attr:`RawFrame.rotation` tells you what to apply. A :class:`ColorConverter` applies it for you. .. GENERATED FROM PYTHON SOURCE LINES 110-120 .. _raw_data_custom_conversion: Doing the conversion yourself ----------------------------- With :attr:`RawFrame.planes` and the metadata beside it (:attr:`~RawFrame.color_space`, :attr:`~RawFrame.color_range`, :attr:`~RawFrame.bit_depth`), the conversion is yours to write. Here it's plain PyTorch ops - it could just as well be a Triton or CUDA kernel, fused with whatever your model needs next. .. GENERATED FROM PYTHON SOURCE LINES 120-150 .. code-block:: Python assert raw_frame.pixel_format in ("yuv420p", "nv12") # 8-bit 4:2:0, CPU and CUDA assert raw_frame.color_space == "bt709" and raw_frame.color_range == "tv" def yuv420_to_rgb(Y, U, V): # BT.709, limited range. Chroma is upsampled by nearest neighbour. height, width = Y.shape def upsample(plane): plane = (plane.float() - 128) * (255 / 224) return plane.repeat_interleave(2, 0).repeat_interleave(2, 1)[:height, :width] y = (Y.float() - 16) * (255 / 219) u, v = upsample(U), upsample(V) rgb = torch.stack( [ y + 1.5748 * v, y - 0.1873 * u - 0.4681 * v, y + 1.8556 * u, ] ) return rgb.round_().clamp_(0, 255).to(torch.uint8) ours = yuv420_to_rgb(Y, U, V) reference = ColorConverter(device=device).convert(raw_frame).data print(f"{ours.shape = }, mean abs diff vs ColorConverter: " f"{(ours.float() - reference.float()).abs().mean():.2f}") .. rst-class:: sphx-glr-script-out .. code-block:: none ours.shape = torch.Size([3, 720, 1280]), mean abs diff vs ColorConverter: 0.29 .. GENERATED FROM PYTHON SOURCE LINES 151-163 Reading or converting the planes on CUDA ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ The planes are views into a PyTorch CUDA allocation that the decoder still owns and may reuse for a later frame. As long as you read them on the stream the decoder ran on, which is the case unless you explicitly request otherwise, there is no problem and you don't need to think about stream synchronization. However, if you read them on a *different* stream than the decoder stream, you have to wait for the decoder's asynchronous copy before reading, and keep the allocator from recycling the buffer while your reads are still queued. See :ref:`sphx_glr_generated_examples_low_level_cuda_streams.py` to learn more. .. GENERATED FROM PYTHON SOURCE LINES 165-175 Formats that can't be viewed ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Not every pixel format can be exposed as a tensor without a copy. :attr:`RawFrame.planes` raises a ``RuntimeError`` for the sub-byte-packed, palettised and float formats, and for frames stored bottom-up. Check :attr:`~RawFrame.pixel_format` first if you are decoding something exotic; on CUDA you never have to, since NVDEC only ever produces a handful of surface formats (``nv12``, ``p010le``, ``p012le``, ``p016le``, ``yuv444p``, ``yuv444p16le``). .. GENERATED FROM PYTHON SOURCE LINES 177-185 Raw HDR frames -------------- Raw planes come at the source's own precision, so a 10-bit HDR video gives ``uint16`` planes with all 10 bits intact - no clipping to 8 bits, and no tone mapping. (For the :class:`~torchcodec.decoders.VideoDecoder` route to HDR, through its ``output_dtype`` parameter, see :ref:`sphx_glr_generated_examples_decoding_hdr_decoding.py`.) .. GENERATED FROM PYTHON SOURCE LINES 185-208 .. code-block:: Python hdr_video_path = temp_dir / "hdr.mp4" subprocess.run( [ "ffmpeg", "-y", "-hide_banner", "-loglevel", "error", "-f", "lavfi", "-i", "testsrc2=size=1280x720:rate=30:duration=1", "-c:v", "libx265", "-pix_fmt", "yuv420p10le", "-preset", "ultrafast", "-x265-params", "colorprim=bt2020:transfer=smpte2084:colormatrix=bt2020nc:range=limited", str(hdr_video_path), ], check=True, capture_output=True, # x265 logs its banner to stderr no matter what ) hdr_demuxer = Demuxer(hdr_video_path) hdr_packet_decoder = hdr_demuxer.streams[0].make_decoder(device=device) hdr_raw = next(decode_raw(hdr_demuxer, hdr_packet_decoder)) hdr_Y = hdr_raw.planes[0] print(f"{hdr_raw.pixel_format = }, {hdr_raw.bit_depth = }, {hdr_Y.dtype = }") print(f"{hdr_raw.color_space = }, {hdr_raw.color_primaries = }, " f"{hdr_raw.color_transfer_characteristic = }") .. rst-class:: sphx-glr-script-out .. code-block:: none hdr_raw.pixel_format = 'p010le', hdr_raw.bit_depth = 10, hdr_Y.dtype = torch.uint16 hdr_raw.color_space = 'bt2020nc', hdr_raw.color_primaries = 'bt2020', hdr_raw.color_transfer_characteristic = 'smpte2084' .. GENERATED FROM PYTHON SOURCE LINES 209-216 A plane's dtype only tells you its storage width, ``uint8`` or ``uint16``; :attr:`RawFrame.bit_depth` is what tells you the range of the values held in it, and so what to shift or scale by. NVDEC surfaces are 16-bit containers holding the samples msb-aligned, so the 10 bits sit at the top and the low 6 are zero. Shift them back down to read the sample values. .. GENERATED FROM PYTHON SOURCE LINES 216-221 .. code-block:: Python shift = 16 - hdr_raw.bit_depth if device == "cuda" else 0 samples = hdr_Y.to(torch.int32) >> shift print(f"luma range: [{samples.min()}, {samples.max()}], " f"{2 ** hdr_raw.bit_depth} levels available") .. rst-class:: sphx-glr-script-out .. code-block:: none luma range: [15, 888], 1024 levels available .. GENERATED FROM PYTHON SOURCE LINES 222-237 Raw audio samples ----------------- The audio side is similar. An :class:`AudioPacketDecoder` hands out :class:`RawAudioSamples` objects, whose :attr:`~RawAudioSamples.data` is always a contiguous ``[num_channels, num_samples]`` tensor - planar and packed sources alike - in whichever dtype holds the source's samples exactly. One thing differs from video: these samples are a *copy*, not a view. :attr:`RawFrame.planes` aliases the buffer of the frame as the decoder produced it, whereas :attr:`~RawAudioSamples.data` is a copy so that all samples come out as ``[num_channels, num_samples]``. It is still the tensor an :class:`AudioConverter` reads, though, so mutating it in place does change what the converter gives you. .. GENERATED FROM PYTHON SOURCE LINES 237-255 .. code-block:: Python audio_path = temp_dir / "audio.wav" subprocess.run( [ "ffmpeg", "-y", "-hide_banner", "-loglevel", "error", "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=44100:duration=1", "-c:a", "pcm_s16le", str(audio_path), ], check=True, ) demuxer = Demuxer(audio_path, streams="audio") packet_decoder = demuxer.streams[0].make_decoder() raw_samples = next(decode_raw(demuxer, packet_decoder)) print(f"{raw_samples.data.dtype = }, {raw_samples.data.shape = }, " f"{raw_samples.sample_rate = }") print(f"value range: [{raw_samples.data.min()}, {raw_samples.data.max()}]") .. rst-class:: sphx-glr-script-out .. code-block:: none raw_samples.data.dtype = torch.int16, raw_samples.data.shape = torch.Size([1, 4096]), raw_samples.sample_rate = 44100 value range: [-4095, 4095] .. GENERATED FROM PYTHON SOURCE LINES 256-261 An ``s16`` source gives ``int16``, an ``s32`` source ``int32``, an ``fltp`` source ``float32``, and so on. The integer ones are *not* normalised to ``[-1, 1]``: that, along with resampling and remixing, is what an :class:`AudioConverter` does. If normalising is all you need, you can always do it manually: .. GENERATED FROM PYTHON SOURCE LINES 261-265 .. code-block:: Python normalised = raw_samples.data.to(torch.float32) / 2 ** 15 # FFmpeg's `sine` source is a quiet tone, so this doesn't reach ±1. print(f"{normalised.dtype = }, " f"range [{normalised.min():.3f}, {normalised.max():.3f}]") .. rst-class:: sphx-glr-script-out .. code-block:: none normalised.dtype = torch.float32, range [-0.125, 0.125] .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 1.240 seconds) .. _sphx_glr_download_generated_examples_low_level_raw_data.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: raw_data.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: raw_data.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: raw_data.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_