apps/ffmpeg-rockchip
1
0
Fork 0
mirror of https://github.com/nyanmisaka/ffmpeg-rockchip.git synced 2025-11-01 04:53:04 +08:00
Code Issues Packages Projects Releases Wiki 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
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-libsrt enable Haivision SRT protocol via libsrt [no]
--enable-libssh enable SFTP protocol via libssh [no] --enable-libssh enable SFTP protocol via libssh [no]
--enable-libtensorflow enable TensorFlow as a DNN module backend --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-libtesseract enable Tesseract, needed for ocr filter [no]
--enable-libtheora enable Theora encoding via libtheora [no] --enable-libtheora enable Theora encoding via libtheora [no]
--enable-libtls enable LibreSSL (via libtls), needed for https support --enable-libtls enable LibreSSL (via libtls), needed for https support
@@ -3402,8 +3402,8 @@ spectrumsynth_filter_deps="avcodec"
spectrumsynth_filter_select="fft" spectrumsynth_filter_select="fft"
spp_filter_deps="gpl avcodec" spp_filter_deps="gpl avcodec"
spp_filter_select="fft idctdsp fdctdsp me_cmp pixblockdsp" spp_filter_select="fft idctdsp fdctdsp me_cmp pixblockdsp"
srcnn_filter_deps="avformat" sr_filter_deps="avformat swscale"
srcnn_filter_select="dnn" sr_filter_select="dnn"
stereo3d_filter_deps="gpl" stereo3d_filter_deps="gpl"
subtitles_filter_deps="avformat avcodec libass" subtitles_filter_deps="avformat avcodec libass"
super2xsai_filter_deps="gpl" super2xsai_filter_deps="gpl"
@@ -6823,7 +6823,7 @@ enabled signature_filter && prepend avfilter_deps "avcodec avformat"
enabled smartblur_filter && prepend avfilter_deps "swscale" enabled smartblur_filter && prepend avfilter_deps "swscale"
enabled spectrumsynth_filter && prepend avfilter_deps "avcodec" enabled spectrumsynth_filter && prepend avfilter_deps "avcodec"
enabled spp_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 subtitles_filter && prepend avfilter_deps "avformat avcodec"
enabled uspp_filter && prepend avfilter_deps "avcodec" enabled uspp_filter && prepend avfilter_deps "avcodec"
enabled zoompan_filter && prepend avfilter_deps "swscale" 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_SOBEL_FILTER) += vf_convolution.o
OBJS-$(CONFIG_SPLIT_FILTER) += split.o OBJS-$(CONFIG_SPLIT_FILTER) += split.o
OBJS-$(CONFIG_SPP_FILTER) += vf_spp.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_SSIM_FILTER) += vf_ssim.o framesync.o
OBJS-$(CONFIG_STEREO3D_FILTER) += vf_stereo3d.o OBJS-$(CONFIG_STEREO3D_FILTER) += vf_stereo3d.o
OBJS-$(CONFIG_STREAMSELECT_FILTER) += f_streamselect.o framesync.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_sobel;
extern AVFilter ff_vf_split; extern AVFilter ff_vf_split;
extern AVFilter ff_vf_spp; extern AVFilter ff_vf_spp;
extern AVFilter ff_vf_srcnn; extern AVFilter ff_vf_sr;
extern AVFilter ff_vf_ssim; extern AVFilter ff_vf_ssim;
extern AVFilter ff_vf_stereo3d; extern AVFilter ff_vf_stereo3d;
extern AVFilter ff_vf_streamselect; extern AVFilter ff_vf_streamselect;

191
libavfilter/dnn_backend_native.c
View File

@@ -25,9 +25,12 @@
#include "dnn_backend_native.h" #include "dnn_backend_native.h"
#include "dnn_srcnn.h" #include "dnn_srcnn.h"
#include "dnn_espcn.h"
#include "libavformat/avio.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{ typedef struct Layer{
LayerType type; LayerType type;
@@ -37,6 +40,7 @@ typedef struct Layer{
typedef struct ConvolutionalParams{ typedef struct ConvolutionalParams{
int32_t input_num, output_num, kernel_size; int32_t input_num, output_num, kernel_size;
ActivationFunc activation;
float* kernel; float* kernel;
float* biases; float* biases;
} ConvolutionalParams; } ConvolutionalParams;
@@ -45,17 +49,22 @@ typedef struct InputParams{
int height, width, channels; int height, width, channels;
} InputParams; } InputParams;
typedef struct DepthToSpaceParams{
int block_size;
} DepthToSpaceParams;
// Represents simple feed-forward convolutional network. // Represents simple feed-forward convolutional network.
typedef struct ConvolutionalNetwork{ typedef struct ConvolutionalNetwork{
Layer* layers; Layer* layers;
int32_t layers_num; int32_t layers_num;
} ConvolutionalNetwork; } 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; ConvolutionalNetwork* network = (ConvolutionalNetwork*)model;
InputParams* input_params; InputParams* input_params;
ConvolutionalParams* conv_params; ConvolutionalParams* conv_params;
DepthToSpaceParams* depth_to_space_params;
int cur_width, cur_height, cur_channels; int cur_width, cur_height, cur_channels;
int32_t layer; int32_t layer;
@@ -63,11 +72,17 @@ static DNNReturnType set_input_output_native(void* model, const DNNData* input,
return DNN_ERROR; return DNN_ERROR;
} }
else{ else{
network->layers[0].output = input->data;
input_params = (InputParams*)network->layers[0].params; input_params = (InputParams*)network->layers[0].params;
input_params->width = cur_width = input->width; input_params->width = cur_width = input->width;
input_params->height = cur_height = input->height; input_params->height = cur_height = input->height;
input_params->channels = cur_channels = input->channels; 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){ for (layer = 1; layer < network->layers_num; ++layer){
@@ -78,8 +93,20 @@ static DNNReturnType set_input_output_native(void* model, const DNNData* input,
return DNN_ERROR; return DNN_ERROR;
} }
cur_channels = conv_params->output_num; cur_channels = conv_params->output_num;
if (layer < network->layers_num - 1){ break;
if (!network->layers[layer].output){ 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); av_freep(&network->layers[layer].output);
} }
network->layers[layer].output = av_malloc(cur_height * cur_width * cur_channels * sizeof(float)); network->layers[layer].output = av_malloc(cur_height * cur_width * cur_channels * sizeof(float));
@@ -87,23 +114,19 @@ static DNNReturnType set_input_output_native(void* model, const DNNData* input,
return DNN_ERROR; return DNN_ERROR;
} }
} }
else{
network->layers[layer].output = output->data; output->data = network->layers[network->layers_num - 1].output;
if (output->width != cur_width || output->height != cur_height || output->channels != cur_channels){ output->height = cur_height;
return DNN_ERROR; output->width = cur_width;
} output->channels = cur_channels;
}
break;
default:
return DNN_ERROR;
}
}
return DNN_SUCCESS; return DNN_SUCCESS;
} }
// Loads model and its parameters that are stored in a binary file with following structure: // 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* ff_dnn_load_model_native(const char* model_filename)
{ {
DNNModel* model = NULL; DNNModel* model = NULL;
@@ -111,7 +134,9 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
AVIOContext* model_file_context; AVIOContext* model_file_context;
int file_size, dnn_size, kernel_size, i; int file_size, dnn_size, kernel_size, i;
int32_t layer; int32_t layer;
LayerType layer_type;
ConvolutionalParams* conv_params; ConvolutionalParams* conv_params;
DepthToSpaceParams* depth_to_space_params;
model = av_malloc(sizeof(DNNModel)); model = av_malloc(sizeof(DNNModel));
if (!model){ if (!model){
@@ -156,18 +181,23 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
} }
for (layer = 1; layer < network->layers_num; ++layer){ for (layer = 1; layer < network->layers_num; ++layer){
layer_type = (int32_t)avio_rl32(model_file_context);
dnn_size += 4;
switch (layer_type){
case CONV:
conv_params = av_malloc(sizeof(ConvolutionalParams)); conv_params = av_malloc(sizeof(ConvolutionalParams));
if (!conv_params){ if (!conv_params){
avio_closep(&model_file_context); avio_closep(&model_file_context);
ff_dnn_free_model_native(&model); ff_dnn_free_model_native(&model);
return NULL; 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->input_num = (int32_t)avio_rl32(model_file_context);
conv_params->output_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); conv_params->kernel_size = (int32_t)avio_rl32(model_file_context);
kernel_size = conv_params->input_num * conv_params->output_num * kernel_size = conv_params->input_num * conv_params->output_num *
conv_params->kernel_size * conv_params->kernel_size; conv_params->kernel_size * conv_params->kernel_size;
dnn_size += 12 + (kernel_size + conv_params->output_num << 2); dnn_size += 16 + (kernel_size + conv_params->output_num << 2);
if (dnn_size > file_size || conv_params->input_num <= 0 || if (dnn_size > file_size || conv_params->input_num <= 0 ||
conv_params->output_num <= 0 || conv_params->kernel_size <= 0){ conv_params->output_num <= 0 || conv_params->kernel_size <= 0){
avio_closep(&model_file_context); avio_closep(&model_file_context);
@@ -189,6 +219,24 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
} }
network->layers[layer].type = CONV; network->layers[layer].type = CONV;
network->layers[layer].params = conv_params; 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;
}
} }
avio_closep(&model_file_context); avio_closep(&model_file_context);
@@ -203,7 +251,8 @@ DNNModel* ff_dnn_load_model_native(const char* model_filename)
return model; 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; ConvolutionalParams* conv_params;
int kernel_size; int kernel_size;
@@ -212,6 +261,7 @@ static int set_up_conv_layer(Layer* layer, const float* kernel, const float* bia
if (!conv_params){ if (!conv_params){
return DNN_ERROR; return DNN_ERROR;
} }
conv_params->activation = activation;
conv_params->input_num = input_num; conv_params->input_num = input_num;
conv_params->output_num = output_num; conv_params->output_num = output_num;
conv_params->kernel_size = size; conv_params->kernel_size = size;
@@ -236,6 +286,7 @@ DNNModel* ff_dnn_load_default_model_native(DNNDefaultModel model_type)
{ {
DNNModel* model = NULL; DNNModel* model = NULL;
ConvolutionalNetwork* network = NULL; ConvolutionalNetwork* network = NULL;
DepthToSpaceParams* depth_to_space_params;
int32_t layer; int32_t layer;
model = av_malloc(sizeof(DNNModel)); model = av_malloc(sizeof(DNNModel));
@@ -253,6 +304,15 @@ DNNModel* ff_dnn_load_default_model_native(DNNDefaultModel model_type)
switch (model_type){ switch (model_type){
case DNN_SRCNN: case DNN_SRCNN:
network->layers_num = 4; 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)); network->layers = av_malloc(network->layers_num * sizeof(Layer));
if (!network->layers){ if (!network->layers){
@@ -272,26 +332,40 @@ DNNModel* ff_dnn_load_default_model_native(DNNDefaultModel model_type)
return NULL; return NULL;
} }
if (set_up_conv_layer(network->layers + 1, conv1_kernel, conv1_biases, 1, 64, 9) != DNN_SUCCESS || switch (model_type){
set_up_conv_layer(network->layers + 2, conv2_kernel, conv2_biases, 64, 32, 1) != DNN_SUCCESS || case DNN_SRCNN:
set_up_conv_layer(network->layers + 3, conv3_kernel, conv3_biases, 32, 1, 5) != DNN_SUCCESS){ 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); ff_dnn_free_model_native(&model);
return NULL; 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; model->set_input_output = &set_input_output_native;
return model; return model;
default:
av_freep(&network);
av_freep(&model);
return NULL;
}
} }
#define CLAMP_TO_EDGE(x, w) ((x) < 0 ? 0 : ((x) >= (w) ? (w - 1) : (x))) #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 y, x, n_filter, ch, kernel_y, kernel_x;
int radius = conv_params->kernel_size >> 1; int radius = conv_params->kernel_size >> 1;
@@ -313,19 +387,53 @@ static void convolve(const float* input, float* output, const ConvolutionalParam
} }
} }
} }
switch (conv_params->activation){
case RELU:
output[n_filter] = FFMAX(output[n_filter], 0.0); 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; 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) DNNReturnType ff_dnn_execute_model_native(const DNNModel* model)
{ {
ConvolutionalNetwork* network = (ConvolutionalNetwork*)model->model; ConvolutionalNetwork* network = (ConvolutionalNetwork*)model->model;
InputParams* input_params; int cur_width, cur_height, cur_channels;
int cur_width, cur_height;
int32_t layer; 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){ if (network->layers_num <= 0 || network->layers[0].type != INPUT || !network->layers[0].output){
return DNN_ERROR; return DNN_ERROR;
@@ -334,6 +442,7 @@ DNNReturnType ff_dnn_execute_model_native(const DNNModel* model)
input_params = (InputParams*)network->layers[0].params; input_params = (InputParams*)network->layers[0].params;
cur_width = input_params->width; cur_width = input_params->width;
cur_height = input_params->height; cur_height = input_params->height;
cur_channels = input_params->channels;
} }
for (layer = 1; layer < network->layers_num; ++layer){ 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){ switch (network->layers[layer].type){
case CONV: 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; break;
case INPUT: case INPUT:
return DNN_ERROR; return DNN_ERROR;
@@ -362,19 +481,13 @@ void ff_dnn_free_model_native(DNNModel** model)
{ {
network = (ConvolutionalNetwork*)(*model)->model; network = (ConvolutionalNetwork*)(*model)->model;
for (layer = 0; layer < network->layers_num; ++layer){ 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; conv_params = (ConvolutionalParams*)network->layers[layer].params;
av_freep(&conv_params->kernel); av_freep(&conv_params->kernel);
av_freep(&conv_params->biases); 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(network);
av_freep(model); av_freep(model);

70
libavfilter/dnn_backend_tf.c
View File

@@ -25,6 +25,7 @@
#include "dnn_backend_tf.h" #include "dnn_backend_tf.h"
#include "dnn_srcnn.h" #include "dnn_srcnn.h"
#include "dnn_espcn.h"
#include "libavformat/avio.h" #include "libavformat/avio.h"
#include <tensorflow/c/c_api.h> #include <tensorflow/c/c_api.h>
@@ -35,9 +36,7 @@ typedef struct TFModel{
TF_Status* status; TF_Status* status;
TF_Output input, output; TF_Output input, output;
TF_Tensor* input_tensor; TF_Tensor* input_tensor;
TF_Tensor* output_tensor; DNNData* output_data;
const DNNData* input_data;
const DNNData* output_data;
} TFModel; } TFModel;
static void free_buffer(void* data, size_t length) 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; 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; TFModel* tf_model = (TFModel*)model;
int64_t input_dims[] = {1, input->height, input->width, input->channels}; 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; TF_SessionOptions* sess_opts;
const TF_Operation* init_op = TF_GraphOperationByName(tf_model->graph, "init"); const TF_Operation* init_op = TF_GraphOperationByName(tf_model->graph, "init");
TF_Tensor* output_tensor;
// Input operation should be named 'x' // Input operation should be named 'x'
tf_model->input.oper = TF_GraphOperationByName(tf_model->graph, "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){ if (!tf_model->input_tensor){
return DNN_ERROR; return DNN_ERROR;
} }
input->data = (float*)TF_TensorData(tf_model->input_tensor);
// Output operation should be named 'y' // Output operation should be named 'y'
tf_model->output.oper = TF_GraphOperationByName(tf_model->graph, "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; return DNN_ERROR;
} }
tf_model->output.index = 0; 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){ if (tf_model->session){
TF_CloseSession(tf_model->session, tf_model->status); 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; return DNN_SUCCESS;
} }
@@ -166,7 +175,7 @@ DNNModel* ff_dnn_load_model_tf(const char* model_filename)
} }
tf_model->session = NULL; tf_model->session = NULL;
tf_model->input_tensor = NULL; tf_model->input_tensor = NULL;
tf_model->output_tensor = NULL; tf_model->output_data = NULL;
graph_def = read_graph(model_filename); graph_def = read_graph(model_filename);
if (!graph_def){ 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->length = srcnn_tf_size;
graph_def->data_deallocator = free_buffer; graph_def->data_deallocator = free_buffer;
break; 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: default:
TF_DeleteBuffer(graph_def); TF_DeleteBuffer(graph_def);
return NULL; return NULL;
@@ -234,7 +254,7 @@ DNNModel* ff_dnn_load_default_model_tf(DNNDefaultModel model_type)
} }
tf_model->session = NULL; tf_model->session = NULL;
tf_model->input_tensor = NULL; tf_model->input_tensor = NULL;
tf_model->output_tensor = NULL; tf_model->output_data = NULL;
tf_model->graph = TF_NewGraph(); tf_model->graph = TF_NewGraph();
tf_model->status = TF_NewStatus(); 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) DNNReturnType ff_dnn_execute_model_tf(const DNNModel* model)
{ {
TFModel* tf_model = (TFModel*)model->model; TFModel* tf_model = (TFModel*)model->model;
TF_Tensor* output_tensor;
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_SessionRun(tf_model->session, NULL, TF_SessionRun(tf_model->session, NULL,
&tf_model->input, &tf_model->input_tensor, 1, &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); NULL, 0, NULL, tf_model->status);
if (TF_GetCode(tf_model->status) != TF_OK){ if (TF_GetCode(tf_model->status) != TF_OK){
return DNN_ERROR; return DNN_ERROR;
} }
else{ else{
memcpy(tf_model->output_data->data, TF_TensorData(tf_model->output_tensor), memcpy(tf_model->output_data->data, TF_TensorData(output_tensor),
tf_model->output_data->height * tf_model->output_data->width * tf_model->output_data->height * tf_model->output_data->width *
tf_model->output_data->channels * sizeof(float)); tf_model->output_data->channels * sizeof(float));
TF_DeleteTensor(output_tensor);
return DNN_SUCCESS; return DNN_SUCCESS;
} }
@@ -300,9 +318,7 @@ void ff_dnn_free_model_tf(DNNModel** model)
if (tf_model->input_tensor){ if (tf_model->input_tensor){
TF_DeleteTensor(tf_model->input_tensor); TF_DeleteTensor(tf_model->input_tensor);
} }
if (tf_model->output_tensor){ av_freep(&tf_model->output_data->data);
TF_DeleteTensor(tf_model->output_tensor);
}
av_freep(&tf_model); av_freep(&tf_model);
av_freep(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_NATIVE, DNN_TF} DNNBackendType;
typedef enum {DNN_SRCNN} DNNDefaultModel; typedef enum {DNN_SRCNN, DNN_ESPCN} DNNDefaultModel;
typedef struct DNNData{ typedef struct DNNData{
float* data; float* data;
@@ -42,7 +42,7 @@ typedef struct DNNModel{
void* model; void* model;
// Sets model input and output, while allocating additional memory for intermediate calculations. // Sets model input and output, while allocating additional memory for intermediate calculations.
// Should be called at least once before model execution. // 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; } DNNModel;
// Stores pointers to functions for loading, executing, freeing DNN models for one of the backends. // 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 * @file
* Default cnn weights for x2 upsampling with srcnn filter. * Default cnn weights for x2 upsampling with srcnn model.
*/ */
#ifndef AVFILTER_DNN_SRCNN_H #ifndef AVFILTER_DNN_SRCNN_H
#define 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.08866338f, 0.055409566f, 0.037196506f, -0.11961404f,
-0.12341991f, 0.29963422f, -0.0911817f, -0.00013613555f, -0.12341991f, 0.29963422f, -0.0911817f, -0.00013613555f,
-0.049023595f, 0.038421184f, -0.077267796f, 0.027273094f, -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 -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.016606892f, -0.011107335f, -0.0048309686f, -0.04867378f,
-0.030040957f, -0.07297248f, -0.019458665f, -0.009738028f, -0.030040957f, -0.07297248f, -0.019458665f, -0.009738028f,
0.6951231f, -0.07369442f, -0.01354204f, 0.010336088f, 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 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.24004751f, 0.1037138f, 0.11173403f, 0.04352092f,
-0.23728481f, 0.12153747f, -0.23676059f, -0.28548065f, -0.23728481f, 0.12153747f, -0.23676059f, -0.28548065f,
-0.612738f, -0.12218937f, -0.06005159f, 0.1850652f, -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 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.12326373f, 0.13270757f, 0.07082674f, 0.051456157f,
0.058445618f, 0.13153197f, 0.0809729f, 0.10153213f, 0.058445618f, 0.13153197f, 0.0809729f, 0.10153213f,
0.055915363f, 0.05228166f, -0.11212896f, 0.07462141f, 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 -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.01733648f, 0.01492609f, 0.019393086f, -0.004445322f,
0.026939709f, 0.00038831023f, 0.004221528f, 0.0050745453f, 0.026939709f, 0.00038831023f, 0.004221528f, 0.0050745453f,
0.0129861f, 0.008007169f, 0.008950762f, 0.005279691f, 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 0.012931146f, 0.0046948805f, 0.013098622f, -0.015422701f
}; };
static const float conv3_biases[] = { static const float srcnn_conv3_biases[] = {
0.05037664f 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,
};