falkTX
/
FFmpeg
mirror of https://github.com/falkTX/FFmpeg.git
1
0
Fork 0
Code Issues 0 Releases 338 Wiki Activity
Browse Source

Adds ESPCN super resolution filter merged with SRCNN filter. 

Signed-off-by: Pedro Arthur <bygrandao@gmail.com>
tags/n4.1
Sergey Lavrushkin Pedro Arthur 7 years ago
parent
d24c9e55f6
commit
575b718990
10 changed files with 13249 additions and 379 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +4
    -4
      configure
  2. +1
    -1
      libavfilter/Makefile
  3. +1
    -1
      libavfilter/allfilters.c
  4. +198
    -85
      libavfilter/dnn_backend_native.c
  5. +45
    -29
      libavfilter/dnn_backend_tf.c
  6. +12637
    -0
      libavfilter/dnn_espcn.h
  7. +2
    -2
      libavfilter/dnn_interface.h
  8. +7
    -7
      libavfilter/dnn_srcnn.h
  9. +354
    -0
      libavfilter/vf_sr.c
  10. +0
    -250
      libavfilter/vf_srcnn.c

+ 4
- 4
configure 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"


+ 1
- 1
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


+ 1
- 1
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;


+ 198
- 85
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){
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){
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;
}
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;
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;
}
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;
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;
}
network->layers[0].type = INPUT;
network->layers[0].params = av_malloc(sizeof(InputParams));
if (!network->layers[0].params){
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;
}

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){
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;
model->set_input_output = &set_input_output_native;

return model;
default:
av_freep(&network);
av_freep(&model);
return NULL;
}
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);


+ 45
- 29
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
- 0
libavfilter/dnn_espcn.h
File diff suppressed because it is too large
View File


+ 2
- 2
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.


+ 7
- 7
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
- 0
libavfilter/vf_sr.c 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,
};


+ 0
- 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,
};


Write Preview
Loading…
Cancel
Save