blob: 88414783bc8001bb39f7f53e02aa6731dae403ab
1 | /* |
2 | * Copyright (c) 2015 Timo Rothenpieler <timo@rothenpieler.org> |
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 | #include "libavutil/opt.h" |
22 | #include "libavutil/imgutils.h" |
23 | #include "avfilter.h" |
24 | #include "formats.h" |
25 | #include "internal.h" |
26 | #include "video.h" |
27 | |
28 | typedef struct ChromakeyContext { |
29 | const AVClass *class; |
30 | |
31 | uint8_t chromakey_rgba[4]; |
32 | uint8_t chromakey_uv[2]; |
33 | |
34 | float similarity; |
35 | float blend; |
36 | |
37 | int is_yuv; |
38 | |
39 | int hsub_log2; |
40 | int vsub_log2; |
41 | } ChromakeyContext; |
42 | |
43 | static uint8_t do_chromakey_pixel(ChromakeyContext *ctx, uint8_t u[9], uint8_t v[9]) |
44 | { |
45 | double diff = 0.0; |
46 | int du, dv, i; |
47 | |
48 | for (i = 0; i < 9; ++i) { |
49 | du = (int)u[i] - ctx->chromakey_uv[0]; |
50 | dv = (int)v[i] - ctx->chromakey_uv[1]; |
51 | |
52 | diff += sqrt((du * du + dv * dv) / (255.0 * 255.0)); |
53 | } |
54 | |
55 | diff /= 9.0; |
56 | |
57 | if (ctx->blend > 0.0001) { |
58 | return av_clipd((diff - ctx->similarity) / ctx->blend, 0.0, 1.0) * 255.0; |
59 | } else { |
60 | return (diff > ctx->similarity) ? 255 : 0; |
61 | } |
62 | } |
63 | |
64 | static av_always_inline void get_pixel_uv(AVFrame *frame, int hsub_log2, int vsub_log2, int x, int y, uint8_t *u, uint8_t *v) |
65 | { |
66 | if (x < 0 || x >= frame->width || y < 0 || y >= frame->height) |
67 | return; |
68 | |
69 | x >>= hsub_log2; |
70 | y >>= vsub_log2; |
71 | |
72 | *u = frame->data[1][frame->linesize[1] * y + x]; |
73 | *v = frame->data[2][frame->linesize[2] * y + x]; |
74 | } |
75 | |
76 | static int do_chromakey_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs) |
77 | { |
78 | AVFrame *frame = arg; |
79 | |
80 | const int slice_start = (frame->height * jobnr) / nb_jobs; |
81 | const int slice_end = (frame->height * (jobnr + 1)) / nb_jobs; |
82 | |
83 | ChromakeyContext *ctx = avctx->priv; |
84 | |
85 | int x, y, xo, yo; |
86 | uint8_t u[9], v[9]; |
87 | |
88 | memset(u, ctx->chromakey_uv[0], sizeof(u)); |
89 | memset(v, ctx->chromakey_uv[1], sizeof(v)); |
90 | |
91 | for (y = slice_start; y < slice_end; ++y) { |
92 | for (x = 0; x < frame->width; ++x) { |
93 | for (yo = 0; yo < 3; ++yo) { |
94 | for (xo = 0; xo < 3; ++xo) { |
95 | get_pixel_uv(frame, ctx->hsub_log2, ctx->vsub_log2, x + xo - 1, y + yo - 1, &u[yo * 3 + xo], &v[yo * 3 + xo]); |
96 | } |
97 | } |
98 | |
99 | frame->data[3][frame->linesize[3] * y + x] = do_chromakey_pixel(ctx, u, v); |
100 | } |
101 | } |
102 | |
103 | return 0; |
104 | } |
105 | |
106 | static int filter_frame(AVFilterLink *link, AVFrame *frame) |
107 | { |
108 | AVFilterContext *avctx = link->dst; |
109 | int res; |
110 | |
111 | if (res = avctx->internal->execute(avctx, do_chromakey_slice, frame, NULL, FFMIN(frame->height, ff_filter_get_nb_threads(avctx)))) |
112 | return res; |
113 | |
114 | return ff_filter_frame(avctx->outputs[0], frame); |
115 | } |
116 | |
117 | #define FIXNUM(x) lrint((x) * (1 << 10)) |
118 | #define RGB_TO_U(rgb) (((- FIXNUM(0.16874) * rgb[0] - FIXNUM(0.33126) * rgb[1] + FIXNUM(0.50000) * rgb[2] + (1 << 9) - 1) >> 10) + 128) |
119 | #define RGB_TO_V(rgb) ((( FIXNUM(0.50000) * rgb[0] - FIXNUM(0.41869) * rgb[1] - FIXNUM(0.08131) * rgb[2] + (1 << 9) - 1) >> 10) + 128) |
120 | |
121 | static av_cold int initialize_chromakey(AVFilterContext *avctx) |
122 | { |
123 | ChromakeyContext *ctx = avctx->priv; |
124 | |
125 | if (ctx->is_yuv) { |
126 | ctx->chromakey_uv[0] = ctx->chromakey_rgba[1]; |
127 | ctx->chromakey_uv[1] = ctx->chromakey_rgba[2]; |
128 | } else { |
129 | ctx->chromakey_uv[0] = RGB_TO_U(ctx->chromakey_rgba); |
130 | ctx->chromakey_uv[1] = RGB_TO_V(ctx->chromakey_rgba); |
131 | } |
132 | |
133 | return 0; |
134 | } |
135 | |
136 | static av_cold int query_formats(AVFilterContext *avctx) |
137 | { |
138 | static const enum AVPixelFormat pixel_fmts[] = { |
139 | AV_PIX_FMT_YUVA420P, |
140 | AV_PIX_FMT_YUVA422P, |
141 | AV_PIX_FMT_YUVA444P, |
142 | AV_PIX_FMT_NONE |
143 | }; |
144 | |
145 | AVFilterFormats *formats = NULL; |
146 | |
147 | formats = ff_make_format_list(pixel_fmts); |
148 | if (!formats) |
149 | return AVERROR(ENOMEM); |
150 | |
151 | return ff_set_common_formats(avctx, formats); |
152 | } |
153 | |
154 | static av_cold int config_input(AVFilterLink *inlink) |
155 | { |
156 | AVFilterContext *avctx = inlink->dst; |
157 | ChromakeyContext *ctx = avctx->priv; |
158 | const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format); |
159 | |
160 | ctx->hsub_log2 = desc->log2_chroma_w; |
161 | ctx->vsub_log2 = desc->log2_chroma_h; |
162 | |
163 | return 0; |
164 | } |
165 | |
166 | static const AVFilterPad chromakey_inputs[] = { |
167 | { |
168 | .name = "default", |
169 | .type = AVMEDIA_TYPE_VIDEO, |
170 | .needs_writable = 1, |
171 | .filter_frame = filter_frame, |
172 | .config_props = config_input, |
173 | }, |
174 | { NULL } |
175 | }; |
176 | |
177 | static const AVFilterPad chromakey_outputs[] = { |
178 | { |
179 | .name = "default", |
180 | .type = AVMEDIA_TYPE_VIDEO, |
181 | }, |
182 | { NULL } |
183 | }; |
184 | |
185 | #define OFFSET(x) offsetof(ChromakeyContext, x) |
186 | #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM |
187 | |
188 | static const AVOption chromakey_options[] = { |
189 | { "color", "set the chromakey key color", OFFSET(chromakey_rgba), AV_OPT_TYPE_COLOR, { .str = "black" }, CHAR_MIN, CHAR_MAX, FLAGS }, |
190 | { "similarity", "set the chromakey similarity value", OFFSET(similarity), AV_OPT_TYPE_FLOAT, { .dbl = 0.01 }, 0.01, 1.0, FLAGS }, |
191 | { "blend", "set the chromakey key blend value", OFFSET(blend), AV_OPT_TYPE_FLOAT, { .dbl = 0.0 }, 0.0, 1.0, FLAGS }, |
192 | { "yuv", "color parameter is in yuv instead of rgb", OFFSET(is_yuv), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, FLAGS }, |
193 | { NULL } |
194 | }; |
195 | |
196 | AVFILTER_DEFINE_CLASS(chromakey); |
197 | |
198 | AVFilter ff_vf_chromakey = { |
199 | .name = "chromakey", |
200 | .description = NULL_IF_CONFIG_SMALL("Turns a certain color into transparency. Operates on YUV colors."), |
201 | .priv_size = sizeof(ChromakeyContext), |
202 | .priv_class = &chromakey_class, |
203 | .init = initialize_chromakey, |
204 | .query_formats = query_formats, |
205 | .inputs = chromakey_inputs, |
206 | .outputs = chromakey_outputs, |
207 | .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS, |
208 | }; |
209 |