Adds ESPCN super resolution filter merged with SRCNN filter.

Signed-off-by: Pedro Arthur <bygrandao@gmail.com>

configure

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

libavfilter/Makefile

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

libavfilter/allfilters.c

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

libavfilter/dnn_backend_native.c

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

libavfilter/dnn_backend_tf.c

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

libavfilter/dnn_interface.h

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

libavfilter/dnn_srcnn.h

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

libavfilter/vf_sr.c

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

libavfilter/vf_srcnn.c

@@ -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,
-};
-