1 /*
2 * Copyright (c) 2019 Xuewei Meng
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 derain filter using deep convolutional networks.
24 * http://openaccess.thecvf.com/content_ECCV_2018/html/Xia_Li_Recurrent_Squeeze-and-Excitation_Context_ECCV_2018_paper.html
25 */
26
27 #include "libavformat/avio.h"
28 #include "libavutil/opt.h"
29 #include "avfilter.h"
30 #include "dnn_filter_common.h"
31 #include "formats.h"
32 #include "internal.h"
33
34 typedef struct DRContext {
35 const AVClass *class;
36 DnnContext dnnctx;
37 int filter_type;
38 } DRContext;
39
40 #define OFFSET(x) offsetof(DRContext, x)
41 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM
42 static const AVOption derain_options[] = {
43 { "filter_type", "filter type(derain/dehaze)", OFFSET(filter_type), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS, "type" },
44 { "derain", "derain filter flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "type" },
45 { "dehaze", "dehaze filter flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "type" },
46 { "dnn_backend", "DNN backend", OFFSET(dnnctx.backend_type), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS, "backend" },
47 { "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" },
48 #if (CONFIG_LIBTENSORFLOW == 1)
49 { "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" },
50 #endif
51 { "model", "path to model file", OFFSET(dnnctx.model_filename), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS },
52 { "input", "input name of the model", OFFSET(dnnctx.model_inputname), AV_OPT_TYPE_STRING, { .str = "x" }, 0, 0, FLAGS },
53 { "output", "output name of the model", OFFSET(dnnctx.model_outputname), AV_OPT_TYPE_STRING, { .str = "y" }, 0, 0, FLAGS },
54 { NULL }
55 };
56
57 AVFILTER_DEFINE_CLASS(derain);
58
query_formats(AVFilterContext * ctx)59 static int query_formats(AVFilterContext *ctx)
60 {
61 AVFilterFormats *formats;
62 const enum AVPixelFormat pixel_fmts[] = {
63 AV_PIX_FMT_RGB24,
64 AV_PIX_FMT_NONE
65 };
66
67 formats = ff_make_format_list(pixel_fmts);
68
69 return ff_set_common_formats(ctx, formats);
70 }
71
filter_frame(AVFilterLink * inlink,AVFrame * in)72 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
73 {
74 AVFilterContext *ctx = inlink->dst;
75 AVFilterLink *outlink = ctx->outputs[0];
76 DRContext *dr_context = ctx->priv;
77 DNNReturnType dnn_result;
78 AVFrame *out;
79
80 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
81 if (!out) {
82 av_log(ctx, AV_LOG_ERROR, "could not allocate memory for output frame\n");
83 av_frame_free(&in);
84 return AVERROR(ENOMEM);
85 }
86 av_frame_copy_props(out, in);
87
88 dnn_result = ff_dnn_execute_model(&dr_context->dnnctx, in, out);
89 if (dnn_result != DNN_SUCCESS){
90 av_log(ctx, AV_LOG_ERROR, "failed to execute model\n");
91 av_frame_free(&in);
92 return AVERROR(EIO);
93 }
94
95 av_frame_free(&in);
96
97 return ff_filter_frame(outlink, out);
98 }
99
init(AVFilterContext * ctx)100 static av_cold int init(AVFilterContext *ctx)
101 {
102 DRContext *dr_context = ctx->priv;
103 return ff_dnn_init(&dr_context->dnnctx, DFT_PROCESS_FRAME, ctx);
104 }
105
uninit(AVFilterContext * ctx)106 static av_cold void uninit(AVFilterContext *ctx)
107 {
108 DRContext *dr_context = ctx->priv;
109 ff_dnn_uninit(&dr_context->dnnctx);
110 }
111
112 static const AVFilterPad derain_inputs[] = {
113 {
114 .name = "default",
115 .type = AVMEDIA_TYPE_VIDEO,
116 .filter_frame = filter_frame,
117 },
118 { NULL }
119 };
120
121 static const AVFilterPad derain_outputs[] = {
122 {
123 .name = "default",
124 .type = AVMEDIA_TYPE_VIDEO,
125 },
126 { NULL }
127 };
128
129 AVFilter ff_vf_derain = {
130 .name = "derain",
131 .description = NULL_IF_CONFIG_SMALL("Apply derain filter to the input."),
132 .priv_size = sizeof(DRContext),
133 .init = init,
134 .uninit = uninit,
135 .query_formats = query_formats,
136 .inputs = derain_inputs,
137 .outputs = derain_outputs,
138 .priv_class = &derain_class,
139 .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
140 };
141