libavfilter: Adds on the fly generation of default DNN models for tensorflow backend...
[ffmpeg.git] / libavfilter / vf_sr.c
1 /*
2  * Copyright (c) 2018 Sergey Lavrushkin
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 /**
22  * @file
23  * Filter implementing image super-resolution using deep convolutional networks.
24  * https://arxiv.org/abs/1501.00092
25  * https://arxiv.org/abs/1609.05158
26  */
27
28 #include "avfilter.h"
29 #include "formats.h"
30 #include "internal.h"
31 #include "libavutil/opt.h"
32 #include "libavformat/avio.h"
33 #include "libswscale/swscale.h"
34 #include "dnn_interface.h"
35
36 typedef enum {SRCNN, ESPCN} SRModel;
37
38 typedef struct SRContext {
39     const AVClass *class;
40
41     SRModel model_type;
42     char* model_filename;
43     DNNBackendType backend_type;
44     DNNModule* dnn_module;
45     DNNModel* model;
46     DNNData input, output;
47     int scale_factor;
48     struct SwsContext* sws_context;
49     int sws_slice_h;
50 } SRContext;
51
52 #define OFFSET(x) offsetof(SRContext, x)
53 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM
54 static const AVOption sr_options[] = {
55     { "model", "specifies what DNN model to use", OFFSET(model_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "model_type" },
56     { "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" },
57     { "espcn", "Efficient Sub-Pixel Convolutional Neural Network model", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "model_type" },
58     { "dnn_backend", "DNN backend used for model execution", OFFSET(backend_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "backend" },
59     { "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" },
60 #if (CONFIG_LIBTENSORFLOW == 1)
61     { "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" },
62 #endif
63     {"scale_factor", "scale factor for SRCNN model", OFFSET(scale_factor), AV_OPT_TYPE_INT, { .i64 = 2 }, 2, 4, FLAGS},
64     { "model_filename", "path to model file specifying network architecture and its parameters", OFFSET(model_filename), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
65     { NULL }
66 };
67
68 AVFILTER_DEFINE_CLASS(sr);
69
70 static av_cold int init(AVFilterContext* context)
71 {
72     SRContext* sr_context = context->priv;
73
74     sr_context->dnn_module = ff_get_dnn_module(sr_context->backend_type);
75     if (!sr_context->dnn_module){
76         av_log(context, AV_LOG_ERROR, "could not create DNN module for requested backend\n");
77         return AVERROR(ENOMEM);
78     }
79     if (!sr_context->model_filename){
80         av_log(context, AV_LOG_VERBOSE, "model file for network was not specified, using default network for x2 upsampling\n");
81         sr_context->scale_factor = 2;
82         switch (sr_context->model_type){
83         case SRCNN:
84             sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_SRCNN);
85             break;
86         case ESPCN:
87             sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_ESPCN);
88         }
89     }
90     else{
91         sr_context->model = (sr_context->dnn_module->load_model)(sr_context->model_filename);
92     }
93     if (!sr_context->model){
94         av_log(context, AV_LOG_ERROR, "could not load DNN model\n");
95         return AVERROR(EIO);
96     }
97
98     return 0;
99 }
100
101 static int query_formats(AVFilterContext* context)
102 {
103     const enum AVPixelFormat pixel_formats[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P,
104                                                 AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_GRAY8,
105                                                 AV_PIX_FMT_NONE};
106     AVFilterFormats* formats_list;
107
108     formats_list = ff_make_format_list(pixel_formats);
109     if (!formats_list){
110         av_log(context, AV_LOG_ERROR, "could not create formats list\n");
111         return AVERROR(ENOMEM);
112     }
113     return ff_set_common_formats(context, formats_list);
114 }
115
116 static int config_props(AVFilterLink* inlink)
117 {
118     AVFilterContext* context = inlink->dst;
119     SRContext* sr_context = context->priv;
120     AVFilterLink* outlink = context->outputs[0];
121     DNNReturnType result;
122     int sws_src_h, sws_src_w, sws_dst_h, sws_dst_w;
123
124     switch (sr_context->model_type){
125     case SRCNN:
126         sr_context->input.width = inlink->w * sr_context->scale_factor;
127         sr_context->input.height = inlink->h * sr_context->scale_factor;
128         break;
129     case ESPCN:
130         sr_context->input.width = inlink->w;
131         sr_context->input.height = inlink->h;
132     }
133     sr_context->input.channels = 1;
134
135     result = (sr_context->model->set_input_output)(sr_context->model->model, &sr_context->input, &sr_context->output);
136     if (result != DNN_SUCCESS){
137         av_log(context, AV_LOG_ERROR, "could not set input and output for the model\n");
138         return AVERROR(EIO);
139     }
140     else{
141         outlink->h = sr_context->output.height;
142         outlink->w = sr_context->output.width;
143         switch (sr_context->model_type){
144         case SRCNN:
145             sr_context->sws_context = sws_getContext(inlink->w, inlink->h, inlink->format,
146                                                      outlink->w, outlink->h, outlink->format, SWS_BICUBIC, NULL, NULL, NULL);
147             if (!sr_context->sws_context){
148                 av_log(context, AV_LOG_ERROR, "could not create SwsContext\n");
149                 return AVERROR(ENOMEM);
150             }
151             sr_context->sws_slice_h = inlink->h;
152             break;
153         case ESPCN:
154             if (inlink->format == AV_PIX_FMT_GRAY8){
155                 sr_context->sws_context = NULL;
156             }
157             else{
158                 sws_src_h = sr_context->input.height;
159                 sws_src_w = sr_context->input.width;
160                 sws_dst_h = sr_context->output.height;
161                 sws_dst_w = sr_context->output.width;
162
163                 switch (inlink->format){
164                 case AV_PIX_FMT_YUV420P:
165                     sws_src_h = AV_CEIL_RSHIFT(sws_src_h, 1);
166                     sws_src_w = AV_CEIL_RSHIFT(sws_src_w, 1);
167                     sws_dst_h = AV_CEIL_RSHIFT(sws_dst_h, 1);
168                     sws_dst_w = AV_CEIL_RSHIFT(sws_dst_w, 1);
169                     break;
170                 case AV_PIX_FMT_YUV422P:
171                     sws_src_w = AV_CEIL_RSHIFT(sws_src_w, 1);
172                     sws_dst_w = AV_CEIL_RSHIFT(sws_dst_w, 1);
173                     break;
174                 case AV_PIX_FMT_YUV444P:
175                     break;
176                 case AV_PIX_FMT_YUV410P:
177                     sws_src_h = AV_CEIL_RSHIFT(sws_src_h, 2);
178                     sws_src_w = AV_CEIL_RSHIFT(sws_src_w, 2);
179                     sws_dst_h = AV_CEIL_RSHIFT(sws_dst_h, 2);
180                     sws_dst_w = AV_CEIL_RSHIFT(sws_dst_w, 2);
181                     break;
182                 case AV_PIX_FMT_YUV411P:
183                     sws_src_w = AV_CEIL_RSHIFT(sws_src_w, 2);
184                     sws_dst_w = AV_CEIL_RSHIFT(sws_dst_w, 2);
185                     break;
186                 default:
187                     av_log(context, AV_LOG_ERROR, "could not create SwsContext for input pixel format");
188                     return AVERROR(EIO);
189                 }
190                 sr_context->sws_context = sws_getContext(sws_src_w, sws_src_h, AV_PIX_FMT_GRAY8,
191                                                          sws_dst_w, sws_dst_h, AV_PIX_FMT_GRAY8, SWS_BICUBIC, NULL, NULL, NULL);
192                 if (!sr_context->sws_context){
193                     av_log(context, AV_LOG_ERROR, "could not create SwsContext\n");
194                     return AVERROR(ENOMEM);
195                 }
196                 sr_context->sws_slice_h = sws_src_h;
197             }
198         }
199
200         return 0;
201     }
202 }
203
204 typedef struct ThreadData{
205     uint8_t* data;
206     int data_linesize, height, width;
207 } ThreadData;
208
209 static int uint8_to_float(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
210 {
211     SRContext* sr_context = context->priv;
212     const ThreadData* td = arg;
213     const int slice_start = (td->height *  jobnr     ) / nb_jobs;
214     const int slice_end   = (td->height * (jobnr + 1)) / nb_jobs;
215     const uint8_t* src = td->data + slice_start * td->data_linesize;
216     float* dst = sr_context->input.data + slice_start * td->width;
217     int y, x;
218
219     for (y = slice_start; y < slice_end; ++y){
220         for (x = 0; x < td->width; ++x){
221             dst[x] = (float)src[x] / 255.0f;
222         }
223         src += td->data_linesize;
224         dst += td->width;
225     }
226
227     return 0;
228 }
229
230 static int float_to_uint8(AVFilterContext* context, void* arg, int jobnr, int nb_jobs)
231 {
232     SRContext* sr_context = context->priv;
233     const ThreadData* td = arg;
234     const int slice_start = (td->height *  jobnr     ) / nb_jobs;
235     const int slice_end   = (td->height * (jobnr + 1)) / nb_jobs;
236     const float* src = sr_context->output.data + slice_start * td->width;
237     uint8_t* dst = td->data + slice_start * td->data_linesize;
238     int y, x;
239
240     for (y = slice_start; y < slice_end; ++y){
241         for (x = 0; x < td->width; ++x){
242             dst[x] = (uint8_t)(255.0f * FFMIN(src[x], 1.0f));
243         }
244         src += td->width;
245         dst += td->data_linesize;
246     }
247
248     return 0;
249 }
250
251 static int filter_frame(AVFilterLink* inlink, AVFrame* in)
252 {
253     AVFilterContext* context = inlink->dst;
254     SRContext* sr_context = context->priv;
255     AVFilterLink* outlink = context->outputs[0];
256     AVFrame* out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
257     ThreadData td;
258     int nb_threads;
259     DNNReturnType dnn_result;
260
261     if (!out){
262         av_log(context, AV_LOG_ERROR, "could not allocate memory for output frame\n");
263         av_frame_free(&in);
264         return AVERROR(ENOMEM);
265     }
266     av_frame_copy_props(out, in);
267     out->height = sr_context->output.height;
268     out->width = sr_context->output.width;
269     switch (sr_context->model_type){
270     case SRCNN:
271         sws_scale(sr_context->sws_context, (const uint8_t **)in->data, in->linesize,
272                   0, sr_context->sws_slice_h, out->data, out->linesize);
273         td.data = out->data[0];
274         td.data_linesize = out->linesize[0];
275         td.height = out->height;
276         td.width = out->width;
277         break;
278     case ESPCN:
279         if (sr_context->sws_context){
280             sws_scale(sr_context->sws_context, (const uint8_t **)(in->data + 1), in->linesize + 1,
281                       0, sr_context->sws_slice_h, out->data + 1, out->linesize + 1);
282             sws_scale(sr_context->sws_context, (const uint8_t **)(in->data + 2), in->linesize + 2,
283                       0, sr_context->sws_slice_h, out->data + 2, out->linesize + 2);
284         }
285         td.data = in->data[0];
286         td.data_linesize = in->linesize[0];
287         td.height = in->height;
288         td.width = in->width;
289     }
290
291     nb_threads = ff_filter_get_nb_threads(context);
292     context->internal->execute(context, uint8_to_float, &td, NULL, FFMIN(td.height, nb_threads));
293     av_frame_free(&in);
294
295     dnn_result = (sr_context->dnn_module->execute_model)(sr_context->model);
296     if (dnn_result != DNN_SUCCESS){
297         av_log(context, AV_LOG_ERROR, "failed to execute loaded model\n");
298         return AVERROR(EIO);
299     }
300
301     td.data = out->data[0];
302     td.data_linesize = out->linesize[0];
303     td.height = out->height;
304     td.width = out->width;
305     context->internal->execute(context, float_to_uint8, &td, NULL, FFMIN(td.height, nb_threads));
306
307     return ff_filter_frame(outlink, out);
308 }
309
310 static av_cold void uninit(AVFilterContext* context)
311 {
312     SRContext* sr_context = context->priv;
313
314     if (sr_context->dnn_module){
315         (sr_context->dnn_module->free_model)(&sr_context->model);
316         av_freep(&sr_context->dnn_module);
317     }
318
319     if (sr_context->sws_context){
320         sws_freeContext(sr_context->sws_context);
321     }
322 }
323
324 static const AVFilterPad sr_inputs[] = {
325     {
326         .name         = "default",
327         .type         = AVMEDIA_TYPE_VIDEO,
328         .config_props = config_props,
329         .filter_frame = filter_frame,
330     },
331     { NULL }
332 };
333
334 static const AVFilterPad sr_outputs[] = {
335     {
336         .name = "default",
337         .type = AVMEDIA_TYPE_VIDEO,
338     },
339     { NULL }
340 };
341
342 AVFilter ff_vf_sr = {
343     .name          = "sr",
344     .description   = NULL_IF_CONFIG_SMALL("Apply DNN-based image super resolution to the input."),
345     .priv_size     = sizeof(SRContext),
346     .init          = init,
347     .uninit        = uninit,
348     .query_formats = query_formats,
349     .inputs        = sr_inputs,
350     .outputs       = sr_outputs,
351     .priv_class    = &sr_class,
352     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
353 };
354