d502080d45f207c208b2f20f0d859a04b39f3808
[ffmpeg.git] / libavfilter / vf_colorbalance.c
1 /*
2  * Copyright (c) 2013 Paul B Mahol
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/pixdesc.h"
23 #include "avfilter.h"
24 #include "drawutils.h"
25 #include "formats.h"
26 #include "internal.h"
27 #include "video.h"
28
29 #define R 0
30 #define G 1
31 #define B 2
32 #define A 3
33
34 typedef struct Range {
35     double shadows;
36     double midtones;
37     double highlights;
38 } Range;
39
40 typedef struct ColorBalanceContext {
41     const AVClass *class;
42     Range cyan_red;
43     Range magenta_green;
44     Range yellow_blue;
45
46     uint16_t lut[3][65536];
47
48     uint8_t rgba_map[4];
49     int step;
50
51     int (*clip)(int x);
52     void (*apply_lut)(AVFilterContext *ctx, AVFrame *in, AVFrame *out);
53 } ColorBalanceContext;
54
55 #define OFFSET(x) offsetof(ColorBalanceContext, x)
56 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
57 static const AVOption colorbalance_options[] = {
58     { "rs", "set red shadows",      OFFSET(cyan_red.shadows),         AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
59     { "gs", "set green shadows",    OFFSET(magenta_green.shadows),    AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
60     { "bs", "set blue shadows",     OFFSET(yellow_blue.shadows),      AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
61     { "rm", "set red midtones",     OFFSET(cyan_red.midtones),        AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
62     { "gm", "set green midtones",   OFFSET(magenta_green.midtones),   AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
63     { "bm", "set blue midtones",    OFFSET(yellow_blue.midtones),     AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
64     { "rh", "set red highlights",   OFFSET(cyan_red.highlights),      AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
65     { "gh", "set green highlights", OFFSET(magenta_green.highlights), AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
66     { "bh", "set blue highlights",  OFFSET(yellow_blue.highlights),   AV_OPT_TYPE_DOUBLE, {.dbl=0}, -1, 1, FLAGS },
67     { NULL }
68 };
69
70 AVFILTER_DEFINE_CLASS(colorbalance);
71
72 static int query_formats(AVFilterContext *ctx)
73 {
74     static const enum AVPixelFormat pix_fmts[] = {
75         AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
76         AV_PIX_FMT_RGBA,  AV_PIX_FMT_BGRA,
77         AV_PIX_FMT_ABGR,  AV_PIX_FMT_ARGB,
78         AV_PIX_FMT_0BGR,  AV_PIX_FMT_0RGB,
79         AV_PIX_FMT_RGB0,  AV_PIX_FMT_BGR0,
80         AV_PIX_FMT_RGB48,  AV_PIX_FMT_BGR48,
81         AV_PIX_FMT_RGBA64, AV_PIX_FMT_BGRA64,
82         AV_PIX_FMT_GBRP,   AV_PIX_FMT_GBRAP,
83         AV_PIX_FMT_GBRP9,
84         AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRAP10,
85         AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRAP12,
86         AV_PIX_FMT_GBRP14,
87         AV_PIX_FMT_GBRP16, AV_PIX_FMT_GBRAP16,
88         AV_PIX_FMT_NONE
89     };
90     AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
91     if (!fmts_list)
92         return AVERROR(ENOMEM);
93     return ff_set_common_formats(ctx, fmts_list);
94 }
95
96 static void apply_lut8_p(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
97 {
98     ColorBalanceContext *s = ctx->priv;
99     const uint8_t *srcg = in->data[0];
100     const uint8_t *srcb = in->data[1];
101     const uint8_t *srcr = in->data[2];
102     const uint8_t *srca = in->data[3];
103     uint8_t *dstg = out->data[0];
104     uint8_t *dstb = out->data[1];
105     uint8_t *dstr = out->data[2];
106     uint8_t *dsta = out->data[3];
107     int i, j;
108
109     for (i = 0; i < out->height; i++) {
110         for (j = 0; j < out->width; j++) {
111             dstg[j] = s->lut[G][srcg[j]];
112             dstb[j] = s->lut[B][srcb[j]];
113             dstr[j] = s->lut[R][srcr[j]];
114             if (in != out && out->linesize[3])
115                 dsta[j] = srca[j];
116         }
117
118         srcg += in->linesize[0];
119         srcb += in->linesize[1];
120         srcr += in->linesize[2];
121         srca += in->linesize[3];
122         dstg += out->linesize[0];
123         dstb += out->linesize[1];
124         dstr += out->linesize[2];
125         dsta += out->linesize[3];
126     }
127 }
128
129 static void apply_lut16_p(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
130 {
131     ColorBalanceContext *s = ctx->priv;
132     const uint16_t *srcg = (const uint16_t *)in->data[0];
133     const uint16_t *srcb = (const uint16_t *)in->data[1];
134     const uint16_t *srcr = (const uint16_t *)in->data[2];
135     const uint16_t *srca = (const uint16_t *)in->data[3];
136     uint16_t *dstg = (uint16_t *)out->data[0];
137     uint16_t *dstb = (uint16_t *)out->data[1];
138     uint16_t *dstr = (uint16_t *)out->data[2];
139     uint16_t *dsta = (uint16_t *)out->data[3];
140     int i, j;
141
142     for (i = 0; i < out->height; i++) {
143         for (j = 0; j < out->width; j++) {
144             dstg[j] = s->lut[G][srcg[j]];
145             dstb[j] = s->lut[B][srcb[j]];
146             dstr[j] = s->lut[R][srcr[j]];
147             if (in != out && out->linesize[3])
148                 dsta[j] = srca[j];
149         }
150
151         srcg += in->linesize[0] / 2;
152         srcb += in->linesize[1] / 2;
153         srcr += in->linesize[2] / 2;
154         srca += in->linesize[3] / 2;
155         dstg += out->linesize[0] / 2;
156         dstb += out->linesize[1] / 2;
157         dstr += out->linesize[2] / 2;
158         dsta += out->linesize[3] / 2;
159     }
160 }
161
162 static void apply_lut8(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
163 {
164     ColorBalanceContext *s = ctx->priv;
165     AVFilterLink *outlink = ctx->outputs[0];
166     const uint8_t *srcrow = in->data[0];
167     const uint8_t roffset = s->rgba_map[R];
168     const uint8_t goffset = s->rgba_map[G];
169     const uint8_t boffset = s->rgba_map[B];
170     const uint8_t aoffset = s->rgba_map[A];
171     const int step = s->step;
172     uint8_t *dstrow;
173     int i, j;
174
175     dstrow = out->data[0];
176     for (i = 0; i < outlink->h; i++) {
177         const uint8_t *src = srcrow;
178         uint8_t *dst = dstrow;
179
180         for (j = 0; j < outlink->w * step; j += step) {
181             dst[j + roffset] = s->lut[R][src[j + roffset]];
182             dst[j + goffset] = s->lut[G][src[j + goffset]];
183             dst[j + boffset] = s->lut[B][src[j + boffset]];
184             if (in != out && step == 4)
185                 dst[j + aoffset] = src[j + aoffset];
186         }
187
188         srcrow += in->linesize[0];
189         dstrow += out->linesize[0];
190     }
191 }
192
193 static void apply_lut16(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
194 {
195     ColorBalanceContext *s = ctx->priv;
196     AVFilterLink *outlink = ctx->outputs[0];
197     const uint16_t *srcrow = (const uint16_t *)in->data[0];
198     const uint8_t roffset = s->rgba_map[R];
199     const uint8_t goffset = s->rgba_map[G];
200     const uint8_t boffset = s->rgba_map[B];
201     const uint8_t aoffset = s->rgba_map[A];
202     const int step = s->step / 2;
203     uint16_t *dstrow;
204     int i, j;
205
206     dstrow = (uint16_t *)out->data[0];
207     for (i = 0; i < outlink->h; i++) {
208         const uint16_t *src = srcrow;
209         uint16_t *dst = dstrow;
210
211         for (j = 0; j < outlink->w * step; j += step) {
212             dst[j + roffset] = s->lut[R][src[j + roffset]];
213             dst[j + goffset] = s->lut[G][src[j + goffset]];
214             dst[j + boffset] = s->lut[B][src[j + boffset]];
215             if (in != out && step == 4)
216                 dst[j + aoffset] = src[j + aoffset];
217         }
218
219         srcrow += in->linesize[0] / 2;
220         dstrow += out->linesize[0] / 2;
221     }
222 }
223
224 static int config_output(AVFilterLink *outlink)
225 {
226     AVFilterContext *ctx = outlink->src;
227     ColorBalanceContext *s = ctx->priv;
228     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(outlink->format);
229     const int depth = desc->comp[0].depth;
230     const int max = 1 << depth;
231     const int planar = av_pix_fmt_count_planes(outlink->format) > 1;
232     double *shadows, *midtones, *highlights, *buffer;
233     int i, r, g, b;
234
235     if (max == 256 && planar) {
236         s->apply_lut = apply_lut8_p;
237     } else if (planar) {
238         s->apply_lut = apply_lut16_p;
239     } else if (max == 256) {
240         s->apply_lut = apply_lut8;
241     } else {
242         s->apply_lut = apply_lut16;
243     }
244
245     buffer = av_malloc(max * 3 * sizeof(*buffer));
246     if (!buffer)
247         return AVERROR(ENOMEM);
248
249     shadows    = buffer + max * 0;
250     midtones   = buffer + max * 1;
251     highlights = buffer + max * 2;
252
253     for (i = 0; i < max; i++) {
254         const double L = 0.333 * (max - 1);
255         const double M = 0.7 * (max - 1);
256         const double H = 1 * (max - 1);
257         double low = av_clipd((i - L) / (-max * 0.25) + 0.5, 0, 1) * M;
258         double mid = av_clipd((i - L) / ( max * 0.25) + 0.5, 0, 1) *
259                      av_clipd((i + L - H) / (-max * 0.25) + 0.5, 0, 1) * M;
260
261         shadows[i] = low;
262         midtones[i] = mid;
263         highlights[max - i] = low;
264     }
265
266     for (i = 0; i < max; i++) {
267         r = g = b = i;
268
269         r = av_clip_uintp2(r + s->cyan_red.shadows         * shadows[r],    depth);
270         r = av_clip_uintp2(r + s->cyan_red.midtones        * midtones[r],   depth);
271         r = av_clip_uintp2(r + s->cyan_red.highlights      * highlights[r], depth);
272
273         g = av_clip_uintp2(g + s->magenta_green.shadows    * shadows[g],    depth);
274         g = av_clip_uintp2(g + s->magenta_green.midtones   * midtones[g],   depth);
275         g = av_clip_uintp2(g + s->magenta_green.highlights * highlights[g], depth);
276
277         b = av_clip_uintp2(b + s->yellow_blue.shadows      * shadows[b],    depth);
278         b = av_clip_uintp2(b + s->yellow_blue.midtones     * midtones[b],   depth);
279         b = av_clip_uintp2(b + s->yellow_blue.highlights   * highlights[b], depth);
280
281         s->lut[R][i] = r;
282         s->lut[G][i] = g;
283         s->lut[B][i] = b;
284     }
285
286     av_free(buffer);
287
288     ff_fill_rgba_map(s->rgba_map, outlink->format);
289     s->step = av_get_padded_bits_per_pixel(desc) >> 3;
290
291     return 0;
292 }
293
294 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
295 {
296     AVFilterContext *ctx = inlink->dst;
297     ColorBalanceContext *s = ctx->priv;
298     AVFilterLink *outlink = ctx->outputs[0];
299     AVFrame *out;
300
301     if (av_frame_is_writable(in)) {
302         out = in;
303     } else {
304         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
305         if (!out) {
306             av_frame_free(&in);
307             return AVERROR(ENOMEM);
308         }
309         av_frame_copy_props(out, in);
310     }
311
312     s->apply_lut(ctx, in, out);
313
314     if (in != out)
315         av_frame_free(&in);
316     return ff_filter_frame(outlink, out);
317 }
318
319 static const AVFilterPad colorbalance_inputs[] = {
320     {
321         .name         = "default",
322         .type         = AVMEDIA_TYPE_VIDEO,
323         .filter_frame = filter_frame,
324     },
325     { NULL }
326 };
327
328 static const AVFilterPad colorbalance_outputs[] = {
329     {
330         .name         = "default",
331         .type         = AVMEDIA_TYPE_VIDEO,
332         .config_props = config_output,
333     },
334     { NULL }
335 };
336
337 AVFilter ff_vf_colorbalance = {
338     .name          = "colorbalance",
339     .description   = NULL_IF_CONFIG_SMALL("Adjust the color balance."),
340     .priv_size     = sizeof(ColorBalanceContext),
341     .priv_class    = &colorbalance_class,
342     .query_formats = query_formats,
343     .inputs        = colorbalance_inputs,
344     .outputs       = colorbalance_outputs,
345     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
346 };