TensorStream

TensorStream is a C++ library for real-time video stream (e.g., RTMP) decoding to CUDA memory which supports some additional features:

CUDA memory conversion to ATen Tensor for using it via Python in PyTorch Deep Learning models
Detecting basic video stream issues related to frames reordering/loss
Video Post Processing (VPP) operations: downscaling/upscaling, color conversion from NV12 to RGB24/BGR24/Y800
Support Linux and Windows

Simple example how to use TensorStream for deep learning tasks:

from tensor_stream import TensorStreamConverter, FourCC

reader = TensorStreamConverter("rtmp://127.0.0.1/live")
reader.initialize()
reader.start()

while need_predictions:
    # read latest available frame from the stream 
    tensor = reader.read(pixel_format=FourCC.BGR24,
                         width=256,
                         height=256)
                         
    # tensor dtype is torch.uint8, device is cuda, shape is (256, 256, 3)
    prediction = model(tensor)
    ...

Initialize tensor stream with a video (e.g., a local file or a network video stream) and start reading it in a separate process.
Get latest available frame from the stream and make a prediction.

Note: All tasks inside TensorStream processed on a GPU, so output tensor also located on the GPU.

Installation
Usage
Docker
Documentation
License

Install TensorStream

Dependencies

NVIDIA CUDA 9.0 or above
FFmpeg and FFmpeg version of headers required to interface with Nvidias codec APIs nv-codec-headers
PyTorch 1.0.1.post2 or above to build C++ extension for Python
Python 3.6 or above to build C++ extension for Python

It is convenient to use TensorStream in Docker containers. The provided Dockerfiles is supplied to create an image with all the necessary dependencies.

Installation from source

TensorStream source code

git clone -b master --single-branch https://github.com/Fonbet/argus-tensor-stream.git
cd argus-tensor-stream

C++ extension for Python

On Linux:

python setup.py install

On Windows:

set FFMPEG_PATH="Path to FFmpeg install folder"
set path=%path%;%FFMPEG_PATH%\bin
set VS150COMNTOOLS="Path to Visual Studio vcvarsall.bat folder"
call "%VS150COMNTOOLS%\vcvarsall.bat" x64 -vcvars_ver=14.11
python setup.py install

To build TensorStream on Windows, Visual Studio 2017 14.11 toolset is required

C++ library:

On Linux:

mkdir build
cd build
cmake ..

On Windows:

set FFMPEG_PATH="Path to FFmpeg install folder"
mkdir build
cd build
cmake -G "Visual Studio 15 2017 Win64" -T v141,version=14.11 ..

Binaries (Linux only)

Extension for Python can be installed via pip:

CUDA 9:

pip install https://tensorstream.argus-ai.com/wheel/cu9/linux/tensor_stream-0.1.8-cp36-cp36m-linux_x86_64.whl

CUDA 10:

pip install https://tensorstream.argus-ai.com/wheel/cu10/linux/tensor_stream-0.1.8-cp36-cp36m-linux_x86_64.whl

Building examples and tests

Examples for Python and C++ can be found in c_examples and python_examples folders. Tests for C++ can be found in tests folder.

Python example

Can be executed via Python after TensorStream C++ extension for Python installation.

cd python_examples
python simple.py

C++ example and unit tests

On Linux

cd c_examples  # tests
mkdir build
cd build
cmake ..

On Windows

set FFMPEG_PATH="Path to FFmpeg install folder"
cd c_examples or tests
mkdir build
cd build
cmake -G "Visual Studio 15 2017 Win64" -T v141,version=14.11 ..

Docker image

Dockerfiles can be found in docker folder. Please note that for different CUDAs different Dockerfiles are required. To distinguish them name suffix is used, i.e. for CUDA 9 Dockerfile name is Dockerfile_cu9, for CUDA 10 Dockerfile_cu10 and so on.

docker build -t tensorstream -f docker/Dockerfile_cu10 .

Run with bash command line and follow installation guide

nvidia-docker run -ti tensorstream bash

Usage

Samples

Simple example demonstrates RTMP to PyTorch tensor conversion. Let's consider some usage scenarios:

Note: You can pass --help to get the list of all available options, their description and default values

Convert an RTMP bitstream to RGB24 PyTorch tensors and dump the result to a dump.yuv file:

python simple.py -i rtmp://184.72.239.149/vod/mp4:bigbuckbunny_1500.mp4 -fc RGB24 -o dump.yuv

Warning: Dumps significantly affect performance

The same scenario with downscaling:

python simple.py -i rtmp://184.72.239.149/vod/mp4:bigbuckbunny_1500.mp4 -fc RGB24 -w 720 -h 480 -o dump.yuv

Number of frames to process can be limited by -n option:

python simple.py -i rtmp://184.72.239.149/vod/mp4:bigbuckbunny_1500.mp4 -fc RGB24 -w 720 -h 480 -o dump.yuv -n 100

Example demonstrates how to use TensorStream in case of several stream consumers:

python many_consumers.py -i rtmp://184.72.239.149/vod/mp4:bigbuckbunny_1500.mp4 -n 100

PyTorch example

Real-time video style transfer example: fast-neural-style.

Documentation

Documentation for Python and C++ API can be found on the site.

License

TensorStream is LGPL-2.1 licensed, see the LICENSE file for details.

saranshkarira / argus-tensor-stream