2 * Copyright (c) 2019 Xuewei Meng
4 * This file is part of FFmpeg.
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.
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.
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
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
27 #include "libavformat/avio.h"
28 #include "libavutil/opt.h"
30 #include "dnn_filter_common.h"
34 typedef struct DRContext {
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" },
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 },
57 AVFILTER_DEFINE_CLASS(derain);
59 static int query_formats(AVFilterContext *ctx)
61 AVFilterFormats *formats;
62 const enum AVPixelFormat pixel_fmts[] = {
67 formats = ff_make_format_list(pixel_fmts);
69 return ff_set_common_formats(ctx, formats);
72 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
74 AVFilterContext *ctx = inlink->dst;
75 AVFilterLink *outlink = ctx->outputs[0];
76 DRContext *dr_context = ctx->priv;
77 DNNReturnType dnn_result;
80 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
82 av_log(ctx, AV_LOG_ERROR, "could not allocate memory for output frame\n");
84 return AVERROR(ENOMEM);
86 av_frame_copy_props(out, in);
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");
97 return ff_filter_frame(outlink, out);
100 static av_cold int init(AVFilterContext *ctx)
102 DRContext *dr_context = ctx->priv;
103 return ff_dnn_init(&dr_context->dnnctx, DFT_PROCESS_FRAME, ctx);
106 static av_cold void uninit(AVFilterContext *ctx)
108 DRContext *dr_context = ctx->priv;
109 ff_dnn_uninit(&dr_context->dnnctx);
112 static const AVFilterPad derain_inputs[] = {
115 .type = AVMEDIA_TYPE_VIDEO,
116 .filter_frame = filter_frame,
121 static const AVFilterPad derain_outputs[] = {
124 .type = AVMEDIA_TYPE_VIDEO,
129 const AVFilter ff_vf_derain = {
131 .description = NULL_IF_CONFIG_SMALL("Apply derain filter to the input."),
132 .priv_size = sizeof(DRContext),
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,