Signed-off-by: Pedro Arthur <bygrandao@gmail.com>tags/n4.1
| @@ -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" | |||||
| srcnn_filter_select="dnn" | |||||
| sr_filter_deps="avformat swscale" | |||||
| 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" | ||||
| @@ -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 | ||||
| @@ -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; | ||||
| @@ -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,32 +93,40 @@ 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){ | |||||
| 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; | break; | ||||
| default: | default: | ||||
| return DNN_ERROR; | 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; | 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,39 +181,62 @@ 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){ | ||||
| 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); | avio_closep(&model_file_context); | ||||
| ff_dnn_free_model_native(&model); | ff_dnn_free_model_native(&model); | ||||
| return NULL; | 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); | 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,45 +304,68 @@ 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)); | |||||
| 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); | ff_dnn_free_model_native(&model); | ||||
| return NULL; | 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); | ff_dnn_free_model_native(&model); | ||||
| return NULL; | 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))) | #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 | |||||
| } | } | ||||
| } | } | ||||
| } | } | ||||
| 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; | 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; | |||||
| int cur_width, cur_height, cur_channels; | |||||
| 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); | ||||
| @@ -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; | |||||
| const DNNData* input_data; | |||||
| const DNNData* output_data; | |||||
| 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; | ||||
| 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_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), | |||||
| 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; | 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){ | |||||
| TF_DeleteTensor(tf_model->output_tensor); | |||||
| } | |||||
| av_freep(&tf_model->output_data->data); | |||||
| av_freep(&tf_model); | av_freep(&tf_model); | ||||
| av_freep(model); | av_freep(model); | ||||
| } | } | ||||
| @@ -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. | ||||
| @@ -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 | ||||
| }; | }; | ||||
| @@ -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, | |||||
| }; | |||||
| @@ -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, | |||||
| }; | |||||