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.
- 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.
git clone -b master --single-branch https://github.com/Fonbet/argus-tensor-stream.git
cd argus-tensor-stream
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
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 ..
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
Examples for Python and C++ can be found in c_examples and python_examples folders. Tests for C++ can be found in tests folder.
Can be executed via Python after TensorStream C++ extension for Python installation.
cd python_examples
python simple.py
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 ..
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
- 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
Real-time video style transfer example: fast-neural-style.
Documentation for Python and C++ API can be found on the site.
TensorStream is LGPL-2.1 licensed, see the LICENSE file for details.
Big Buck Bunny is licensed under the Creative Commons Attribution 3.0 license. (c) copyright 2008, Blender Foundation / www.bigbuckbunny.org