1 /*
2  * Copyright (c) 2021 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 <float.h>
22 
23 #include "libavutil/opt.h"
24 #include "libavutil/imgutils.h"
25 #include "avfilter.h"
26 #include "drawutils.h"
27 #include "formats.h"
28 #include "internal.h"
29 #include "video.h"
30 
31 #define R 0
32 #define G 1
33 #define B 2
34 
35 typedef struct ColorTemperatureContext {
36     const AVClass *class;
37 
38     float temperature;
39     float mix;
40     float preserve;
41 
42     float color[3];
43 
44     int step;
45     int depth;
46     uint8_t rgba_map[4];
47 
48     int (*do_slice)(AVFilterContext *s, void *arg,
49                     int jobnr, int nb_jobs);
50 } ColorTemperatureContext;
51 
saturate(float input)52 static float saturate(float input)
53 {
54     return av_clipf(input, 0.f, 1.f);
55 }
56 
kelvin2rgb(float k,float * rgb)57 static void kelvin2rgb(float k, float *rgb)
58 {
59     float kelvin = k / 100.0f;
60 
61     if (kelvin <= 66.0f) {
62         rgb[0] = 1.0f;
63         rgb[1] = saturate(0.39008157876901960784f * logf(kelvin) - 0.63184144378862745098f);
64     } else {
65         const float t = fmaxf(kelvin - 60.0f, 0.0f);
66         rgb[0] = saturate(1.29293618606274509804f * powf(t, -0.1332047592f));
67         rgb[1] = saturate(1.12989086089529411765f * powf(t, -0.0755148492f));
68     }
69 
70     if (kelvin >= 66.0f)
71         rgb[2] = 1.0f;
72     else if (kelvin <= 19.0f)
73         rgb[2] = 0.0f;
74     else
75         rgb[2] = saturate(0.54320678911019607843f * logf(kelvin - 10.0f) - 1.19625408914f);
76 }
77 
lerpf(float v0,float v1,float f)78 static float lerpf(float v0, float v1, float f)
79 {
80     return v0 + (v1 - v0) * f;
81 }
82 
83 #define PROCESS()                                                   \
84     nr = r * color[0];                                              \
85     ng = g * color[1];                                              \
86     nb = b * color[2];                                              \
87                                                                     \
88     nr = lerpf(r, nr, mix);                                         \
89     ng = lerpf(g, ng, mix);                                         \
90     nb = lerpf(b, nb, mix);                                         \
91                                                                     \
92     l0 = (FFMAX3(r, g, b) + FFMIN3(r, g, b)) + FLT_EPSILON;         \
93     l1 = (FFMAX3(nr, ng, nb) + FFMIN3(nr, ng, nb)) + FLT_EPSILON;   \
94     l = l0 / l1;                                                    \
95                                                                     \
96     r = nr * l;                                                     \
97     g = ng * l;                                                     \
98     b = nb * l;                                                     \
99                                                                     \
100     nr = lerpf(nr, r, preserve);                                    \
101     ng = lerpf(ng, g, preserve);                                    \
102     nb = lerpf(nb, b, preserve);
103 
temperature_slice8(AVFilterContext * ctx,void * arg,int jobnr,int nb_jobs)104 static int temperature_slice8(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
105 {
106     ColorTemperatureContext *s = ctx->priv;
107     AVFrame *frame = arg;
108     const int width = frame->width;
109     const int height = frame->height;
110     const float mix = s->mix;
111     const float preserve = s->preserve;
112     const float *color = s->color;
113     const int slice_start = (height * jobnr) / nb_jobs;
114     const int slice_end = (height * (jobnr + 1)) / nb_jobs;
115     const int glinesize = frame->linesize[0];
116     const int blinesize = frame->linesize[1];
117     const int rlinesize = frame->linesize[2];
118     uint8_t *gptr = frame->data[0] + slice_start * glinesize;
119     uint8_t *bptr = frame->data[1] + slice_start * blinesize;
120     uint8_t *rptr = frame->data[2] + slice_start * rlinesize;
121 
122     for (int y = slice_start; y < slice_end; y++) {
123         for (int x = 0; x < width; x++) {
124             float g = gptr[x];
125             float b = bptr[x];
126             float r = rptr[x];
127             float nr, ng, nb;
128             float l0, l1, l;
129 
130             PROCESS()
131 
132             gptr[x] = av_clip_uint8(ng);
133             bptr[x] = av_clip_uint8(nb);
134             rptr[x] = av_clip_uint8(nr);
135         }
136 
137         gptr += glinesize;
138         bptr += blinesize;
139         rptr += rlinesize;
140     }
141 
142     return 0;
143 }
144 
temperature_slice16(AVFilterContext * ctx,void * arg,int jobnr,int nb_jobs)145 static int temperature_slice16(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
146 {
147     ColorTemperatureContext *s = ctx->priv;
148     AVFrame *frame = arg;
149     const int depth = s->depth;
150     const int width = frame->width;
151     const int height = frame->height;
152     const float preserve = s->preserve;
153     const float mix = s->mix;
154     const float *color = s->color;
155     const int slice_start = (height * jobnr) / nb_jobs;
156     const int slice_end = (height * (jobnr + 1)) / nb_jobs;
157     const int glinesize = frame->linesize[0] / sizeof(uint16_t);
158     const int blinesize = frame->linesize[1] / sizeof(uint16_t);
159     const int rlinesize = frame->linesize[2] / sizeof(uint16_t);
160     uint16_t *gptr = (uint16_t *)frame->data[0] + slice_start * glinesize;
161     uint16_t *bptr = (uint16_t *)frame->data[1] + slice_start * blinesize;
162     uint16_t *rptr = (uint16_t *)frame->data[2] + slice_start * rlinesize;
163 
164     for (int y = slice_start; y < slice_end; y++) {
165         for (int x = 0; x < width; x++) {
166             float g = gptr[x];
167             float b = bptr[x];
168             float r = rptr[x];
169             float nr, ng, nb;
170             float l0, l1, l;
171 
172             PROCESS()
173 
174             gptr[x] = av_clip_uintp2_c(ng, depth);
175             bptr[x] = av_clip_uintp2_c(nb, depth);
176             rptr[x] = av_clip_uintp2_c(nr, depth);
177         }
178 
179         gptr += glinesize;
180         bptr += blinesize;
181         rptr += rlinesize;
182     }
183 
184     return 0;
185 }
186 
temperature_slice8p(AVFilterContext * ctx,void * arg,int jobnr,int nb_jobs)187 static int temperature_slice8p(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
188 {
189     ColorTemperatureContext *s = ctx->priv;
190     AVFrame *frame = arg;
191     const int step = s->step;
192     const int width = frame->width;
193     const int height = frame->height;
194     const float mix = s->mix;
195     const float preserve = s->preserve;
196     const float *color = s->color;
197     const uint8_t roffset = s->rgba_map[R];
198     const uint8_t goffset = s->rgba_map[G];
199     const uint8_t boffset = s->rgba_map[B];
200     const int slice_start = (height * jobnr) / nb_jobs;
201     const int slice_end = (height * (jobnr + 1)) / nb_jobs;
202     const int linesize = frame->linesize[0];
203     uint8_t *ptr = frame->data[0] + slice_start * linesize;
204 
205     for (int y = slice_start; y < slice_end; y++) {
206         for (int x = 0; x < width; x++) {
207             float g = ptr[x * step + goffset];
208             float b = ptr[x * step + boffset];
209             float r = ptr[x * step + roffset];
210             float nr, ng, nb;
211             float l0, l1, l;
212 
213             PROCESS()
214 
215             ptr[x * step + goffset] = av_clip_uint8(ng);
216             ptr[x * step + boffset] = av_clip_uint8(nb);
217             ptr[x * step + roffset] = av_clip_uint8(nr);
218         }
219 
220         ptr += linesize;
221     }
222 
223     return 0;
224 }
225 
temperature_slice16p(AVFilterContext * ctx,void * arg,int jobnr,int nb_jobs)226 static int temperature_slice16p(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
227 {
228     ColorTemperatureContext *s = ctx->priv;
229     AVFrame *frame = arg;
230     const int step = s->step;
231     const int depth = s->depth;
232     const int width = frame->width;
233     const int height = frame->height;
234     const float preserve = s->preserve;
235     const float mix = s->mix;
236     const float *color = s->color;
237     const uint8_t roffset = s->rgba_map[R];
238     const uint8_t goffset = s->rgba_map[G];
239     const uint8_t boffset = s->rgba_map[B];
240     const int slice_start = (height * jobnr) / nb_jobs;
241     const int slice_end = (height * (jobnr + 1)) / nb_jobs;
242     const int linesize = frame->linesize[0] / sizeof(uint16_t);
243     uint16_t *ptr = (uint16_t *)frame->data[0] + slice_start * linesize;
244 
245     for (int y = slice_start; y < slice_end; y++) {
246         for (int x = 0; x < width; x++) {
247             float g = ptr[x * step + goffset];
248             float b = ptr[x * step + boffset];
249             float r = ptr[x * step + roffset];
250             float nr, ng, nb;
251             float l0, l1, l;
252 
253             PROCESS()
254 
255             ptr[x * step + goffset] = av_clip_uintp2_c(ng, depth);
256             ptr[x * step + boffset] = av_clip_uintp2_c(nb, depth);
257             ptr[x * step + roffset] = av_clip_uintp2_c(nr, depth);
258         }
259 
260         ptr += linesize;
261     }
262 
263     return 0;
264 }
265 
filter_frame(AVFilterLink * inlink,AVFrame * frame)266 static int filter_frame(AVFilterLink *inlink, AVFrame *frame)
267 {
268     AVFilterContext *ctx = inlink->dst;
269     ColorTemperatureContext *s = ctx->priv;
270 
271     kelvin2rgb(s->temperature, s->color);
272 
273     ctx->internal->execute(ctx, s->do_slice, frame, NULL,
274                            FFMIN(frame->height, ff_filter_get_nb_threads(ctx)));
275 
276     return ff_filter_frame(ctx->outputs[0], frame);
277 }
278 
query_formats(AVFilterContext * ctx)279 static av_cold int query_formats(AVFilterContext *ctx)
280 {
281     static const enum AVPixelFormat pixel_fmts[] = {
282         AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
283         AV_PIX_FMT_RGBA, AV_PIX_FMT_BGRA,
284         AV_PIX_FMT_ARGB, AV_PIX_FMT_ABGR,
285         AV_PIX_FMT_0RGB, AV_PIX_FMT_0BGR,
286         AV_PIX_FMT_RGB0, AV_PIX_FMT_BGR0,
287         AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP,
288         AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRP12,
289         AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
290         AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
291         AV_PIX_FMT_RGB48,  AV_PIX_FMT_BGR48,
292         AV_PIX_FMT_RGBA64, AV_PIX_FMT_BGRA64,
293         AV_PIX_FMT_NONE
294     };
295 
296     AVFilterFormats *formats = NULL;
297 
298     formats = ff_make_format_list(pixel_fmts);
299     if (!formats)
300         return AVERROR(ENOMEM);
301 
302     return ff_set_common_formats(ctx, formats);
303 }
304 
config_input(AVFilterLink * inlink)305 static av_cold int config_input(AVFilterLink *inlink)
306 {
307     AVFilterContext *ctx = inlink->dst;
308     ColorTemperatureContext *s = ctx->priv;
309     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
310     int planar = desc->flags & AV_PIX_FMT_FLAG_PLANAR;
311 
312     s->step = desc->nb_components;
313     if (inlink->format == AV_PIX_FMT_RGB0 ||
314         inlink->format == AV_PIX_FMT_0RGB ||
315         inlink->format == AV_PIX_FMT_BGR0 ||
316         inlink->format == AV_PIX_FMT_0BGR)
317         s->step = 4;
318 
319     s->depth = desc->comp[0].depth;
320     s->do_slice = s->depth <= 8 ? temperature_slice8 : temperature_slice16;
321     if (!planar)
322         s->do_slice = s->depth <= 8 ? temperature_slice8p : temperature_slice16p;
323 
324     ff_fill_rgba_map(s->rgba_map, inlink->format);
325 
326     return 0;
327 }
328 
329 static const AVFilterPad inputs[] = {
330     {
331         .name           = "default",
332         .type           = AVMEDIA_TYPE_VIDEO,
333         .filter_frame   = filter_frame,
334         .config_props   = config_input,
335         .needs_writable = 1,
336     },
337     { NULL }
338 };
339 
340 static const AVFilterPad outputs[] = {
341     {
342         .name = "default",
343         .type = AVMEDIA_TYPE_VIDEO,
344     },
345     { NULL }
346 };
347 
348 #define OFFSET(x) offsetof(ColorTemperatureContext, x)
349 #define VF AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
350 
351 static const AVOption colortemperature_options[] = {
352     { "temperature", "set the temperature in Kelvin",          OFFSET(temperature), AV_OPT_TYPE_FLOAT, {.dbl=6500}, 1000,  40000, VF },
353     { "mix",         "set the mix with filtered output",       OFFSET(mix),         AV_OPT_TYPE_FLOAT, {.dbl=1},       0,      1, VF },
354     { "pl",          "set the amount of preserving lightness", OFFSET(preserve),    AV_OPT_TYPE_FLOAT, {.dbl=0},       0,      1, VF },
355     { NULL }
356 };
357 
358 AVFILTER_DEFINE_CLASS(colortemperature);
359 
360 AVFilter ff_vf_colortemperature = {
361     .name          = "colortemperature",
362     .description   = NULL_IF_CONFIG_SMALL("Adjust color temperature of video."),
363     .priv_size     = sizeof(ColorTemperatureContext),
364     .priv_class    = &colortemperature_class,
365     .query_formats = query_formats,
366     .inputs        = inputs,
367     .outputs       = outputs,
368     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
369     .process_command = ff_filter_process_command,
370 };
371