{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3BxT9HiUjMyv"
      },
      "source": [
        "# Hardware-Accelerated Video Decoding and Encoding\n",
        "\n",
        "**Author**: [Moto Hira](moto@meta.com)\n",
        "\n",
        "This tutorial shows how to use NVIDIA's hardware video decoder (NVDEC) and encoder (NVENC) with TorchAudio.\n",
        "\n",
        "Using hardware encoder/decoder improves the speed of loading and saving certain types of videos. Using them in TorchAduio requires additional FFmpeg configuration. This tutorial goes over how to compile FFmpeg, and compare the speed it takes to process video.\n",
        "\n",
        "**WARNING**\n",
        "\n",
        "> This tutorial instsalls FFmpeg in system directory.\n",
        "> If you run this tutorial on your system, please adjust the build configuration accordingly.\n",
        "\n",
        "**NOTE**\n",
        "\n",
        "> This tutorial was authored in Google Colab, and is tailored to Google Colab's specifications. Please check out this tutorial in [Google Colab](https://colab.research.google.com/drive/1DDah_IaGULEO66CfQWltRqaVheBkiXdN#sandboxMode=true).\n",
        "\n",
        "To use NVENC/NVDEC with TorchAudio, the following items are required.\n",
        "\n",
        "1. NVIDIA GPU with hardware video decoder/encoder.\n",
        "2. FFmpeg libraries compiled with NVDEC/NVENC support. †\n",
        "3. PyTorch / TorchAudio with CUDA support.\n",
        "\n",
        "TorchAudio's official binary distributions are compiled with FFmpeg 4 libraries, and they contain the logic required for hardware-based decoding/encoding.\n",
        "\n",
        "In the following sections, we build FFmpeg 4 libraries with NVDEC/NVENC support, then we demonstrate the performance imrovement using TorchAudio's `StreamReader`/`StreamWriter`.\n",
        "\n",
        "† For details on NVDEC/NVENC and FFmpeg, please refer to the following articles.\n",
        "\n",
        "* https://docs.nvidia.com/video-technologies/video-codec-sdk/nvdec-video-decoder-api-prog-guide/\n",
        "* https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/#compiling-ffmpeg\n",
        "* https://developer.nvidia.com/blog/nvidia-ffmpeg-transcoding-guide/\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08XIDshiiFp-"
      },
      "source": [
        "## Check the available GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gCegZlbgutNM",
        "outputId": "0eca8ead-4671-4b83-ca9b-a602642844a5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Fri Oct  7 13:01:26 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   56C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Update PyTorch and TorchAudio with nightly builds\n",
        "\n",
        "Until TorchAudio 0.13 is released, we need to use the nightly builds of PyTorch and TorchAudio."
      ],
      "metadata": {
        "id": "8EI-DwaeQbjp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fACIAl1dgYt-",
        "outputId": "7ff2cb88-33e5-4b8c-9542-4f53cce09cfa"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/, https://download.pytorch.org/whl/nightly/cu116\n",
            "Collecting torch\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu116/torch-1.14.0.dev20221007%2Bcu116-cp37-cp37m-linux_x86_64.whl (2286.1 MB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hCollecting torchaudio\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu116/torchaudio-0.13.0.dev20221006%2Bcu116-cp37-cp37m-linux_x86_64.whl (4.2 MB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (4.1.1)\n",
            "Collecting torch\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu116/torch-1.13.0.dev20221006%2Bcu116-cp37-cp37m-linux_x86_64.whl (1983.0 MB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hInstalling collected packages: torch, torchaudio\n",
            "Successfully installed torch-1.13.0.dev20221006+cu116 torchaudio-0.13.0.dev20221006+cu116\n"
          ]
        }
      ],
      "source": [
        "!pip uninstall -y -q torch torchaudio torchvision torchtext\n",
        "!pip install --progress-bar off --pre torch torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu116 2> /dev/null"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RAziwx_MqJzr"
      },
      "source": [
        "## Build FFmpeg libraries with Nvidia NVDEC/NVENC support\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcKsO-zurtWM"
      },
      "source": [
        "### Install NVIDIA Video Codec Headers\n",
        "\n",
        "To build FFmpeg with NVDEC/NVENC, we first install the headers that FFmpeg uses to interact with Video Codec SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vgbusLNCq2mV",
        "outputId": "04b9c2b6-bc18-4dd0-ea8f-4b320559ebb3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'nv-codec-headers'...\n",
            "remote: Enumerating objects: 819, done.\u001b[K\n",
            "remote: Counting objects: 100% (819/819), done.\u001b[K\n",
            "remote: Compressing objects: 100% (697/697), done.\u001b[K\n",
            "remote: Total 819 (delta 439), reused 0 (delta 0)\u001b[K\n",
            "Receiving objects: 100% (819/819), 156.42 KiB | 410.00 KiB/s, done.\n",
            "Resolving deltas: 100% (439/439), done.\n",
            "Note: checking out 'n11.0.10.1'.\n",
            "\n",
            "You are in 'detached HEAD' state. You can look around, make experimental\n",
            "changes and commit them, and you can discard any commits you make in this\n",
            "state without impacting any branches by performing another checkout.\n",
            "\n",
            "If you want to create a new branch to retain commits you create, you may\n",
            "do so (now or later) by using -b with the checkout command again. Example:\n",
            "\n",
            "  git checkout -b <new-branch-name>\n",
            "\n",
            "HEAD is now at 315ad74 add cuMemcpy\n",
            "sed 's#@@PREFIX@@#/usr/local#' ffnvcodec.pc.in > ffnvcodec.pc\n",
            "install -m 0755 -d '/usr/local/include/ffnvcodec'\n",
            "install -m 0644 include/ffnvcodec/*.h '/usr/local/include/ffnvcodec'\n",
            "install -m 0755 -d '/usr/local/lib/pkgconfig'\n",
            "install -m 0644 ffnvcodec.pc '/usr/local/lib/pkgconfig'\n"
          ]
        }
      ],
      "source": [
        "!git clone https://git.videolan.org/git/ffmpeg/nv-codec-headers.git\n",
        "# Note: Google Colab's GPU has NVENC API ver 11.0, so we checkout 11.0 tag.\n",
        "!cd nv-codec-headers && git checkout n11.0.10.1 && sudo make install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PWT0m_LGq_GE"
      },
      "source": [
        "### Download FFmpeg source code\n",
        "\n",
        "Next we download the source code of FFmpeg 4. Any version later than 4.1 should work. We use 4.4.2 here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "AKERwitAzfLo"
      },
      "outputs": [],
      "source": [
        "!wget -q https://github.com/FFmpeg/FFmpeg/archive/refs/tags/n4.4.2.tar.gz\n",
        "!tar -xf n4.4.2.tar.gz\n",
        "!mv FFmpeg-n4.4.2 ffmpeg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpWDpxBvDy-b"
      },
      "source": [
        "### Install FFmpeg build and runtime dependencies\n",
        "\n",
        "In the later test, we use H264 video codec and HTTPS protocol, so we install the libraries for them here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vn50oz2UvvPP",
        "outputId": "b2e695d9-c953-4370-fb8d-92cdd8a832be"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STRIP\tinstall-libavutil-shared\n",
            "\n",
            "... Omitted for brevity ...\n",
            "\n",
            "Setting up libx264-dev:amd64 (2:0.152.2854+gite9a5903-2) ...\n",
            "Setting up yasm (1.3.0-2build1) ...\n",
            "Setting up libunbound2:amd64 (1.6.7-1ubuntu2.5) ...\n",
            "Setting up libp11-kit-dev:amd64 (0.23.9-2ubuntu0.1) ...\n",
            "Setting up libtasn1-6-dev:amd64 (4.13-2) ...\n",
            "Setting up libtasn1-doc (4.13-2) ...\n",
            "Setting up libgnutlsxx28:amd64 (3.5.18-1ubuntu1.6) ...\n",
            "Setting up libgnutls-dane0:amd64 (3.5.18-1ubuntu1.6) ...\n",
            "Setting up libgnutls-openssl27:amd64 (3.5.18-1ubuntu1.6) ...\n",
            "Setting up libgmpxx4ldbl:amd64 (2:6.1.2+dfsg-2) ...\n",
            "Setting up libidn2-dev:amd64 (2.0.4-1.1ubuntu0.2) ...\n",
            "Setting up libidn2-0-dev (2.0.4-1.1ubuntu0.2) ...\n",
            "Setting up libgmp-dev:amd64 (2:6.1.2+dfsg-2) ...\n",
            "Setting up nettle-dev:amd64 (3.4.1-0ubuntu0.18.04.1) ...\n",
            "Setting up libgnutls28-dev:amd64 (3.5.18-1ubuntu1.6) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "Processing triggers for libc-bin (2.27-3ubuntu1.6) ...\n"
          ]
        }
      ],
      "source": [
        "!apt -qq update\n",
        "!apt -qq install -y yasm libx264-dev libgnutls28-dev"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wL3lAZWPTJaN"
      },
      "source": [
        "### Configure FFmpeg build with Nvidia CUDA hardware support\n",
        "\n",
        "Next we configure FFmpeg build. Note the following:\n",
        "\n",
        "1. We provide flags like `-I/usr/local/cuda/include`, `-L/usr/local/cuda/lib64` to let the build process know where the CUDA libraries are found.\n",
        "2. We provide flags like `--enable-nvdec` and `--enable-nvenc` to enable NVDEC/NVENC. Please check out the Transcoding Guide† for the detail.\n",
        "3. We also provide NVCC flags with compute capability 37.   \n",
        "This is because by default the configuration script verifies NVCC by compiling sample code targeting compute capability 30, which is too old for CUDA 11.\n",
        "4. Many features are disabled to reduce the compilation time.\n",
        "5. We install the library in `/usr/lib/`, which is one of the active search path of the dynamic loader.   \n",
        "Doing so allows the resulting libraries to be found without requiring a restart of the current session. This might be an undesirable location, e.g. when one isn't using a disposable VM.\n",
        "\n",
        "† NVIDIA FFmpeg Transcoding Guide https://developer.nvidia.com/blog/nvidia-ffmpeg-transcoding-guide/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KZr8bMdztRz1",
        "outputId": "d542338b-bfcd-4176-c349-4a9235fcf411"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: ccap=37\n",
            "install prefix            /usr/\n",
            "source path               .\n",
            "C compiler                gcc\n",
            "C library                 glibc\n",
            "ARCH                      x86 (generic)\n",
            "big-endian                no\n",
            "runtime cpu detection     yes\n",
            "standalone assembly       yes\n",
            "x86 assembler             yasm\n",
            "MMX enabled               yes\n",
            "MMXEXT enabled            yes\n",
            "3DNow! enabled            yes\n",
            "3DNow! extended enabled   yes\n",
            "SSE enabled               yes\n",
            "SSSE3 enabled             yes\n",
            "AESNI enabled             yes\n",
            "AVX enabled               yes\n",
            "AVX2 enabled              yes\n",
            "AVX-512 enabled           yes\n",
            "XOP enabled               yes\n",
            "FMA3 enabled              yes\n",
            "FMA4 enabled              yes\n",
            "i686 features enabled     yes\n",
            "CMOV is fast              yes\n",
            "EBX available             yes\n",
            "EBP available             yes\n",
            "debug symbols             yes\n",
            "strip symbols             yes\n",
            "optimize for size         no\n",
            "optimizations             yes\n",
            "static                    no\n",
            "shared                    yes\n",
            "postprocessing support    no\n",
            "network support           yes\n",
            "threading support         pthreads\n",
            "safe bitstream reader     yes\n",
            "texi2html enabled         no\n",
            "perl enabled              yes\n",
            "pod2man enabled           yes\n",
            "makeinfo enabled          no\n",
            "makeinfo supports HTML    no\n",
            "\n",
            "External libraries:\n",
            "alsa                    libx264                 lzma\n",
            "bzlib                   libxcb                  zlib\n",
            "gnutls                  libxcb_shape\n",
            "iconv                   libxcb_xfixes\n",
            "\n",
            "External libraries providing hardware acceleration:\n",
            "cuda                    cuvid                   nvenc\n",
            "cuda_llvm               ffnvcodec               v4l2_m2m\n",
            "cuda_nvcc               nvdec\n",
            "\n",
            "Libraries:\n",
            "avcodec                 avformat                swscale\n",
            "avdevice                avutil\n",
            "avfilter                swresample\n",
            "\n",
            "Programs:\n",
            "ffmpeg                  ffprobe\n",
            "\n",
            "Enabled decoders:\n",
            "aac                     hevc                    rawvideo\n",
            "av1                     mjpeg                   vc1\n",
            "h263                    mpeg1video              vp8\n",
            "h264                    mpeg2video              vp9\n",
            "h264_cuvid              mpeg4\n",
            "\n",
            "Enabled encoders:\n",
            "h264_nvenc              libx264\n",
            "\n",
            "Enabled hwaccels:\n",
            "av1_nvdec               mpeg1_nvdec             vp8_nvdec\n",
            "h264_nvdec              mpeg2_nvdec             vp9_nvdec\n",
            "hevc_nvdec              mpeg4_nvdec             wmv3_nvdec\n",
            "mjpeg_nvdec             vc1_nvdec\n",
            "\n",
            "Enabled parsers:\n",
            "h263                    mpeg4video              vp9\n",
            "\n",
            "Enabled demuxers:\n",
            "mov\n",
            "\n",
            "Enabled muxers:\n",
            "mov                     mp4\n",
            "\n",
            "Enabled protocols:\n",
            "file                    tcp\n",
            "https                   tls\n",
            "\n",
            "Enabled filters:\n",
            "aformat                 hflip                   transpose\n",
            "anull                   null                    trim\n",
            "atrim                   scale                   vflip\n",
            "format                  testsrc2\n",
            "\n",
            "Enabled bsfs:\n",
            "aac_adtstoasc           null                    vp9_superframe_split\n",
            "h264_mp4toannexb        vp9_superframe\n",
            "\n",
            "Enabled indevs:\n",
            "lavfi\n",
            "\n",
            "Enabled outdevs:\n",
            "\n",
            "License: nonfree and unredistributable\n"
          ]
        }
      ],
      "source": [
        "# NOTE:\n",
        "# When the configure script of FFmpeg 4 checks nvcc, it uses compute\n",
        "# capability of 30 (3.0) by default. CUDA 11, however, does not support \n",
        "# compute capability 30.\n",
        "# Here, we use 37, which is supported by CUDA 11 and both K80 and T4.\n",
        "#\n",
        "# Tesla K80: 37\n",
        "# NVIDIA T4: 75\n",
        "\n",
        "%env ccap=37\n",
        "\n",
        "# NOTE:\n",
        "# We disable most of the features to speed up compilation\n",
        "# The necessary components are\n",
        "# - demuxer: mov, aac\n",
        "# - decoder: h264, h264_nvdec\n",
        "# - muxer: mp4\n",
        "# - encoder: libx264, h264_nvenc\n",
        "# - gnutls (HTTPS)\n",
        "#\n",
        "# Additionally, we use FFmpeg's virtual input device to generate \n",
        "# test video data. This requires\n",
        "# - input device: lavfi\n",
        "# - filter: testsrc2\n",
        "# - decoder: rawvideo\n",
        "#\n",
        "\n",
        "!cd ffmpeg && ./configure \\\n",
        "  --prefix='/usr/' \\\n",
        "  --extra-cflags='-I/usr/local/cuda/include' \\\n",
        "  --extra-ldflags='-L/usr/local/cuda/lib64' \\\n",
        "  --nvccflags=\"-gencode arch=compute_${ccap},code=sm_${ccap} -O2\" \\\n",
        "  --disable-doc \\\n",
        "  --disable-static \\\n",
        "  --disable-bsfs \\\n",
        "  --disable-decoders \\\n",
        "  --disable-encoders \\\n",
        "  --disable-filters \\\n",
        "  --disable-demuxers \\\n",
        "  --disable-devices \\\n",
        "  --disable-muxers \\\n",
        "  --disable-parsers \\\n",
        "  --disable-postproc \\\n",
        "  --disable-protocols \\\n",
        "  --enable-decoder=aac \\\n",
        "  --enable-decoder=h264 \\\n",
        "  --enable-decoder=h264_cuvid \\\n",
        "  --enable-decoder=rawvideo \\\n",
        "  --enable-indev=lavfi \\\n",
        "  --enable-encoder=libx264 \\\n",
        "  --enable-encoder=h264_nvenc \\\n",
        "  --enable-demuxer=mov \\\n",
        "  --enable-muxer=mp4 \\\n",
        "  --enable-filter=scale \\\n",
        "  --enable-filter=testsrc2 \\\n",
        "  --enable-protocol=file \\\n",
        "  --enable-protocol=https \\\n",
        "  --enable-gnutls \\\n",
        "  --enable-shared \\\n",
        "  --enable-gpl \\\n",
        "  --enable-nonfree \\\n",
        "  --enable-cuda-nvcc \\\n",
        "  --enable-libx264 \\\n",
        "  --enable-nvenc \\\n",
        "  --enable-cuvid \\\n",
        "  --enable-nvdec"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFTCT16ufL-e"
      },
      "source": [
        "### Build and install FFmpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ki9PN3V8XbSh",
        "outputId": "47b40f5f-ec63-4827-fdac-a7e17f6eb1d8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INSTALL\tlibavdevice/libavdevice.so\n",
            "INSTALL\tlibavfilter/libavfilter.so\n",
            "INSTALL\tlibavformat/libavformat.so\n",
            "INSTALL\tlibavcodec/libavcodec.so\n",
            "INSTALL\tlibswresample/libswresample.so\n",
            "INSTALL\tlibswscale/libswscale.so\n",
            "INSTALL\tlibavutil/libavutil.so\n",
            "INSTALL\tinstall-progs-yes\n",
            "INSTALL\tffmpeg\n",
            "INSTALL\tffprobe\n"
          ]
        }
      ],
      "source": [
        "!cd ffmpeg && make clean && make -j > /dev/null 2>&1\n",
        "!cd ffmpeg && make install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l1A9bI0BkWCf"
      },
      "source": [
        "## Check FFmpeg installation\n",
        "\n",
        "Let's do a quick sanity check to confirm that the FFmpeg we built works."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4fu9POO7kbQL",
        "outputId": "e55f7443-3a7f-471c-c894-5677fd843ee7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " VFS..D h264                 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10\n",
            " V..... h264_cuvid           Nvidia CUVID H264 decoder (codec h264)\n"
          ]
        }
      ],
      "source": [
        "!ffprobe -hide_banner -decoders | grep h264"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "499UCyKtFPrO",
        "outputId": "2cd0f6a5-cde9-42b7-c30b-95c920a44720"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " V..... libx264              libx264 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10 (codec h264)\n",
            " V....D h264_nvenc           NVIDIA NVENC H.264 encoder (codec h264)\n"
          ]
        }
      ],
      "source": [
        "!ffmpeg -hide_banner -encoders | grep 264"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following command fetches video from remote server, decode with NVDEC (cuvid) and re-encode with NVENC. If this command does not work, then there is an issue with FFmpeg installation, and TorchAudio would not be able to use them either."
      ],
      "metadata": {
        "id": "yZunKfgUb-i_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A--jYQBrfcOr",
        "outputId": "f7f32af2-845d-4326-c9a7-6f3e05ca7231"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4':\n",
            "  Metadata:\n",
            "    major_brand     : mp42\n",
            "    minor_version   : 512\n",
            "    compatible_brands: mp42iso2avc1mp41\n",
            "    encoder         : Lavf58.76.100\n",
            "  Duration: 00:03:26.04, start: 0.000000, bitrate: 1294 kb/s\n",
            "  Stream #0:0(eng): Video: h264 (High) (avc1 / 0x31637661), yuv420p(tv, bt709), 960x540 [SAR 1:1 DAR 16:9], 1156 kb/s, 29.97 fps, 29.97 tbr, 30k tbn, 59.94 tbc (default)\n",
            "    Metadata:\n",
            "      handler_name    : ?Mainconcept Video Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n",
            "  Stream #0:1(eng): Audio: aac (LC) (mp4a / 0x6134706D), 48000 Hz, stereo, fltp, 128 kb/s (default)\n",
            "    Metadata:\n",
            "      handler_name    : #Mainconcept MP4 Sound Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n",
            "Stream mapping:\n",
            "  Stream #0:0 -> #0:0 (h264 (h264_cuvid) -> h264 (h264_nvenc))\n",
            "  Stream #0:1 -> #0:1 (copy)\n",
            "Press [q] to stop, [?] for help\n",
            "Output #0, mp4, to 'test.mp4':\n",
            "  Metadata:\n",
            "    major_brand     : mp42\n",
            "    minor_version   : 512\n",
            "    compatible_brands: mp42iso2avc1mp41\n",
            "    encoder         : Lavf58.76.100\n",
            "  Stream #0:0(eng): Video: h264 (Main) (avc1 / 0x31637661), cuda(tv, bt709, progressive), 360x240 [SAR 1:1 DAR 3:2], q=2-31, 5000 kb/s, 29.97 fps, 30k tbn (default)\n",
            "    Metadata:\n",
            "      handler_name    : ?Mainconcept Video Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n",
            "      encoder         : Lavc58.134.100 h264_nvenc\n",
            "    Side data:\n",
            "      cpb: bitrate max/min/avg: 0/0/5000000 buffer size: 10000000 vbv_delay: N/A\n",
            "  Stream #0:1(eng): Audio: aac (LC) (mp4a / 0x6134706D), 48000 Hz, stereo, fltp, 128 kb/s (default)\n",
            "    Metadata:\n",
            "      handler_name    : #Mainconcept MP4 Sound Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n",
            "frame= 6175 fps=1712 q=11.0 Lsize=   37935kB time=00:03:26.01 bitrate=1508.5kbits/s speed=57.1x    \n",
            "video:34502kB audio:3234kB subtitle:0kB other streams:0kB global headers:0kB muxing overhead: 0.526932%\n"
          ]
        }
      ],
      "source": [
        "!ffmpeg -hide_banner -y -vsync 0 -hwaccel cuvid -hwaccel_output_format cuda -c:v h264_cuvid -resize 360x240 -i \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\" -c:a copy -c:v h264_nvenc -b:v 5M test.mp4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYzJWphMG9nh"
      },
      "source": [
        "## Benchmarking GPU Encoding and Decoding\n",
        "\n",
        "Now that FFmpeg and the resulting libraries are ready to use, we test NVDEC/NVENC with TorchAudio. For the basics of TorchAudio's streaming APIs, please refer to [Media I/O tutorials](https://pytorch.org/audio/main/tutorials.io.html).\n",
        "\n",
        "**Note**\n",
        "\n",
        "If you rebuild FFmpeg after importing torchaudio, you'll need to restart the session to activate the newly built FFmpeg libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wDJ-qdM45fV6",
        "outputId": "9242a9bd-62e8-4217-8113-a18a62fbbe94"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.13.0.dev20221006+cu116\n",
            "0.13.0.dev20221006+cu116\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torchaudio\n",
        "\n",
        "print(torch.__version__)\n",
        "print(torchaudio.__version__)\n",
        "\n",
        "from torchaudio.io import StreamReader, StreamWriter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "zKMACIJNIXns"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "\n",
        "pd.set_option('display.max_rows', None)\n",
        "pd.set_option('display.max_columns', None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSCZFqt95GZm"
      },
      "source": [
        "## Benchmark NVDEC with `StreamReader`\n",
        "\n",
        "First we test hardware decoding, and we fetch video from multiple locations (local file, network file, AWS S3) and use NVDEC to decod them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lbA3S6p_Z2lI",
        "outputId": "818e6aef-048a-49d3-a29b-1e3f421dc8ef"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting boto3\n",
            "  Downloading boto3-1.24.88-py3-none-any.whl (132 kB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hCollecting s3transfer<0.7.0,>=0.6.0\n",
            "  Downloading s3transfer-0.6.0-py3-none-any.whl (79 kB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hCollecting jmespath<2.0.0,>=0.7.1\n",
            "  Downloading jmespath-1.0.1-py3-none-any.whl (20 kB)\n",
            "Collecting botocore<1.28.0,>=1.27.88\n",
            "  Downloading botocore-1.27.88-py3-none-any.whl (9.2 MB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hCollecting urllib3<1.27,>=1.25.4\n",
            "  Downloading urllib3-1.26.12-py2.py3-none-any.whl (140 kB)\n",
            "\u001b[?25l\n",
            "\u001b[?25hRequirement already satisfied: python-dateutil<3.0.0,>=2.1 in /usr/local/lib/python3.7/dist-packages (from botocore<1.28.0,>=1.27.88->boto3) (2.8.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil<3.0.0,>=2.1->botocore<1.28.0,>=1.27.88->boto3) (1.15.0)\n",
            "Installing collected packages: urllib3, jmespath, botocore, s3transfer, boto3\n",
            "  Attempting uninstall: urllib3\n",
            "    Found existing installation: urllib3 1.24.3\n",
            "    Uninstalling urllib3-1.24.3:\n",
            "      Successfully uninstalled urllib3-1.24.3\n",
            "Successfully installed boto3-1.24.88 botocore-1.27.88 jmespath-1.0.1 s3transfer-0.6.0 urllib3-1.26.12\n"
          ]
        }
      ],
      "source": [
        "!pip3 install --progress-bar off boto3 2> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g3TMksRk5ara",
        "outputId": "17191485-7188-4986-e540-c77df1c5afe0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.24.88\n"
          ]
        }
      ],
      "source": [
        "import boto3\n",
        "from botocore import UNSIGNED\n",
        "from botocore.config import Config\n",
        "\n",
        "print(boto3.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "d0GYPyGGTz-q"
      },
      "outputs": [],
      "source": [
        "!wget -q -O input.mp4 \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "go6xfv49doJ9"
      },
      "source": [
        "First, we define the functions we'll use for testing.\n",
        "\n",
        "Funcion `test_decode` decodes the given source from start to end, and it reports the elapsed time, and returns one image frmae as a sample."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "GsmoRAY75T6J"
      },
      "outputs": [],
      "source": [
        "result = torch.zeros((4, 2))\n",
        "samples = [[None, None] for _ in range(4)]\n",
        "\n",
        "\n",
        "def test_decode(src, config, i_sample):\n",
        "  print(\"=\" * 40)\n",
        "  print(\"* Configuration:\", config)\n",
        "  print(\"* Source:\", src)\n",
        "  print(\"=\" * 40)\n",
        "\n",
        "  s = StreamReader(src)\n",
        "  s.add_video_stream(5, **config)\n",
        "\n",
        "  t0 = time.monotonic()\n",
        "  num_frames = 0\n",
        "  for i, (chunk, ) in enumerate(s.stream()):\n",
        "    if i == 0:\n",
        "      print(' - Chunk:', chunk.shape, chunk.device, chunk.dtype)\n",
        "    if i == i_sample:\n",
        "      sample = chunk[0]\n",
        "    num_frames += chunk.shape[0]\n",
        "  elapsed = time.monotonic() - t0\n",
        "\n",
        "  print()\n",
        "  print(f\" - Processed {num_frames} frames.\")\n",
        "  print(f\" - Elapsed: {elapsed} seconds.\")\n",
        "  print()\n",
        "\n",
        "  return elapsed, sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGCQUW9gskBZ"
      },
      "source": [
        "### Decode MP4 from local file\n",
        "\n",
        "For the first test, we compare the time it takes for CPU and NVDEC to decode 250MB of MP4 video."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "mlrA9oVM3qvt"
      },
      "outputs": [],
      "source": [
        "local_src = \"input.mp4\"\n",
        "\n",
        "i_sample = 520"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9QScuPXO3_8e"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z6kivDretHtS",
        "outputId": "fc14c8a1-6f26-4370-b532-d90098d0f775"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 46.691246449000005 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "vAEDVY0D37bJ"
      },
      "outputs": [],
      "source": [
        "result[0, 0] = elapsed\n",
        "samples[0][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jn8Wj9TN3-iY"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V6jwiJD-tW0M",
        "outputId": "0f40c5b2-5e04-4ca1-9902-5a9ed9dc5820"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 6.117441406000012 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cuda_conf = {\n",
        "    \"decoder\": \"h264_cuvid\",  # Use CUDA HW decoder\n",
        "    \"hw_accel\": \"cuda:0\",  # Then keep the memory on CUDA:0\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "yIA6UAaw39VC"
      },
      "outputs": [],
      "source": [
        "result[0, 1] = elapsed\n",
        "samples[0][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYBrsBKstY0B"
      },
      "source": [
        "### Decode MP4 from network\n",
        "\n",
        "Let's run the same test on the source retrieved via network on-the-fly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "ICa21rL84QVb"
      },
      "outputs": [],
      "source": [
        "network_src = \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\"\n",
        "i_sample = 750"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RsHY7D84YZ3"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6AGfH8X8tboy",
        "outputId": "3039e4b6-35a2-4b55-e7a7-d8b1b63e84f6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 46.460909987000036 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test_decode(network_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "4u8glKk14bYT"
      },
      "outputs": [],
      "source": [
        "result[1, 0] = elapsed\n",
        "samples[1][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jRAnSGM4cYS"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "95bky0TqtfMI",
        "outputId": "14e2cea4-fe27-405b-feec-e276c466542a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 4.23164078800005 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test_decode(network_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "kir92rPk4eLr"
      },
      "outputs": [],
      "source": [
        "result[1, 1] = elapsed\n",
        "samples[1][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owPsmnyTMiQm"
      },
      "source": [
        "### Decode MP4 directly from S3\n",
        "\n",
        "Using file-like object input, we can fetch a video stored on AWS S3 and decode it without saving it on local file system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "XBYc5AiGORUU"
      },
      "outputs": [],
      "source": [
        "bucket = \"pytorch\"\n",
        "key = \"torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\"\n",
        "\n",
        "s3_client = boto3.client(\"s3\", config=Config(signature_version=UNSIGNED))\n",
        "i_sample = 115"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BG1fHJWFND4C"
      },
      "source": [
        "#### Defining Helper class\n",
        "\n",
        "StreamReader supports file-like objects with `read` method. In addition to this,\n",
        "if the file-like object has `seek` method, StreamReader attempts to use it for more reliable detection of media formats.\n",
        "\n",
        "However, the seek method of `boto3`'s S3 client response object only raises errors to let users know that seek operation is not supported. Therefore we wrap it with a class that does not have `seek` method. This way, StreamReader won't try to use the `seek` method.\n",
        "\n",
        "**Note**\n",
        "\n",
        "Due to the nature of streaming, when using file-like object without seek method, some formats are not supported. For example, MP4 formats contain metadata at the beginning of file or at the end. If metadata is located at the end, without `seek` method, StreamReader cannot decode streams."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "Odq5pMJ4M2ph"
      },
      "outputs": [],
      "source": [
        "# Wrapper to hide the native `seek` method of boto3, which\n",
        "# only raises an error.\n",
        "class UnseekableWrapper:\n",
        "  def __init__(self, obj):\n",
        "    self.obj = obj\n",
        "\n",
        "  def read(self, n):\n",
        "    return self.obj.read(n)\n",
        "\n",
        "  def __str__(self):\n",
        "    return str(self.obj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ftfQQLIpOBgR"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "rjDH0g0qOAhh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e8323115-a3f6-4525-9645-e20d77fadc95"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: <botocore.response.StreamingBody object at 0x7fb991dddfd0>\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 40.758733775999985 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = s3_client.get_object(Bucket=bucket, Key=key)\n",
        "src = UnseekableWrapper(response[\"Body\"])\n",
        "elapsed, sample = test_decode(src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "JfjvXqBqQGyu"
      },
      "outputs": [],
      "source": [
        "result[2, 0] = elapsed\n",
        "samples[2][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lh3_FqdSODsM"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "17-0MVlZOm-f",
        "outputId": "9604cb57-cdb1-4e72-8754-47dfa737f0f7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: <botocore.response.StreamingBody object at 0x7fb991d390d0>\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 4.620101478000038 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = s3_client.get_object(Bucket=bucket, Key=key)\n",
        "src = UnseekableWrapper(response[\"Body\"])\n",
        "elapsed, sample = test_decode(src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "VpDLx494QIN4"
      },
      "outputs": [],
      "source": [
        "result[2, 1] = elapsed\n",
        "samples[2][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jbCyeWFsvM6r"
      },
      "source": [
        "### Decoding and resizing\n",
        "\n",
        "In the next test, we add preprocessing. NVDEC supports several preprocessing schemes, which are also performed on the chosen hardware. For CPU, we apply the same kind of software preprocessing through FFmpeg's filter graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "XKpuH_tXvMSF"
      },
      "outputs": [],
      "source": [
        "i_sample = 1085"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRL3C-bm4ijE"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OEtAVQVYwVCy",
        "outputId": "29fd9dca-5378-48a3-9ef3-26ab1b892107"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264', 'filter_desc': 'scale=360:240'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 240, 360]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 19.082725973000038 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "    \"filter_desc\": \"scale=360:240\",  # Software filter\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "uzZMz0rW4j73"
      },
      "outputs": [],
      "source": [
        "result[3, 0] = elapsed\n",
        "samples[3][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xkrBNRNP4laK"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "XajhZb-G4mwm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "29038cfd-2d4e-4b4b-a40f-ef6731ddff96"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'decoder_option': {'resize': '360x240'}, 'hw_accel': 'cuda:0'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 240, 360]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 4.157032522999998 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cuda_conf = {\n",
        "    \"decoder\": \"h264_cuvid\",  # Use CUDA HW decoder\n",
        "    \"decoder_option\": {\n",
        "        \"resize\": \"360x240\",  # Then apply HW preprocessing (resize)\n",
        "    },\n",
        "    \"hw_accel\": \"cuda:0\",  # Then keep the memory on CUDA:0\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "5mXoy38TwYVx"
      },
      "outputs": [],
      "source": [
        "result[3, 1] = elapsed\n",
        "samples[3][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEA2jrrKqmr4"
      },
      "source": [
        "### Results\n",
        "\n",
        "The following table summarizes the time it took to decode the same media with CPU and NVDEC.\n",
        "We see significant speedup with NVDEC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "id": "Ni0bihmogZAb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b586825a-1164-4b9e-f5ca-79febf95628a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "                                       CPU     NVDEC\n",
            "Decoding (local file)            46.691246  6.117441\n",
            "Decoding (network file)          46.460911  4.231641\n",
            "Decoding (file-like object, S3)  40.758736  4.620101\n",
            "Decoding + Resize                19.082726  4.157032\n"
          ]
        }
      ],
      "source": [
        "res = pd.DataFrame(\n",
        "    result.numpy(),\n",
        "    index=[\"Decoding (local file)\", \"Decoding (network file)\", \"Decoding (file-like object, S3)\", \"Decoding + Resize\"],\n",
        "    columns=[\"CPU\", \"NVDEC\"],\n",
        ")\n",
        "print(res)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i71SJL8PmZfE"
      },
      "source": [
        "The following code shows some frames generated by CPU decoding and NVDEC. They produce seemingly identical results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "nWTT5ih5v1s6"
      },
      "outputs": [],
      "source": [
        "def yuv_to_rgb(img):\n",
        "  img = img.cpu().to(torch.float)\n",
        "  y = img[..., 0, :, :]\n",
        "  u = img[..., 1, :, :]\n",
        "  v = img[..., 2, :, :]\n",
        "\n",
        "  y /= 255\n",
        "  u = u / 255 - 0.5\n",
        "  v = v / 255 - 0.5\n",
        "\n",
        "  r = y + 1.14 * v\n",
        "  g = y + -0.396 * u - 0.581 * v\n",
        "  b = y + 2.029 * u\n",
        "\n",
        "  rgb = torch.stack([r, g, b], -1)\n",
        "  rgb = (rgb * 255).clamp(0, 255).to(torch.uint8)\n",
        "  return rgb.numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "5AL9u6_xmRQa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "89bdea08-4d2c-4192-874c-6e1b316ea63a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[]"
            ]
          },
          "metadata": {},
          "execution_count": 40
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 921.6x1382.4 with 8 Axes>"
            ],
            "image/png": "[encoded data, 674568 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "f, axs = plt.subplots(4, 2, figsize=[12.8, 19.2])\n",
        "for i in range(4):\n",
        "  for j in range(2):\n",
        "    axs[i][j].imshow(yuv_to_rgb(samples[i][j]))\n",
        "    axs[i][j].set_title(\n",
        "        f\"{'CPU' if j == 0 else 'NVDEC'}{' with resize' if i == 3 else ''}\")\n",
        "plt.plot(block=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eXzKSVrHk1vS"
      },
      "source": [
        "## Benchmark NVENC with `StreamWriter`\n",
        "\n",
        "Next, we benchmark encoding speed with StreamWriter and NVENC.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "k5WkyQEp0hjv"
      },
      "outputs": [],
      "source": [
        "def test_encode(data, dst, **config):\n",
        "  print(\"=\" * 40)\n",
        "  print(\"* Configuration:\", config)\n",
        "  print(\"* Destination:\", dst)\n",
        "  print(\"=\" * 40)\n",
        "\n",
        "  s = StreamWriter(dst)\n",
        "  s.add_video_stream(**config)\n",
        "\n",
        "  t0 = time.monotonic()\n",
        "  with s.open():\n",
        "    s.write_video_chunk(0, data)\n",
        "  elapsed = time.monotonic() - t0\n",
        "\n",
        "  print()\n",
        "  print(f\" - Processed {len(data)} frames.\")\n",
        "  print(f\" - Elapsed: {elapsed} seconds.\")\n",
        "  print()\n",
        "  return elapsed\n",
        "\n",
        "result = torch.zeros((3, 3))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We use ``StreamReader`` to generate test data."
      ],
      "metadata": {
        "id": "wx3cQ7YwX880"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "nHynkI-oK9mX"
      },
      "outputs": [],
      "source": [
        "def get_data(frame_rate, height, width, format, duration=15):\n",
        "  src = f\"testsrc2=rate={frame_rate}:size={width}x{height}:duration={duration}\"\n",
        "  s = StreamReader(src=src, format=\"lavfi\")\n",
        "  s.add_basic_video_stream(-1, format=format)\n",
        "  s.process_all_packets()\n",
        "  video, = s.pop_chunks()\n",
        "  return video"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IQafz-HEQtu"
      },
      "source": [
        "### Encode MP4 - 360P\n",
        "\n",
        "For the first test, we compare the time it takes for CPU and NVENC to encode 15 seconds of video with small resolution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "2dZk8Ui7EZew"
      },
      "outputs": [],
      "source": [
        "pict_config = {\n",
        "    \"height\": 360,\n",
        "    \"width\": 640,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZWLHmf8H3Q5"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "id": "wP648bHCH9p_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "30af3ebc-a27c-4c40-a533-b4ac8e977527"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 360p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 3.280829835000077 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 0] = test_encode(video, \"360p_cpu.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P731iWsIH45U"
      },
      "source": [
        "#### CUDA (from CPU Tensor)\n",
        "\n",
        "Now we test NVENC. This time, the data is sent from CPU memory to GPU memory as part of encoding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "zQ9W_UVTIf-E",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3a9ccbfb-d243-4f1f-a790-0595a1de8bba"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}}\n",
            "* Destination: 360p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.34294435300000714 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",  # Use NVENC\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},  # Run encoding on the cuda:0 device\n",
        "}\n",
        "\n",
        "result[1, 0] = test_encode(video, \"360p_cuda.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QirJIHuSm-WT"
      },
      "source": [
        "#### CUDA (from CUDA Tensor)\n",
        "\n",
        "If the data is already present on CUDA, then we can pass it to GPU encoder directly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "EHjyMYpJnEb3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b3a4c529-9b6e-468e-a836-6e5f5fe51d56"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 360p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.2424524550000342 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",  # GPU Encoder\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},  # Run encoding on the cuda:0 device\n",
        "    \"hw_accel\": device,  # Data comes from cuda:0 device\n",
        "}\n",
        "\n",
        "result[2, 0] = test_encode(video.to(torch.device(device)), \"360p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBCBSKkRJeiA"
      },
      "source": [
        "### Encode MP4 - 720P\n",
        "\n",
        "Let's run the same tests on video with larger resolution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "CUdzN9BzJlo_"
      },
      "outputs": [],
      "source": [
        "pict_config = {\n",
        "    \"height\": 720,\n",
        "    \"width\": 1280,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wTjAcXGAJsE8"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FoEM76wFJta5",
        "outputId": "abe666de-8d02-469a-beea-74fcff494b89"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 720p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 11.638768525999922 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 1] = test_encode(video, \"720p_cpu.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lgbIN67BKARA"
      },
      "source": [
        "#### CUDA (from CPU Tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fgNR79olJ7Yo",
        "outputId": "65484cc4-ab54-4373-b1a4-50bfea20e73d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 720p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.8508033889999069 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[1, 1] = test_encode(video, \"720p_cuda.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oD_XUn4-KDod"
      },
      "source": [
        "#### CUDA (from CUDA Tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2NX4SU5xJ7zc",
        "outputId": "130ad5a6-a197-46dc-d799-d230e2b469c0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 720p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.6384492569999338 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},\n",
        "    \"hw_accel\": device,\n",
        "}\n",
        "\n",
        "result[2, 1] = test_encode(video.to(torch.device(device)), \"720p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Encode MP4 - 1080P\n",
        "\n",
        "We make the video with even larger."
      ],
      "metadata": {
        "id": "FoMDEhtYO-xD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pict_config = {\n",
        "    \"height\": 1080,\n",
        "    \"width\": 1920,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ],
      "metadata": {
        "id": "klcLfXipO5OX"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CPU"
      ],
      "metadata": {
        "id": "o28eOGRuPD3n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 2] = test_encode(video, \"1080p_cpu.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6UuiAEowOedt",
        "outputId": "2bcc2c9b-dc27-4dd5-b469-44e760a0068a"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 1080p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 27.020421489 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CUDA (from CPU Tensor)"
      ],
      "metadata": {
        "id": "ACJ3oEwgPFAh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[1, 2] = test_encode(video, \"1080p_cuda.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dzYNn8HtOeVi",
        "outputId": "53b8e871-360f-483a-c7d6-c217a913dae9"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 1080p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 1.60377999800005 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CUDA (from CUDA Tensor)"
      ],
      "metadata": {
        "id": "YcEH15BYPOQ_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},\n",
        "    \"hw_accel\": device,\n",
        "}\n",
        "\n",
        "result[2, 2] = test_encode(video.to(torch.device(device)), \"1080p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cUZOWVllOeRl",
        "outputId": "b7608edc-ed60-4bc9-ab44-fa0b333a4264"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 1080p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 1.4101193979998925 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Result\n",
        "\n",
        "Here is the result."
      ],
      "metadata": {
        "id": "Z9lEeklfPY-Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "labels = [\"CPU\", \"CUDA (from CPU Tensor)\", \"CUDA (from CUDA Tensor)\"]\n",
        "columns = [\"360P\", \"720P\", \"1080P\"]\n",
        "res = pd.DataFrame(\n",
        "    result.numpy(),\n",
        "    index=labels,\n",
        "    columns=columns,\n",
        ")\n",
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jDWUQFjVPaUV",
        "outputId": "eda3af22-a346-4e0a-fd48-f1227e524ddf"
      },
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "                             360P       720P      1080P\n",
            "CPU                      3.280830  11.638768  27.020422\n",
            "CUDA (from CPU Tensor)   0.342944   0.850803   1.603780\n",
            "CUDA (from CUDA Tensor)  0.242452   0.638449   1.410119\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "plt.plot(result.T)\n",
        "plt.legend(labels)\n",
        "plt.xticks([i for i in range(3)], columns)\n",
        "plt.grid(visible=True, axis='y')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "gEccR5zwTYRQ",
        "outputId": "486af2b2-9137-4197-bc6b-b17af0b2075a"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "[encoded data, 15732 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The resulting videos look like the following."
      ],
      "metadata": {
        "id": "mVeuM2gmkKld"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "\n",
        "HTML('''\n",
        "<div>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cpu.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda_hw.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "</div>\n",
        "''')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 431
        },
        "id": "LNPWCXxOgnqz",
        "outputId": "cd9449d8-d464-4c0e-88c3-e71d7b3261e5"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "<div>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cpu.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda_hw.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "</div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 61
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yuYzPpV9zc4M"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "We looked at how to build FFmpeg libraries with NVDEC/NVENC support and use them from TorchAudio. NVDEC/NVENC provide significant speed up when saving/loading a video."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "VcKsO-zurtWM",
        "PWT0m_LGq_GE",
        "TpWDpxBvDy-b"
      ],
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
