Load TensorFlow and PyTorch Models in FFmpeg SR Filter

This article provides a practical guide on how to load and execute deep learning models using the FFmpeg super-resolution (sr) filter. It covers the preparation of TensorFlow and PyTorch model files, the required format conversions, and the exact FFmpeg commands needed to run super-resolution processing on your video files.

FFmpeg DNN Backends

The FFmpeg sr (super-resolution) filter relies on FFmpeg’s internal Deep Neural Network (DNN) interface. This interface supports three backends: 1. TensorFlow: Loads frozen TensorFlow models in .pb format. 2. OpenVINO: Loads Intel OpenVINO models in .xml and .bin formats. 3. Native: Loads simple models converted into FFmpeg’s native .model format.

Because FFmpeg does not natively support PyTorch (.pt or .pth) files, PyTorch models must be converted to an intermediate format before they can be loaded.

Loading a TensorFlow Model

To use a TensorFlow model with the sr filter, the model must be exported as a frozen graph (a .pb file) with fixed input and output tensor names.

Run the following command to apply the TensorFlow model to a video:

ffmpeg -i input.mp4 -vf "sr=dnn_backend=tensorflow:model=super_resolution_model.pb" output.mp4

Key Parameters: * dnn_backend=tensorflow: Instructs FFmpeg to use the TensorFlow engine. * model=super_resolution_model.pb: The path to your frozen TensorFlow model file.

Loading a PyTorch Model

FFmpeg cannot load PyTorch .pth files directly. To use a PyTorch model, you must convert it to ONNX format first, and then convert that ONNX model into either a TensorFlow .pb file or an OpenVINO format.

Step 1: Export PyTorch to ONNX

Use Python to export your PyTorch model to ONNX:

import torch

# Load your PyTorch model
model = MySRModel()
model.load_state_dict(torch.load("model.pth"))
model.eval()

# Create dummy input matching the model's expected input shape
dummy_input = torch.randn(1, 3, 360, 640) 

# Export to ONNX
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

Step 2: Convert ONNX to TensorFlow (.pb)

You can use the onnx-tf library in Python to convert the ONNX file into a TensorFlow format, which you can then freeze into a .pb file.

Step 3: Run the Converted Model in FFmpeg

Once converted, run the model using the appropriate backend. If you converted the PyTorch model to OpenVINO format (which often yields the best CPU performance), use the following command:

ffmpeg -i input.mp4 -vf "sr=dnn_backend=openvino:model=model.xml" output.mp4

Specifying Options and Scale Factor

The sr filter allows you to define the scale factor and specify the input/output layer names if your model uses non-standard tensor names.

ffmpeg -i input.mp4 -vf "sr=dnn_backend=tensorflow:model=model.pb:scale_factor=2:input=input_tensor_name:output=output_tensor_name" output.mp4