virtualenv/FFmpeg
1
0
Fork 0
mirror of https://github.com/FFmpeg/FFmpeg.git synced 2024-11-21 10:55:51 +02:00
Code Issues Releases Activity

Adds ESPCN super resolution filter merged with SRCNN filter.

Browse Source 
Signed-off-by: Pedro Arthur <bygrandao@gmail.com>
This commit is contained in:
Sergey Lavrushkin 2018-06-14 00:37:12 +03:00 committed by Pedro Arthur
parent d24c9e55f6
commit 575b718990
10 changed files with 13255 additions and 385 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
configure vendored
View File

@ -260,7 +260,7 @@ External library support:
--enable-libsrt enable Haivision SRT protocol via libsrt [no]
--enable-libssh enable SFTP protocol via libssh [no]
--enable-libtensorflow enable TensorFlow as a DNN module backend
for DNN based filters like srcnn [no]
for DNN based filters like sr [no]
--enable-libtesseract enable Tesseract, needed for ocr filter [no]
--enable-libtheora enable Theora encoding via libtheora [no]
--enable-libtls enable LibreSSL (via libtls), needed for https support
@ -3402,8 +3402,8 @@ spectrumsynth_filter_deps="avcodec"
spectrumsynth_filter_select="fft"
spp_filter_deps="gpl avcodec"
spp_filter_select="fft idctdsp fdctdsp me_cmp pixblockdsp"
srcnn_filter_deps="avformat"
srcnn_filter_select="dnn"
sr_filter_deps="avformat swscale"
sr_filter_select="dnn"
stereo3d_filter_deps="gpl"
subtitles_filter_deps="avformat avcodec libass"
super2xsai_filter_deps="gpl"
@ -6823,7 +6823,7 @@ enabled signature_filter && prepend avfilter_deps "avcodec avformat"
enabled smartblur_filter && prepend avfilter_deps "swscale"
enabled spectrumsynth_filter && prepend avfilter_deps "avcodec"
enabled spp_filter && prepend avfilter_deps "avcodec"
enabled srcnn_filter && prepend avfilter_deps "avformat"
enabled sr_filter && prepend avfilter_deps "avformat"
enabled subtitles_filter && prepend avfilter_deps "avformat avcodec"
enabled uspp_filter && prepend avfilter_deps "avcodec"
enabled zoompan_filter && prepend avfilter_deps "swscale"

2
libavfilter/Makefile
View File

@ -340,7 +340,7 @@ OBJS-$(CONFIG_SMARTBLUR_FILTER) += vf_smartblur.o
OBJS-$(CONFIG_SOBEL_FILTER) += vf_convolution.o
OBJS-$(CONFIG_SPLIT_FILTER) += split.o
OBJS-$(CONFIG_SPP_FILTER) += vf_spp.o
OBJS-$(CONFIG_SRCNN_FILTER) += vf_srcnn.o
OBJS-$(CONFIG_SR_FILTER) += vf_sr.o
OBJS-$(CONFIG_SSIM_FILTER) += vf_ssim.o framesync.o
OBJS-$(CONFIG_STEREO3D_FILTER) += vf_stereo3d.o
OBJS-$(CONFIG_STREAMSELECT_FILTER) += f_streamselect.o framesync.o

2
libavfilter/allfilters.c
View File

@ -328,7 +328,7 @@ extern AVFilter ff_vf_smartblur;
extern AVFilter ff_vf_sobel;
extern AVFilter ff_vf_split;
extern AVFilter ff_vf_spp;
extern AVFilter ff_vf_srcnn;
extern AVFilter ff_vf_sr;
extern AVFilter ff_vf_ssim;
extern AVFilter ff_vf_stereo3d;
extern AVFilter ff_vf_streamselect;

295
libavfilter/dnn_backend_native.c
View File

@ -25,9 +25,12 @@
#include "dnn_backend_native.h"
#include "dnn_srcnn.h"
#include "dnn_espcn.h"
#include "libavformat/avio.h"
typedef enum {INPUT, CONV} LayerType;
typedef enum {INPUT, CONV, DEPTH_TO_SPACE} LayerType;
typedef enum {RELU, TANH, SIGMOID} ActivationFunc;
typedef struct Layer{
LayerType type;
@ -37,6 +40,7 @@ typedef struct Layer{
typedef struct ConvolutionalParams{
int32_t input_num, output_num, kernel_size;
ActivationFunc activation;
float* kernel;
float* biases;
} ConvolutionalParams;
@ -45,17 +49,22 @@ typedef struct InputParams{
int height, width, channels;
} InputParams;
typedef struct DepthToSpaceParams{
int block_size;
} DepthToSpaceParams;
// Represents simple feed-forward convolutional network.
typedef struct ConvolutionalNetwork{
Layer* layers;
int32_t layers_num;
} ConvolutionalNetwork;
static DNNReturnType set_input_output_native(void* model, const DNNData* input, const DNNData* output)
static DNNReturnType set_input_output_native(void* model, DNNData* input, DNNData* output)
{
ConvolutionalNetwork* network = (ConvolutionalNetwork*)model;
InputParams* input_params;
ConvolutionalParams* conv_params;
DepthToSpaceParams* depth_to_space_params;
int cur_width, cur_height, cur_channels;
int32_t layer;
@ -63,11 +72,17 @@ static DNNReturnType set_input_output_native(void* model, const DNNData* input,
return DNN_ERROR;
}
else{
network->layers[0].output = input->data;
input_params = (InputParams*)network->layers[0].params;
input_params->width = cur_width = input->width;
input_params->height = cur_height = input->height;
input_params->channels = cur_channels = input->channels;
if (input->data){
av_freep(&input->data);
}
network->layers[0].output = input->data = av_malloc(cur_height * cur_width * cur_channels * sizeof(float));
if (!network->layers[0].output){
return DNN_ERROR;
}
}
for (layer = 1; layer < network->layers_num; ++layer){
@ -78,32 +93,40 @@ static DNNReturnType set_input_output_native(void* model, const DNNData* input,
return DNN_ERROR;
}
cur_channels = conv_params->output_num;
if (layer < network->layers_num - 1){
if (!network->layers[layer].output){
av_freep(&network->layers[layer].output);
}
network->layers[layer].output = av_malloc(cur_height * cur_width * cur_channels * sizeof(float));
if (!network->layers[layer].output){
return DNN_ERROR;
}
}
else{
network->layers[layer].output = output->data;
if (output->width != cur_width || output->height != cur_height || output->channels != cur_channels){
return DNN_ERROR;
}
break;
case DEPTH_TO_SPACE:
depth_to_space_params = (DepthToSpaceParams*)network->layers[layer].params;
if (cur_channels % (depth_to_space_params->block_size * depth_to_space_params->block_size) != 0){
return DNN_ERROR;
}
cur_channels = cur_channels / (depth_to_space_params->block_size * depth_to_space_params->block_size);
cur_height *= depth_to_space_params->block_size;
cur_width *= depth_to_space_params->block_size;
break;
default:
return DNN_ERROR;
}
if (network->layers[layer].output){
av_freep(&network->layers[layer].output);
}
network->layers[layer].output = av_malloc(cur_height * cur_width * cur_channels * sizeof(float));
if (!network->layers[layer].output){
return DNN_ERROR;
}
}
output->data = network->layers[network->layers_num - 1].output;
output->height = cur_height;
output->width = cur_width;
output->channels = cur_channels;
return DNN_SUCCESS;
}
// Loads model and its parameters that are stored in a binary file with following structure:
// layers_num,conv_input_num,conv_output_num,conv_kernel_size,conv_kernel,conv_biases,conv_input_num...
// layers_num,layer_type,layer_parameterss,layer_type,layer_parameters...
// For CONV layer: activation_function, input_num, output_num, kernel_size, kernel, biases
// For DEPTH_TO_SPACE layer: block_size
DNNModel* ff_dnn_load_model_native(const char* model_filename)
{
DNNModel* model = NULL;
@ -111,7 +134,9 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
AVIOContext* model_file_context;
int file_size, dnn_size, kernel_size, i;
int32_t layer;
LayerType layer_type;
ConvolutionalParams* conv_params;
DepthToSpaceParams* depth_to_space_params;
model = av_malloc(sizeof(DNNModel));
if (!model){
@ -156,39 +181,62 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
}
for (layer = 1; layer < network->layers_num; ++layer){
conv_params = av_malloc(sizeof(ConvolutionalParams));
if (!conv_params){
layer_type = (int32_t)avio_rl32(model_file_context);
dnn_size += 4;
switch (layer_type){
case CONV:
conv_params = av_malloc(sizeof(ConvolutionalParams));
if (!conv_params){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
conv_params->activation = (int32_t)avio_rl32(model_file_context);
conv_params->input_num = (int32_t)avio_rl32(model_file_context);
conv_params->output_num = (int32_t)avio_rl32(model_file_context);
conv_params->kernel_size = (int32_t)avio_rl32(model_file_context);
kernel_size = conv_params->input_num * conv_params->output_num *
conv_params->kernel_size * conv_params->kernel_size;
dnn_size += 16 + (kernel_size + conv_params->output_num << 2);
if (dnn_size > file_size || conv_params->input_num <= 0 ||
conv_params->output_num <= 0 || conv_params->kernel_size <= 0){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
conv_params->kernel = av_malloc(kernel_size * sizeof(float));
conv_params->biases = av_malloc(conv_params->output_num * sizeof(float));
if (!conv_params->kernel || !conv_params->biases){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
for (i = 0; i < kernel_size; ++i){
conv_params->kernel[i] = av_int2float(avio_rl32(model_file_context));
}
for (i = 0; i < conv_params->output_num; ++i){
conv_params->biases[i] = av_int2float(avio_rl32(model_file_context));
}
network->layers[layer].type = CONV;
network->layers[layer].params = conv_params;
break;
case DEPTH_TO_SPACE:
depth_to_space_params = av_malloc(sizeof(DepthToSpaceParams));
if (!depth_to_space_params){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
depth_to_space_params->block_size = (int32_t)avio_rl32(model_file_context);
dnn_size += 4;
network->layers[layer].type = DEPTH_TO_SPACE;
network->layers[layer].params = depth_to_space_params;
break;
default:
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
conv_params->input_num = (int32_t)avio_rl32(model_file_context);
conv_params->output_num = (int32_t)avio_rl32(model_file_context);
conv_params->kernel_size = (int32_t)avio_rl32(model_file_context);
kernel_size = conv_params->input_num * conv_params->output_num *
conv_params->kernel_size * conv_params->kernel_size;
dnn_size += 12 + (kernel_size + conv_params->output_num << 2);
if (dnn_size > file_size || conv_params->input_num <= 0 ||
conv_params->output_num <= 0 || conv_params->kernel_size <= 0){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
conv_params->kernel = av_malloc(kernel_size * sizeof(float));
conv_params->biases = av_malloc(conv_params->output_num * sizeof(float));
if (!conv_params->kernel || !conv_params->biases){
avio_closep(&model_file_context);
ff_dnn_free_model_native(&model);
return NULL;
}
for (i = 0; i < kernel_size; ++i){
conv_params->kernel[i] = av_int2float(avio_rl32(model_file_context));
}
for (i = 0; i < conv_params->output_num; ++i){
conv_params->biases[i] = av_int2float(avio_rl32(model_file_context));
}
network->layers[layer].type = CONV;
network->layers[layer].params = conv_params;
}
avio_closep(&model_file_context);
@ -203,7 +251,8 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
return model;
}
static int set_up_conv_layer(Layer* layer, const float* kernel, const float* biases, int32_t input_num, int32_t output_num, int32_t size)
static int set_up_conv_layer(Layer* layer, const float* kernel, const float* biases, ActivationFunc activation,
int32_t input_num, int32_t output_num, int32_t size)
{
ConvolutionalParams* conv_params;
int kernel_size;
@ -212,6 +261,7 @@ static int set_up_conv_layer(Layer* layer, const float* kernel, const float* bia
if (!conv_params){
return DNN_ERROR;
}
conv_params->activation = activation;
conv_params->input_num = input_num;
conv_params->output_num = output_num;
conv_params->kernel_size = size;
@ -236,6 +286,7 @@ DNNModel* ff_dnn_load_default_model_native(DNNDefaultModel model_type)
{
DNNModel* model = NULL;
ConvolutionalNetwork* network = NULL;
DepthToSpaceParams* depth_to_space_params;
int32_t layer;
model = av_malloc(sizeof(DNNModel));
@ -253,45 +304,68 @@ DNNModel* ff_dnn_load_default_model_native(DNNDefaultModel model_type)
switch (model_type){
case DNN_SRCNN:
network->layers_num = 4;
network->layers = av_malloc(network->layers_num * sizeof(Layer));
if (!network->layers){
av_freep(&network);
av_freep(&model);
return NULL;
}
for (layer = 0; layer < network->layers_num; ++layer){
network->layers[layer].output = NULL;
network->layers[layer].params = NULL;
}
network->layers[0].type = INPUT;
network->layers[0].params = av_malloc(sizeof(InputParams));
if (!network->layers[0].params){
ff_dnn_free_model_native(&model);
return NULL;
}
if (set_up_conv_layer(network->layers + 1, conv1_kernel, conv1_biases, 1, 64, 9) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 2, conv2_kernel, conv2_biases, 64, 32, 1) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 3, conv3_kernel, conv3_biases, 32, 1, 5) != DNN_SUCCESS){
ff_dnn_free_model_native(&model);
return NULL;
}
model->set_input_output = &set_input_output_native;
return model;
break;
case DNN_ESPCN:
network->layers_num = 5;
break;
default:
av_freep(&network);
av_freep(&model);
return NULL;
}
network->layers = av_malloc(network->layers_num * sizeof(Layer));
if (!network->layers){
av_freep(&network);
av_freep(&model);
return NULL;
}
for (layer = 0; layer < network->layers_num; ++layer){
network->layers[layer].output = NULL;
network->layers[layer].params = NULL;
}
network->layers[0].type = INPUT;
network->layers[0].params = av_malloc(sizeof(InputParams));
if (!network->layers[0].params){
ff_dnn_free_model_native(&model);
return NULL;
}
switch (model_type){
case DNN_SRCNN:
if (set_up_conv_layer(network->layers + 1, srcnn_conv1_kernel, srcnn_conv1_biases, RELU, 1, 64, 9) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 2, srcnn_conv2_kernel, srcnn_conv2_biases, RELU, 64, 32, 1) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 3, srcnn_conv3_kernel, srcnn_conv3_biases, RELU, 32, 1, 5) != DNN_SUCCESS){
ff_dnn_free_model_native(&model);
return NULL;
}
break;
case DNN_ESPCN:
if (set_up_conv_layer(network->layers + 1, espcn_conv1_kernel, espcn_conv1_biases, TANH, 1, 64, 5) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 2, espcn_conv2_kernel, espcn_conv2_biases, TANH, 64, 32, 3) != DNN_SUCCESS ||
set_up_conv_layer(network->layers + 3, espcn_conv3_kernel, espcn_conv3_biases, SIGMOID, 32, 4, 3) != DNN_SUCCESS){
ff_dnn_free_model_native(&model);
return NULL;
}
network->layers[4].type = DEPTH_TO_SPACE;
depth_to_space_params = av_malloc(sizeof(DepthToSpaceParams));
if (!depth_to_space_params){
ff_dnn_free_model_native(&model);
return NULL;
}
depth_to_space_params->block_size = 2;
network->layers[4].params = depth_to_space_params;
}
model->set_input_output = &set_input_output_native;
return model;
}
#define CLAMP_TO_EDGE(x, w) ((x) < 0 ? 0 : ((x) >= (w) ? (w - 1) : (x)))
static void convolve(const float* input, float* output, const ConvolutionalParams* conv_params, int32_t width, int32_t height)
static void convolve(const float* input, float* output, const ConvolutionalParams* conv_params, int width, int height)
{
int y, x, n_filter, ch, kernel_y, kernel_x;
int radius = conv_params->kernel_size >> 1;
@ -313,19 +387,53 @@ static void convolve(const float* input, float* output, const ConvolutionalParam
}
}
}
output[n_filter] = FFMAX(output[n_filter], 0.0);
switch (conv_params->activation){
case RELU:
output[n_filter] = FFMAX(output[n_filter], 0.0);
break;
case TANH:
output[n_filter] = 2.0f / (1.0f + exp(-2.0f * output[n_filter])) - 1.0f;
break;
case SIGMOID:
output[n_filter] = 1.0f / (1.0f + exp(-output[n_filter]));
}
}
output += conv_params->output_num;
}
}
}
static void depth_to_space(const float* input, float* output, int block_size, int width, int height, int channels)
{
int y, x, by, bx, ch;
int new_channels = channels / (block_size * block_size);
int output_linesize = width * channels;
int by_linesize = output_linesize / block_size;
int x_linesize = new_channels * block_size;
for (y = 0; y < height; ++y){
for (x = 0; x < width; ++x){
for (by = 0; by < block_size; ++by){
for (bx = 0; bx < block_size; ++bx){
for (ch = 0; ch < new_channels; ++ch){
output[by * by_linesize + x * x_linesize + bx * new_channels + ch] = input[ch];
}
input += new_channels;
}
}
}
output += output_linesize;
}
}
DNNReturnType ff_dnn_execute_model_native(const DNNModel* model)
{
ConvolutionalNetwork* network = (ConvolutionalNetwork*)model->model;
InputParams* input_params;
int cur_width, cur_height;
int cur_width, cur_height, cur_channels;
int32_t layer;
InputParams* input_params;
ConvolutionalParams* conv_params;
DepthToSpaceParams* depth_to_space_params;
if (network->layers_num <= 0 || network->layers[0].type != INPUT || !network->layers[0].output){
return DNN_ERROR;
@ -334,6 +442,7 @@ DNNReturnType ff_dnn_execute_model_native(const DNNModel* model)
input_params = (InputParams*)network->layers[0].params;
cur_width = input_params->width;
cur_height = input_params->height;
cur_channels = input_params->channels;
}
for (layer = 1; layer < network->layers_num; ++layer){
@ -342,7 +451,17 @@ DNNReturnType ff_dnn_execute_model_native(const DNNModel* model)
}
switch (network->layers[layer].type){
case CONV:
convolve(network->layers[layer - 1].output, network->layers[layer].output, (ConvolutionalParams*)network->layers[layer].params, cur_width, cur_height);
conv_params = (ConvolutionalParams*)network->layers[layer].params;
convolve(network->layers[layer - 1].output, network->layers[layer].output, conv_params, cur_width, cur_height);
cur_channels = conv_params->output_num;
break;
case DEPTH_TO_SPACE:
depth_to_space_params = (DepthToSpaceParams*)network->layers[layer].params;
depth_to_space(network->layers[layer - 1].output, network->layers[layer].output,
depth_to_space_params->block_size, cur_width, cur_height, cur_channels);
cur_height *= depth_to_space_params->block_size;
cur_width *= depth_to_space_params->block_size;
cur_channels /= depth_to_space_params->block_size * depth_to_space_params->block_size;
break;
case INPUT:
return DNN_ERROR;
@ -362,19 +481,13 @@ void ff_dnn_free_model_native(DNNModel** model)
{
network = (ConvolutionalNetwork*)(*model)->model;
for (layer = 0; layer < network->layers_num; ++layer){
switch (network->layers[layer].type){
case CONV:
if (layer < network->layers_num - 1){
av_freep(&network->layers[layer].output);
}
av_freep(&network->layers[layer].output);
if (network->layers[layer].type == CONV){
conv_params = (ConvolutionalParams*)network->layers[layer].params;
av_freep(&conv_params->kernel);
av_freep(&conv_params->biases);
av_freep(&conv_params);
break;
case INPUT:
av_freep(&network->layers[layer].params);
}
av_freep(&network->layers[layer].params);
}
av_freep(network);
av_freep(model);

74
libavfilter/dnn_backend_tf.c
View File

@ -25,6 +25,7 @@
#include "dnn_backend_tf.h"
#include "dnn_srcnn.h"
#include "dnn_espcn.h"
#include "libavformat/avio.h"
#include <tensorflow/c/c_api.h>
@ -35,9 +36,7 @@ typedef struct TFModel{
TF_Status* status;
TF_Output input, output;
TF_Tensor* input_tensor;
TF_Tensor* output_tensor;
const DNNData* input_data;
const DNNData* output_data;
DNNData* output_data;
} TFModel;
static void free_buffer(void* data, size_t length)
@ -78,13 +77,13 @@ static TF_Buffer* read_graph(const char* model_filename)
return graph_buf;
}
static DNNReturnType set_input_output_tf(void* model, const DNNData* input, const DNNData* output)
static DNNReturnType set_input_output_tf(void* model, DNNData* input, DNNData* output)
{
TFModel* tf_model = (TFModel*)model;
int64_t input_dims[] = {1, input->height, input->width, input->channels};
int64_t output_dims[] = {1, output->height, output->width, output->channels};
TF_SessionOptions* sess_opts;
const TF_Operation* init_op = TF_GraphOperationByName(tf_model->graph, "init");
TF_Tensor* output_tensor;
// Input operation should be named 'x'
tf_model->input.oper = TF_GraphOperationByName(tf_model->graph, "x");
@ -100,6 +99,7 @@ static DNNReturnType set_input_output_tf(void* model, const DNNData* input, cons
if (!tf_model->input_tensor){
return DNN_ERROR;
}
input->data = (float*)TF_TensorData(tf_model->input_tensor);
// Output operation should be named 'y'
tf_model->output.oper = TF_GraphOperationByName(tf_model->graph, "y");
@ -107,17 +107,6 @@ static DNNReturnType set_input_output_tf(void* model, const DNNData* input, cons
return DNN_ERROR;
}
tf_model->output.index = 0;
if (tf_model->output_tensor){
TF_DeleteTensor(tf_model->output_tensor);
}
tf_model->output_tensor = TF_AllocateTensor(TF_FLOAT, output_dims, 4,
output_dims[1] * output_dims[2] * output_dims[3] * sizeof(float));
if (!tf_model->output_tensor){
return DNN_ERROR;
}
tf_model->input_data = input;
tf_model->output_data = output;
if (tf_model->session){
TF_CloseSession(tf_model->session, tf_model->status);
@ -144,6 +133,26 @@ static DNNReturnType set_input_output_tf(void* model, const DNNData* input, cons
}
}
// Execute network to get output height, width and number of channels
TF_SessionRun(tf_model->session, NULL,
&tf_model->input, &tf_model->input_tensor, 1,
&tf_model->output, &output_tensor, 1,
NULL, 0, NULL, tf_model->status);
if (TF_GetCode(tf_model->status) != TF_OK){
return DNN_ERROR;
}
else{
output->height = TF_Dim(output_tensor, 1);
output->width = TF_Dim(output_tensor, 2);
output->channels = TF_Dim(output_tensor, 3);
output->data = av_malloc(output->height * output->width * output->channels * sizeof(float));
if (!output->data){
return DNN_ERROR;
}
tf_model->output_data = output;
TF_DeleteTensor(output_tensor);
}
return DNN_SUCCESS;
}
@ -166,7 +175,7 @@ DNNModel* ff_dnn_load_model_tf(const char* model_filename)
}
tf_model->session = NULL;
tf_model->input_tensor = NULL;
tf_model->output_tensor = NULL;
tf_model->output_data = NULL;
graph_def = read_graph(model_filename);
if (!graph_def){
@ -215,6 +224,17 @@ DNNModel* ff_dnn_load_default_model_tf(DNNDefaultModel model_type)
graph_def->length = srcnn_tf_size;
graph_def->data_deallocator = free_buffer;
break;
case DNN_ESPCN:
graph_data = av_malloc(espcn_tf_size);
if (!graph_data){
TF_DeleteBuffer(graph_def);
return NULL;
}
memcpy(graph_data, espcn_tf_model, espcn_tf_size);
graph_def->data = (void*)graph_data;
graph_def->length = espcn_tf_size;
graph_def->data_deallocator = free_buffer;
break;
default:
TF_DeleteBuffer(graph_def);
return NULL;
@ -234,7 +254,7 @@ DNNModel* ff_dnn_load_default_model_tf(DNNDefaultModel model_type)
}
tf_model->session = NULL;
tf_model->input_tensor = NULL;
tf_model->output_tensor = NULL;
tf_model->output_data = NULL;
tf_model->graph = TF_NewGraph();
tf_model->status = TF_NewStatus();
@ -259,23 +279,21 @@ DNNModel* ff_dnn_load_default_model_tf(DNNDefaultModel model_type)
DNNReturnType ff_dnn_execute_model_tf(const DNNModel* model)
{
TFModel* tf_model = (TFModel*)model->model;
memcpy(TF_TensorData(tf_model->input_tensor), tf_model->input_data->data,
tf_model->input_data->height * tf_model->input_data->width *
tf_model->input_data->channels * sizeof(float));
TF_Tensor* output_tensor;
TF_SessionRun(tf_model->session, NULL,
&tf_model->input, &tf_model->input_tensor, 1,
&tf_model->output, &tf_model->output_tensor, 1,
&tf_model->output, &output_tensor, 1,
NULL, 0, NULL, tf_model->status);
if (TF_GetCode(tf_model->status) != TF_OK){
return DNN_ERROR;
}
else{
memcpy(tf_model->output_data->data, TF_TensorData(tf_model->output_tensor),
tf_model->output_data->height * tf_model->output_data->width *
tf_model->output_data->channels * sizeof(float));
memcpy(tf_model->output_data->data, TF_TensorData(output_tensor),
tf_model->output_data->height * tf_model->output_data->width *
tf_model->output_data->channels * sizeof(float));
TF_DeleteTensor(output_tensor);
return DNN_SUCCESS;
}
@ -300,9 +318,7 @@ void ff_dnn_free_model_tf(DNNModel** model)
if (tf_model->input_tensor){
TF_DeleteTensor(tf_model->input_tensor);
}
if (tf_model->output_tensor){
TF_DeleteTensor(tf_model->output_tensor);
}
av_freep(&tf_model->output_data->data);
av_freep(&tf_model);
av_freep(model);
}

12637
libavfilter/dnn_espcn.h Normal file
View File

File diff suppressed because it is too large Load Diff

4
libavfilter/dnn_interface.h
View File

@ -30,7 +30,7 @@ typedef enum {DNN_SUCCESS, DNN_ERROR} DNNReturnType;
typedef enum {DNN_NATIVE, DNN_TF} DNNBackendType;
typedef enum {DNN_SRCNN} DNNDefaultModel;
typedef enum {DNN_SRCNN, DNN_ESPCN} DNNDefaultModel;
typedef struct DNNData{
float* data;
@ -42,7 +42,7 @@ typedef struct DNNModel{
void* model;
// Sets model input and output, while allocating additional memory for intermediate calculations.
// Should be called at least once before model execution.
DNNReturnType (*set_input_output)(void* model, const DNNData* input, const DNNData* output);
DNNReturnType (*set_input_output)(void* model, DNNData* input, DNNData* output);
} DNNModel;
// Stores pointers to functions for loading, executing, freeing DNN models for one of the backends.

14
libavfilter/dnn_srcnn.h
View File

@ -20,13 +20,13 @@
/**
* @file
* Default cnn weights for x2 upsampling with srcnn filter.
* Default cnn weights for x2 upsampling with srcnn model.
*/
#ifndef AVFILTER_DNN_SRCNN_H
#define AVFILTER_DNN_SRCNN_H
static const float conv1_kernel[] = {
static const float srcnn_conv1_kernel[] = {
-0.08866338f, 0.055409566f, 0.037196506f, -0.11961404f,
-0.12341991f, 0.29963422f, -0.0911817f, -0.00013613555f,
-0.049023595f, 0.038421184f, -0.077267796f, 0.027273094f,
@ -1325,7 +1325,7 @@ static const float conv1_kernel[] = {
-0.013759381f, 0.026358005f, 0.088238746f, 0.082134426f
};
static const float conv1_biases[] = {
static const float srcnn_conv1_biases[] = {
-0.016606892f, -0.011107335f, -0.0048309686f, -0.04867378f,
-0.030040957f, -0.07297248f, -0.019458665f, -0.009738028f,
0.6951231f, -0.07369442f, -0.01354204f, 0.010336088f,
@ -1344,7 +1344,7 @@ static const float conv1_biases[] = {
0.054407462f, -0.08068252f, -0.009446503f, -0.04663234f
};
static const float conv2_kernel[] = {
static const float srcnn_conv2_kernel[] = {
-0.24004751f, 0.1037138f, 0.11173403f, 0.04352092f,
-0.23728481f, 0.12153747f, -0.23676059f, -0.28548065f,
-0.612738f, -0.12218937f, -0.06005159f, 0.1850652f,
@ -1859,7 +1859,7 @@ static const float conv2_kernel[] = {
0.11089696f, -0.08941251f, -0.3529318f, 0.0654588f
};
static const float conv2_biases[] = {
static const float srcnn_conv2_biases[] = {
0.12326373f, 0.13270757f, 0.07082674f, 0.051456157f,
0.058445618f, 0.13153197f, 0.0809729f, 0.10153213f,
0.055915363f, 0.05228166f, -0.11212896f, 0.07462141f,
@ -1870,7 +1870,7 @@ static const float conv2_biases[] = {
-0.086404406f, 0.06046943f, -0.1733751f, 0.2654999f
};
static const float conv3_kernel[] = {
static const float srcnn_conv3_kernel[] = {
-0.01733648f, 0.01492609f, 0.019393086f, -0.004445322f,
0.026939709f, 0.00038831023f, 0.004221528f, 0.0050745453f,
0.0129861f, 0.008007169f, 0.008950762f, 0.005279691f,
@ -2073,7 +2073,7 @@ static const float conv3_kernel[] = {
0.012931146f, 0.0046948805f, 0.013098622f, -0.015422701f
};
static const float conv3_biases[] = {
static const float srcnn_conv3_biases[] = {
0.05037664f
};

354
libavfilter/vf_sr.c Normal file
View File

@ -0,0 +1,354 @@
/*
* Copyright (c) 2018 Sergey Lavrushkin
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
/**
* @file
* Filter implementing image super-resolution using deep convolutional networks.
* https://arxiv.org/abs/1501.00092
* https://arxiv.org/abs/1609.05158
*/
#include "avfilter.h"
#include "formats.h"
#include "internal.h"
#include "libavutil/opt.h"
#include "libavformat/avio.h"
#include "libswscale/swscale.h"
#include "dnn_interface.h"
typedef enum {SRCNN, ESPCN} SRModel;
typedef struct SRContext {
const AVClass *class;
SRModel model_type;
char* model_filename;
DNNBackendType backend_type;
DNNModule* dnn_module;
DNNModel* model;
DNNData input, output;
int scale_factor;
struct SwsContext* sws_context;
int sws_slice_h;
} SRContext;
#define OFFSET(x) offsetof(SRContext, x)
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM
static const AVOption sr_options[] = {
{ "model", "specifies what DNN model to use", OFFSET(model_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "model_type" },
{ "srcnn", "Super-Resolution Convolutional Neural Network model (scale factor should be specified for custom SRCNN model)", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "model_type" },
{ "espcn", "Efficient Sub-Pixel Convolutional Neural Network model", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "model_type" },
{ "dnn_backend", "DNN backend used for model execution", OFFSET(backend_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "backend" },
{ "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" },
#if (CONFIG_LIBTENSORFLOW == 1)
{ "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" },
#endif
{"scale_factor", "scale factor for SRCNN model", OFFSET(scale_factor), AV_OPT_TYPE_INT, { .i64 = 2 }, 2, 4, FLAGS},
{ "model_filename", "path to model file specifying network architecture and its parameters", OFFSET(model_filename), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
{ NULL }
};
AVFILTER_DEFINE_CLASS(sr);
static av_cold int init(AVFilterContext* context)
{
SRContext* sr_context = context->priv;
sr_context->dnn_module = ff_get_dnn_module(sr_context->backend_type);
if (!sr_context->dnn_module){
av_log(context, AV_LOG_ERROR, "could not create DNN module for requested backend\n");
return AVERROR(ENOMEM);
}
if (!sr_context->model_filename){
av_log(context, AV_LOG_VERBOSE, "model file for network was not specified, using default network for x2 upsampling\n");
sr_context->scale_factor = 2;
switch (sr_context->model_type){
case SRCNN:
sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_SRCNN);
break;
case ESPCN:
sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_ESPCN);
}
}
else{
sr_context->model = (sr_context->dnn_module->load_model)(sr_context->model_filename);
}
if (!sr_context->model){
av_log(context, AV_LOG_ERROR, "could not load DNN model\n");
return AVERROR(EIO);
}
return 0;
}
static int query_formats(AVFilterContext* context)
{
const enum AVPixelFormat pixel_formats[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P,
AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_GRAY8,
AV_PIX_FMT_NONE};
AVFilterFormats* formats_list;
formats_list = ff_make_format_list(pixel_formats);
if (!formats_list){
av_log(context, AV_LOG_ERROR, "could not create formats list\n");
return AVERROR(ENOMEM);
}
return ff_set_common_formats(context, formats_list);
}
static int config_props(AVFilterLink* inlink)
{
AVFilterContext* context = inlink->dst;
SRContext* sr_context = context->priv;
AVFilterLink* outlink = context->outputs[0];
DNNReturnType result;
int sws_src_h, sws_src_w, sws_dst_h, sws_dst_w;
switch (sr_context->model_type){
case SRCNN:
sr_context->input.width = inlink->w * sr_context->scale_factor;
sr_context->input.height = inlink->h * sr_context->scale_factor;
break;
case ESPCN:
sr_context->input.width = inlink->w;
sr_context->input.height = inlink->h;
}
sr_context->input.channels = 1;
result = (sr_context->model->set_input_output)(sr_context->model->model, &sr_context->input, &sr_context->output);
if (result != DNN_SUCCESS){
av_log(context, AV_LOG_ERROR, "could not set input and output for the model\n");
return AVERROR(EIO);
}
else{
outlink->h = sr_context->output.height;
outlink->w = sr_context->output.width;
switch (sr_context->model_type){
case SRCNN:
sr_context->sws_context = sws_getContext(inlink->w, inlink->h, inlink->format,
outlink->w, outlink->h, outlink->format, SWS_BICUBIC, NULL, NULL, NULL);
if (!sr_context->sws_context){
av_log(context, AV_LOG_ERROR, "could not create SwsContext\n");
return AVERROR(ENOMEM);
}
sr_context->sws_slice_h = inlink->h;
break;
case ESPCN:
if (inlink->format == AV_PIX_FMT_GRAY8){
sr_context->sws_context = NULL;
}
else{
sws_src_h = sr_context->input.height;
sws_src_w = sr_context->input.width;
sws_dst_h = sr_context->output.height;
sws_dst_w = sr_context->output.width;
switch (inlink->format){
case AV_PIX_FMT_YUV420P:
sws_src_h = (sws_src_h >> 1) + (sws_src_h % 2 != 0 ? 1 : 0);
sws_src_w = (sws_src_w >> 1) + (sws_src_w % 2 != 0 ? 1 : 0);
sws_dst_h = (sws_dst_h >> 1) + (sws_dst_h % 2 != 0 ? 1 : 0);
sws_dst_w = (sws_dst_w >> 1) + (sws_dst_w % 2 != 0 ? 1 : 0);
break;
case AV_PIX_FMT_YUV422P:
sws_src_w = (sws_src_w >> 1) + (sws_src_w % 2 != 0 ? 1 : 0);
sws_dst_w = (sws_dst_w >> 1) + (sws_dst_w % 2 != 0 ? 1 : 0);
break;
case AV_PIX_FMT_YUV444P:
break;
case AV_PIX_FMT_YUV410P:
sws_src_h = (sws_src_h >> 2) + (sws_src_h % 4 != 0 ? 1 : 0);
sws_src_w = (sws_src_w >> 2) + (sws_src_w % 4 != 0 ? 1 : 0);
sws_dst_h = (sws_dst_h >> 2) + (sws_dst_h % 4 != 0 ? 1 : 0);
sws_dst_w = (sws_dst_w >> 2) + (sws_dst_w % 4 != 0 ? 1 : 0);
break;
case AV_PIX_FMT_YUV411P:
sws_src_w = (sws_src_w >> 2) + (sws_src_w % 4 != 0 ? 1 : 0);
sws_dst_w = (sws_dst_w >> 2) + (sws_dst_w % 4 != 0 ? 1 : 0);
break;
default:
av_log(context, AV_LOG_ERROR, "could not create SwsContext for input pixel format");
return AVERROR(EIO);
}
sr_context->sws_context = sws_getContext(sws_src_w, sws_src_h, AV_PIX_FMT_GRAY8,
sws_dst_w, sws_dst_h, AV_PIX_FMT_GRAY8, SWS_BICUBIC, NULL, NULL, NULL);
if (!sr_context->sws_context){
av_log(context, AV_LOG_ERROR, "could not create SwsContext\n");
return AVERROR(ENOMEM);
}
sr_context->sws_slice_h = sws_src_h;
}
}
return 0;
}
}
typedef struct ThreadData{
uint8_t* data;
int data_linesize, height, width;
} ThreadData;
static int uint8_to_float(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
{
SRContext* sr_context = context->priv;
const ThreadData* td = arg;
const int slice_start = (td->height * jobnr ) / nb_jobs;
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs;
const uint8_t* src = td->data + slice_start * td->data_linesize;
float* dst = sr_context->input.data + slice_start * td->width;
int y, x;
for (y = slice_start; y < slice_end; ++y){
for (x = 0; x < td->width; ++x){
dst[x] = (float)src[x] / 255.0f;
}
src += td->data_linesize;
dst += td->width;
}
return 0;
}
static int float_to_uint8(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
{
SRContext* sr_context = context->priv;
const ThreadData* td = arg;
const int slice_start = (td->height * jobnr ) / nb_jobs;
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs;
const float* src = sr_context->output.data + slice_start * td->width;
uint8_t* dst = td->data + slice_start * td->data_linesize;
int y, x;
for (y = slice_start; y < slice_end; ++y){
for (x = 0; x < td->width; ++x){
dst[x] = (uint8_t)(255.0f * FFMIN(src[x], 1.0f));
}
src += td->width;
dst += td->data_linesize;
}
return 0;
}
static int filter_frame(AVFilterLink* inlink, AVFrame* in)
{
AVFilterContext* context = inlink->dst;
SRContext* sr_context = context->priv;
AVFilterLink* outlink = context->outputs[0];
AVFrame* out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
ThreadData td;
int nb_threads;
DNNReturnType dnn_result;
if (!out){
av_log(context, AV_LOG_ERROR, "could not allocate memory for output frame\n");
av_frame_free(&in);
return AVERROR(ENOMEM);
}
av_frame_copy_props(out, in);
out->height = sr_context->output.height;
out->width = sr_context->output.width;
switch (sr_context->model_type){
case SRCNN:
sws_scale(sr_context->sws_context, in->data, in->linesize,
0, sr_context->sws_slice_h, out->data, out->linesize);
td.data = out->data[0];
td.data_linesize = out->linesize[0];
td.height = out->height;
td.width = out->width;
break;
case ESPCN:
if (sr_context->sws_context){
sws_scale(sr_context->sws_context, in->data + 1, in->linesize + 1,
0, sr_context->sws_slice_h, out->data + 1, out->linesize + 1);
sws_scale(sr_context->sws_context, in->data + 2, in->linesize + 2,
0, sr_context->sws_slice_h, out->data + 2, out->linesize + 2);
}
td.data = in->data[0];
td.data_linesize = in->linesize[0];
td.height = in->height;
td.width = in->width;
}
nb_threads = ff_filter_get_nb_threads(context);
context->internal->execute(context, uint8_to_float, &td, NULL, FFMIN(td.height, nb_threads));
av_frame_free(&in);
dnn_result = (sr_context->dnn_module->execute_model)(sr_context->model);
if (dnn_result != DNN_SUCCESS){
av_log(context, AV_LOG_ERROR, "failed to execute loaded model\n");
return AVERROR(EIO);
}
td.data = out->data[0];
td.data_linesize = out->linesize[0];
td.height = out->height;
td.width = out->width;
context->internal->execute(context, float_to_uint8, &td, NULL, FFMIN(td.height, nb_threads));
return ff_filter_frame(outlink, out);
}
static av_cold void uninit(AVFilterContext* context)
{
SRContext* sr_context = context->priv;
if (sr_context->dnn_module){
(sr_context->dnn_module->free_model)(&sr_context->model);
av_freep(&sr_context->dnn_module);
}
if (sr_context->sws_context){
sws_freeContext(sr_context->sws_context);
}
}
static const AVFilterPad sr_inputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
.config_props = config_props,
.filter_frame = filter_frame,
},
{ NULL }
};
static const AVFilterPad sr_outputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
},
{ NULL }
};
AVFilter ff_vf_sr = {
.name = "sr",
.description = NULL_IF_CONFIG_SMALL("Apply DNN-based image super resolution to the input."),
.priv_size = sizeof(SRContext),
.init = init,
.uninit = uninit,
.query_formats = query_formats,
.inputs = sr_inputs,
.outputs = sr_outputs,
.priv_class = &sr_class,
.flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
};

250
libavfilter/vf_srcnn.c
View File

@ -1,250 +0,0 @@
/*
* Copyright (c) 2018 Sergey Lavrushkin
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
/**
* @file
* Filter implementing image super-resolution using deep convolutional networks.
* https://arxiv.org/abs/1501.00092
*/
#include "avfilter.h"
#include "formats.h"
#include "internal.h"
#include "libavutil/opt.h"
#include "libavformat/avio.h"
#include "dnn_interface.h"
typedef struct SRCNNContext {
const AVClass *class;
char* model_filename;
float* input_output_buf;
DNNBackendType backend_type;
DNNModule* dnn_module;
DNNModel* model;
DNNData input_output;
} SRCNNContext;
#define OFFSET(x) offsetof(SRCNNContext, x)
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM
static const AVOption srcnn_options[] = {
{ "dnn_backend", "DNN backend used for model execution", OFFSET(backend_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "backend" },
{ "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" },
#if (CONFIG_LIBTENSORFLOW == 1)
{ "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" },
#endif
{ "model_filename", "path to model file specifying network architecture and its parameters", OFFSET(model_filename), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
{ NULL }
};
AVFILTER_DEFINE_CLASS(srcnn);
static av_cold int init(AVFilterContext* context)
{
SRCNNContext* srcnn_context = context->priv;
srcnn_context->dnn_module = ff_get_dnn_module(srcnn_context->backend_type);
if (!srcnn_context->dnn_module){
av_log(context, AV_LOG_ERROR, "could not create DNN module for requested backend\n");
return AVERROR(ENOMEM);
}
if (!srcnn_context->model_filename){
av_log(context, AV_LOG_VERBOSE, "model file for network was not specified, using default network for x2 upsampling\n");
srcnn_context->model = (srcnn_context->dnn_module->load_default_model)(DNN_SRCNN);
}
else{
srcnn_context->model = (srcnn_context->dnn_module->load_model)(srcnn_context->model_filename);
}
if (!srcnn_context->model){
av_log(context, AV_LOG_ERROR, "could not load DNN model\n");
return AVERROR(EIO);
}
return 0;
}
static int query_formats(AVFilterContext* context)
{
const enum AVPixelFormat pixel_formats[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P,
AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_GRAY8,
AV_PIX_FMT_NONE};
AVFilterFormats* formats_list;
formats_list = ff_make_format_list(pixel_formats);
if (!formats_list){
av_log(context, AV_LOG_ERROR, "could not create formats list\n");
return AVERROR(ENOMEM);
}
return ff_set_common_formats(context, formats_list);
}
static int config_props(AVFilterLink* inlink)
{
AVFilterContext* context = inlink->dst;
SRCNNContext* srcnn_context = context->priv;
DNNReturnType result;
srcnn_context->input_output_buf = av_malloc(inlink->h * inlink->w * sizeof(float));
if (!srcnn_context->input_output_buf){
av_log(context, AV_LOG_ERROR, "could not allocate memory for input/output buffer\n");
return AVERROR(ENOMEM);
}
srcnn_context->input_output.data = srcnn_context->input_output_buf;
srcnn_context->input_output.width = inlink->w;
srcnn_context->input_output.height = inlink->h;
srcnn_context->input_output.channels = 1;
result = (srcnn_context->model->set_input_output)(srcnn_context->model->model, &srcnn_context->input_output, &srcnn_context->input_output);
if (result != DNN_SUCCESS){
av_log(context, AV_LOG_ERROR, "could not set input and output for the model\n");
return AVERROR(EIO);
}
else{
return 0;
}
}
typedef struct ThreadData{
uint8_t* out;
int out_linesize, height, width;
} ThreadData;
static int uint8_to_float(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
{
SRCNNContext* srcnn_context = context->priv;
const ThreadData* td = arg;
const int slice_start = (td->height * jobnr ) / nb_jobs;
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs;
const uint8_t* src = td->out + slice_start * td->out_linesize;
float* dst = srcnn_context->input_output_buf + slice_start * td->width;
int y, x;
for (y = slice_start; y < slice_end; ++y){
for (x = 0; x < td->width; ++x){
dst[x] = (float)src[x] / 255.0f;
}
src += td->out_linesize;
dst += td->width;
}
return 0;
}
static int float_to_uint8(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
{
SRCNNContext* srcnn_context = context->priv;
const ThreadData* td = arg;
const int slice_start = (td->height * jobnr ) / nb_jobs;
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs;
const float* src = srcnn_context->input_output_buf + slice_start * td->width;
uint8_t* dst = td->out + slice_start * td->out_linesize;
int y, x;
for (y = slice_start; y < slice_end; ++y){
for (x = 0; x < td->width; ++x){
dst[x] = (uint8_t)(255.0f * FFMIN(src[x], 1.0f));
}
src += td->width;
dst += td->out_linesize;
}
return 0;
}
static int filter_frame(AVFilterLink* inlink, AVFrame* in)
{
AVFilterContext* context = inlink->dst;
SRCNNContext* srcnn_context = context->priv;
AVFilterLink* outlink = context->outputs[0];
AVFrame* out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
ThreadData td;
int nb_threads;
DNNReturnType dnn_result;
if (!out){
av_log(context, AV_LOG_ERROR, "could not allocate memory for output frame\n");
av_frame_free(&in);
return AVERROR(ENOMEM);
}
av_frame_copy_props(out, in);
av_frame_copy(out, in);
av_frame_free(&in);
td.out = out->data[0];
td.out_linesize = out->linesize[0];
td.height = out->height;
td.width = out->width;
nb_threads = ff_filter_get_nb_threads(context);
context->internal->execute(context, uint8_to_float, &td, NULL, FFMIN(td.height, nb_threads));
dnn_result = (srcnn_context->dnn_module->execute_model)(srcnn_context->model);
if (dnn_result != DNN_SUCCESS){
av_log(context, AV_LOG_ERROR, "failed to execute loaded model\n");
return AVERROR(EIO);
}
context->internal->execute(context, float_to_uint8, &td, NULL, FFMIN(td.height, nb_threads));
return ff_filter_frame(outlink, out);
}
static av_cold void uninit(AVFilterContext* context)
{
SRCNNContext* srcnn_context = context->priv;
if (srcnn_context->dnn_module){
(srcnn_context->dnn_module->free_model)(&srcnn_context->model);
av_freep(&srcnn_context->dnn_module);
}
av_freep(&srcnn_context->input_output_buf);
}
static const AVFilterPad srcnn_inputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
.config_props = config_props,
.filter_frame = filter_frame,
},
{ NULL }
};
static const AVFilterPad srcnn_outputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
},
{ NULL }
};
AVFilter ff_vf_srcnn = {
.name = "srcnn",
.description = NULL_IF_CONFIG_SMALL("Apply super resolution convolutional neural network to the input. Use bicubic upsamping with corresponding scaling factor before."),
.priv_size = sizeof(SRCNNContext),
.init = init,
.uninit = uninit,
.query_formats = query_formats,
.inputs = srcnn_inputs,
.outputs = srcnn_outputs,
.priv_class = &srcnn_class,
.flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
};