.

Author: ahmadsharif1

examples/basic_cuda_example.py

@@ -4,14 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 """
-Accelerated video decoding with NVDEC
+CUDA Decoding on Nvidia GPUs
 =====================================
 
-.. _nvdec_tutorial:
+.. _ndecoderec_tutorial:
 
-**Author**: `Ahmad Sharif <[email protected]>`__
-
-TorchCodec can use Nvidia hardware to speed-up video decoding. This is called "CUDA Decoding".
+TorchCodec can use supported Nvidia hardware (see support matrix here
+<https://developer.nvidia.com/video-encode-and-decode-gpu-support-matrix-new>) to speed-up
+video decoding. This is called "CUDA Decoding".
 CUDA Decoding can be faster than CPU Decoding for the actual decoding step and also for
 subsequent transform steps like scaling, cropping or rotating. This is because the decode step leaves
 the decoded tensor in GPU memory so the GPU doesn't have to fetch from main memory before
@@ -30,34 +30,34 @@
 Here are situations where CUDA Decoding may not make sense:
 
 #. You want bit-exact results compared to CPU Decoding
-#. If you have small resolution videos and the PCI-e transfer latency is large
-#. If your GPU is already busy and CPU is not
+#. You have small resolution videos and the PCI-e transfer latency is large
+#. Your GPU is already busy and CPU is not
 
 It's best to experiment with CUDA Decoding to see if it improves your use-case. With
 TorchCodec you can simply pass in a device parameter to the
 :class:`~torchcodec.decoders.VideoDecoder` class to use CUDA Decoding.
 
 
-In order use CUDA Decoding will need the following installed in your environment:
+In order to use CUDA Decoding will need the following installed in your environment:
 
-#. CUDA-enabled pytorch
-#. FFmpeg binaries that support NVDEC-enabled codecs
+#. An Nvidia GPU that supports decoding the video format you want to decode. See
+   the support matrix here <https://developer.nvidia.com/video-encode-and-decode-gpu-support-matrix-new>
+#. `CUDA-enabled pytorch <https://pytorch.org/get-started/locally/>`_
+#. FFmpeg binaries that support NdecoderEC-enabled codecs
 #. libnpp and nvrtc (these are usually installed when you install the full cuda-toolkit)
 
 
-FFmpeg versions 5, 6 and 7 from conda-forge are built with NVDEC support and you can
+FFmpeg versions 5, 6 and 7 from conda-forge are built with NdecoderEC support and you can
 install them with conda. For example, to install FFmpeg version 7:
 
 .. code-block:: bash
-   # These libraries are needed for CUDA decoding with TorchCodec
    conda install ffmpeg=7 -c conda-forge
    conda install libnpp cuda-nvrtc -c nvidia
 """
 
 # %%
-######################################################################
 # Checking if Pytorch has CUDA enabled
-######################################################################
+# -------------------------------------
 #
 # .. note::
 #
@@ -72,9 +72,8 @@
 
 
 # %%
-######################################################################
 # Downloading the video
-######################################################################
+# -------------------------------------
 #
 # We will use the following video which has the following properties;
 #
@@ -98,16 +97,15 @@
 
 
 # %%
-######################################################################
 # CUDA Decoding using VideoDecoder
-######################################################################
+# -------------------------------------
 #
 # To use CUDA decoder, you need to pass in a cuda device to the decoder.
 #
 from torchcodec.decoders import VideoDecoder
 
-vd = VideoDecoder(video_file, device="cuda:0")
-frame = vd[0]
+decoder = VideoDecoder(video_file, device="cuda")
+frame = decoder[0]
 
 # %%
 #
@@ -123,13 +121,11 @@
 
 
 # %%
-######################################################################
 # Visualizing Frames
-######################################################################
+# -------------------------------------
 #
 # Let's look at the frames decoded by CUDA decoder and compare them
 # against equivalent results from the CPU decoders.
-import matplotlib.pyplot as plt
 
 
 def get_frames(timestamps: list[float], device: str):
@@ -148,25 +144,43 @@ def get_numpy_images(frames):
 
 timestamps = [12, 19, 45, 131, 180]
 cpu_frames = get_frames(timestamps, device="cpu")
-cuda_frames = get_frames(timestamps, device="cuda:0")
+cuda_frames = get_frames(timestamps, device="cuda")
 cpu_numpy_images = get_numpy_images(cpu_frames)
 cuda_numpy_images = get_numpy_images(cuda_frames)
 
 
-def plot_cpu_and_cuda_images():
-    n_rows = len(timestamps)
-    fig, axes = plt.subplots(n_rows, 2, figsize=[12.8, 16.0])
-    for i in range(n_rows):
-        axes[i][0].imshow(cpu_numpy_images[i])
-        axes[i][1].imshow(cuda_numpy_images[i])
+def plot(
+    frames1: List[torch.Tensor],
+    frames2: List[torch.Tensor],
+    title1: Optional[str] = None,
+    title2: Optional[str] = None,
+):
+    try:
+        import matplotlib.pyplot as plt
+        from torchvision.transforms.v2.functional import to_pil_image
+        from torchvision.utils import make_grid
+    except ImportError:
+        print("Cannot plot, please run `pip install torchvision matplotlib`")
+        return
+
+    plt.rcParams["savefig.bbox"] = "tight"
+
+    fig, ax = plt.subplots(1, 2)
+
+    ax[0].imshow(to_pil_image(make_grid(frames1)))
+    ax[0].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
+    if title1 is not None:
+        ax[0].set_title(title1)
+
+    ax[1].imshow(to_pil_image(make_grid(frames2)))
+    ax[1].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
+    if title2 is not None:
+        ax[1].set_title(title2)
 
-    axes[0][0].set_title("CPU decoder")
-    axes[0][1].set_title("CUDA decoder")
-    plt.setp(axes, xticks=[], yticks=[])
     plt.tight_layout()
 
 
-plot_cpu_and_cuda_images()
+plot(cpu_frames, cuda_frames, "CPU decoder", "CUDA decoder")
 
 # %%
 #