1 /*
2 * Copyright (c) 2022 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 /**
22 * @file
23 * Compute a look-up table from map of colors.
24 */
25
26 #include "libavutil/attributes.h"
27 #include "libavutil/avassert.h"
28 #include "libavutil/common.h"
29 #include "libavutil/opt.h"
30 #include "avfilter.h"
31 #include "internal.h"
32 #include "framesync.h"
33 #include "video.h"
34
35 #define MAX_SIZE 64
36
37 enum KernelType {
38 EUCLIDEAN,
39 WEUCLIDEAN,
40 NB_KERNELS,
41 };
42
43 typedef struct ColorMapContext {
44 const AVClass *class;
45 int w, h;
46 int size;
47 int nb_maps;
48 int changed[2];
49
50 float source[MAX_SIZE][4];
51 float ttarget[MAX_SIZE][4];
52 float target[MAX_SIZE][4];
53 float icoeff[4][4];
54 float coeff[MAX_SIZE][4];
55
56 int target_type;
57 int kernel_type;
58 float (*kernel)(const float *x, const float *y);
59
60 FFFrameSync fs;
61
62 double A[(MAX_SIZE + 4) * (MAX_SIZE + 4)];
63 double b[MAX_SIZE + 4];
64 int pivot[MAX_SIZE + 4];
65 } ColorMapContext;
66
67 #define OFFSET(x) offsetof(ColorMapContext, x)
68 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
69
70 static const AVOption colormap_options[] = {
71 { "patch_size", "set patch size", OFFSET(w), AV_OPT_TYPE_IMAGE_SIZE, {.str = "64x64"}, 0, 0, FLAGS },
72 { "nb_patches", "set number of patches", OFFSET(size), AV_OPT_TYPE_INT, {.i64 = 0}, 0, MAX_SIZE, FLAGS },
73 { "type", "set the target type used", OFFSET(target_type), AV_OPT_TYPE_INT, {.i64=1}, 0, 1, FLAGS, "type" },
74 { "relative", "the target colors are relative", 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 1, FLAGS, "type" },
75 { "absolute", "the target colors are absolute", 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 1, FLAGS, "type" },
76 { "kernel", "set the kernel used for measuring color difference", OFFSET(kernel_type), AV_OPT_TYPE_INT, {.i64=0}, 0, NB_KERNELS-1, FLAGS, "kernel" },
77 { "euclidean", "square root of sum of squared differences", 0, AV_OPT_TYPE_CONST, {.i64=EUCLIDEAN}, 0, 0, FLAGS, "kernel" },
78 { "weuclidean", "weighted square root of sum of squared differences",0, AV_OPT_TYPE_CONST, {.i64=WEUCLIDEAN}, 0, 0, FLAGS, "kernel" },
79 { NULL }
80 };
81
gauss_make_triangular(double * A,int * p,int n)82 static int gauss_make_triangular(double *A, int *p, int n)
83 {
84 p[n - 1] = n - 1;
85 for (int k = 0; k < n; k++) {
86 double t1;
87 int m = k;
88
89 for (int i = k + 1; i < n; i++)
90 if (fabs(A[k + n * i]) > fabs(A[k + n * m]))
91 m = i;
92 p[k] = m;
93 t1 = A[k + n * m];
94 A[k + n * m] = A[k + n * k];
95 A[k + n * k] = t1;
96 if (t1 != 0) {
97 for (int i = k + 1; i < n; i++)
98 A[k + n * i] /= -t1;
99 if (k != m)
100 for (int i = k + 1; i < n; i++) {
101 double t2 = A[i + n * m];
102 A[i + n * m] = A[i + n * k];
103 A[i + n * k] = t2;
104 }
105 for (int j = k + 1; j < n; j++)
106 for (int i = k + 1; i < n; i++)
107 A[i + n * j] += A[k + j * n] * A[i + k * n];
108 } else {
109 return 0;
110 }
111 }
112
113 return 1;
114 }
115
gauss_solve_triangular(const double * A,const int * p,double * b,int n)116 static void gauss_solve_triangular(const double *A, const int *p, double *b, int n)
117 {
118 for(int k = 0; k < n - 1; k++) {
119 int m = p[k];
120 double t = b[m];
121 b[m] = b[k];
122 b[k] = t;
123 for (int i = k + 1; i < n; i++)
124 b[i] += A[k + n * i] * t;
125 }
126
127 for(int k = n - 1; k > 0; k--) {
128 double t = b[k] /= A[k + n * k];
129 for (int i = 0; i < k; i++)
130 b[i] -= A[k + n * i] * t;
131 }
132
133 b[0] /= A[0 + 0 * n];
134 }
135
gauss_solve(double * A,double * b,int n)136 static int gauss_solve(double *A, double *b, int n)
137 {
138 int p[3] = { 0 };
139
140 av_assert2(n <= FF_ARRAY_ELEMS(p));
141
142 if (!gauss_make_triangular(A, p, n))
143 return 1;
144
145 gauss_solve_triangular(A, p, b, n);
146
147 return 0;
148 }
149
150 #define P2(x) ((x)*(x))
151
euclidean_kernel(const float * x,const float * y)152 static float euclidean_kernel(const float *x, const float *y)
153 {
154 const float d2 = P2(x[0]-y[0]) +
155 P2(x[1]-y[1]) +
156 P2(x[2]-y[2]);
157 return sqrtf(d2);
158 }
159
weuclidean_kernel(const float * x,const float * y)160 static float weuclidean_kernel(const float *x, const float *y)
161 {
162 const float rm = (x[0] + y[0]) * 0.5f;
163 const float d2 = P2(x[0]-y[0]) * (2.f + rm) +
164 P2(x[1]-y[1]) * 4.f +
165 P2(x[2]-y[2]) * (3.f - rm);
166 return sqrtf(d2);
167 }
168
build_map(AVFilterContext * ctx)169 static void build_map(AVFilterContext *ctx)
170 {
171 ColorMapContext *s = ctx->priv;
172
173 for (int j = 0; j < s->nb_maps; j++) {
174 s->target[j][0] = s->target_type == 0 ? s->source[j][0] + s->ttarget[j][0] : s->ttarget[j][0];
175 s->target[j][1] = s->target_type == 0 ? s->source[j][1] + s->ttarget[j][1] : s->ttarget[j][1];
176 s->target[j][2] = s->target_type == 0 ? s->source[j][2] + s->ttarget[j][2] : s->ttarget[j][2];
177 }
178
179 for (int c = 0; c < 3; c++) {
180 for (int j = 0; j < s->nb_maps; j++)
181 s->coeff[j][c] = 0.f;
182
183 for (int j = 0; j < 4; j++) {
184 s->icoeff[j][c] = 0;
185 s->icoeff[j][c] = 0;
186 s->icoeff[j][c] = 0;
187 }
188
189 s->icoeff[c+1][c] = 1.f;
190
191 switch (s->nb_maps) {
192 case 1:
193 {
194 float div = fabsf(s->source[0][c]) < 1e-6f ? 1e-6f : s->source[0][c];
195 s->icoeff[c][1+c] = s->target[0][c] / div;
196 }
197 break;
198 case 2:
199 {
200 double A[2 * 2] = { 1, s->source[0][c],
201 1, s->source[1][c] };
202 double b[2] = { s->target[0][c], s->target[1][c] };
203
204 if (gauss_solve(A, b, 2))
205 continue;
206
207 s->icoeff[0 ][c] = b[0];
208 s->icoeff[1+c][c] = b[1];
209 }
210 break;
211 case 3:
212 {
213 const uint8_t idx[3][3] = {{ 0, 1, 2 },
214 { 1, 0, 2 },
215 { 2, 0, 1 }};
216 const uint8_t didx[3][4] = {{ 0, 1, 2, 2 },
217 { 0, 2, 1, 2 },
218 { 0, 2, 2, 1 }};
219 const int C0 = idx[c][0];
220 const int C1 = idx[c][1];
221 const int C2 = idx[c][2];
222 double A[3 * 3] = { 1, s->source[0][C0], s->source[0][C1] + s->source[0][C2],
223 1, s->source[1][C0], s->source[1][C1] + s->source[1][C2],
224 1, s->source[2][C0], s->source[2][C1] + s->source[2][C2] };
225 double b[3] = { s->target[0][c], s->target[1][c], s->target[2][c] };
226
227 if (gauss_solve(A, b, 3))
228 continue;
229
230 s->icoeff[0][c] = b[didx[c][0]];
231 s->icoeff[1][c] = b[didx[c][1]];
232 s->icoeff[2][c] = b[didx[c][2]];
233 s->icoeff[3][c] = b[didx[c][3]];
234 }
235 break;
236 case 4:
237 {
238 double A[4 * 4] = { 1, s->source[0][0], s->source[0][1], s->source[0][2],
239 1, s->source[1][0], s->source[1][1], s->source[1][2],
240 1, s->source[2][0], s->source[2][1], s->source[2][2],
241 1, s->source[3][0], s->source[3][1], s->source[3][2] };
242 double b[4] = { s->target[0][c], s->target[1][c], s->target[2][c], s->target[3][c] };
243 int pivot[4];
244
245 if (!gauss_make_triangular(A, pivot, 4))
246 continue;
247 gauss_solve_triangular(A, pivot, b, 4);
248
249 s->icoeff[0][c] = b[0];
250 s->icoeff[1][c] = b[1];
251 s->icoeff[2][c] = b[2];
252 s->icoeff[3][c] = b[3];
253 }
254 break;
255 default:
256 {
257 const int N = s->nb_maps;
258 const int N4 = N + 4;
259 double *A = s->A;
260 double *b = s->b;
261 int *pivot = s->pivot;
262
263 for (int j = 0; j < N; j++)
264 for (int i = j; i < N; i++)
265 A[j*N4+i] = A[i*N4+j] = s->kernel(s->source[i], s->source[j]);
266
267 for (int i = 0; i < N; i++)
268 A[i*N4+N+0] = A[(N+0)*N4+i] = 1;
269 for (int i = 0; i < N; i++)
270 A[i*N4+N+1] = A[(N+1)*N4+i] = s->source[i][0];
271 for (int i = 0; i < N; i++)
272 A[i*N4+N+2] = A[(N+2)*N4+i] = s->source[i][1];
273 for (int i = 0; i < N; i++)
274 A[i*N4+N+3] = A[(N+3)*N4+i] = s->source[i][2];
275
276 for (int j = N; j < N4; j++)
277 for (int i = N;i < N4; i++)
278 A[j * N4 + i] = 0.;
279
280 if (gauss_make_triangular(A, pivot, N4)) {
281 for (int i = 0; i < N; i++)
282 b[i] = s->target[i][c];
283 for (int i = N; i < N + 4; i++)
284 b[i] = 0;
285
286 gauss_solve_triangular(A, pivot, b, N4);
287
288 for (int i = 0; i < N; i++)
289 s->coeff[i][c] = b[i];
290
291 for (int i = 0; i < 4; i++)
292 s->icoeff[i][c] = b[N + i];
293 }
294 }
295 }
296 }
297 }
298
299 typedef struct ThreadData {
300 AVFrame *in, *out;
301 } ThreadData;
302
colormap_slice(AVFilterContext * ctx,void * arg,int jobnr,int nb_jobs)303 static int colormap_slice(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
304 {
305 ColorMapContext *s = ctx->priv;
306 ThreadData *td = arg;
307 AVFrame *in = td->in;
308 AVFrame *out = td->out;
309 const int maps = s->nb_maps;
310 const int width = out->width;
311 const int height = out->height;
312 const int slice_start = (height * jobnr) / nb_jobs;
313 const int slice_end = (height * (jobnr + 1)) / nb_jobs;
314 const int sr_linesize = in->linesize[2] / 4;
315 const int dr_linesize = out->linesize[2] / 4;
316 const int sg_linesize = in->linesize[0] / 4;
317 const int dg_linesize = out->linesize[0] / 4;
318 const int sb_linesize = in->linesize[1] / 4;
319 const int db_linesize = out->linesize[1] / 4;
320 const float *sr = (float *)in->data[2] + slice_start * sr_linesize;
321 const float *sg = (float *)in->data[0] + slice_start * sg_linesize;
322 const float *sb = (float *)in->data[1] + slice_start * sb_linesize;
323 float *r = (float *)out->data[2] + slice_start * dr_linesize;
324 float *g = (float *)out->data[0] + slice_start * dg_linesize;
325 float *b = (float *)out->data[1] + slice_start * db_linesize;
326 float (*kernel)(const float *x, const float *y) = s->kernel;
327 const float *icoeff[4] = { s->icoeff[0], s->icoeff[1], s->icoeff[2], s->icoeff[3] };
328
329 for (int y = slice_start; y < slice_end; y++) {
330 for (int x = 0; x < width; x++) {
331 const float input[3] = { sr[x], sg[x], sb[x] };
332 float srv, sgv, sbv;
333 float rv, gv, bv;
334
335 srv = sr[x];
336 sgv = sg[x];
337 sbv = sb[x];
338
339 rv = icoeff[0][0];
340 gv = icoeff[0][1];
341 bv = icoeff[0][2];
342
343 rv += icoeff[1][0] * srv + icoeff[2][0] * sgv + icoeff[3][0] * sbv;
344 gv += icoeff[1][1] * srv + icoeff[2][1] * sgv + icoeff[3][1] * sbv;
345 bv += icoeff[1][2] * srv + icoeff[2][2] * sgv + icoeff[3][2] * sbv;
346
347 for (int z = 0; z < maps && maps > 4; z++) {
348 const float *coeff = s->coeff[z];
349 const float cr = coeff[0];
350 const float cg = coeff[1];
351 const float cb = coeff[2];
352 const float f = kernel(input, s->source[z]);
353
354 rv += f * cr;
355 gv += f * cg;
356 bv += f * cb;
357 }
358
359 r[x] = rv;
360 g[x] = gv;
361 b[x] = bv;
362 }
363
364 sg += sg_linesize;
365 g += dg_linesize;
366 sb += sb_linesize;
367 b += db_linesize;
368 sr += sr_linesize;
369 r += dr_linesize;
370 }
371
372 return 0;
373 }
374
import_map(AVFilterLink * inlink,AVFrame * in)375 static int import_map(AVFilterLink *inlink, AVFrame *in)
376 {
377 AVFilterContext *ctx = inlink->dst;
378 ColorMapContext *s = ctx->priv;
379 const int is_target = FF_INLINK_IDX(inlink) > 1;
380 const int pw = s->w;
381 const int pw2 = s->w / 2;
382 const int ph = s->h;
383 const int ph2 = s->h / 2;
384 int changed = 0;
385 int idx;
386
387 for (int plane = 0; plane < 3; plane++) {
388 const int c = plane == 0 ? 1 : plane == 1 ? 2 : 0;
389
390 idx = 0;
391 for (int y = ph2; y < in->height && idx < MAX_SIZE; y += ph) {
392 const float *src = (const float *)(in->data[plane] + y * in->linesize[plane]);
393
394 for (int x = pw2; x < in->width && idx < MAX_SIZE; x += pw) {
395 float value = src[x];
396
397 if (is_target) {
398 if (s->ttarget[idx][c] != value)
399 changed = 1;
400 s->ttarget[idx][c] = value;
401 } else {
402 if (s->source[idx][c] != value)
403 changed = 1;
404 s->source[idx][c] = value;
405 }
406
407 idx++;
408 }
409 }
410 }
411
412 if (changed)
413 s->changed[is_target] = 1;
414 if (!s->size)
415 s->size = FFMIN(idx, MAX_SIZE);
416 if (!is_target)
417 s->nb_maps = FFMIN(idx, s->size);
418
419 return 0;
420 }
421
process_frame(FFFrameSync * fs)422 static int process_frame(FFFrameSync *fs)
423 {
424 AVFilterContext *ctx = fs->parent;
425 ColorMapContext *s = fs->opaque;
426 AVFilterLink *outlink = ctx->outputs[0];
427 AVFrame *in, *out, *source, *target;
428 ThreadData td;
429 int ret;
430
431 switch (s->kernel_type) {
432 case EUCLIDEAN:
433 s->kernel = euclidean_kernel;
434 break;
435 case WEUCLIDEAN:
436 s->kernel = weuclidean_kernel;
437 break;
438 default:
439 return AVERROR_BUG;
440 }
441
442 if ((ret = ff_framesync_get_frame(&s->fs, 0, &in, 1)) < 0 ||
443 (ret = ff_framesync_get_frame(&s->fs, 1, &source, 0)) < 0 ||
444 (ret = ff_framesync_get_frame(&s->fs, 2, &target, 0)) < 0)
445 return ret;
446
447 import_map(ctx->inputs[1], source);
448 import_map(ctx->inputs[2], target);
449
450 if (s->changed[0] || s->changed[1]) {
451 build_map(ctx);
452 s->changed[0] = s->changed[1] = 0;
453 }
454
455 if (!ctx->is_disabled) {
456 if (av_frame_is_writable(in)) {
457 out = in;
458 } else {
459 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
460 if (!out) {
461 av_frame_free(&in);
462 return AVERROR(ENOMEM);
463 }
464 av_frame_copy_props(out, in);
465 }
466
467 td.in = in;
468 td.out = out;
469 ff_filter_execute(ctx, colormap_slice, &td, NULL,
470 FFMIN(in->height, ff_filter_get_nb_threads(ctx)));
471
472 if (out != in)
473 av_frame_free(&in);
474 } else {
475 out = in;
476 }
477
478 out->pts = av_rescale_q(s->fs.pts, s->fs.time_base, outlink->time_base);
479
480 return ff_filter_frame(outlink, out);
481 }
482
config_output(AVFilterLink * outlink)483 static int config_output(AVFilterLink *outlink)
484 {
485 AVFilterContext *ctx = outlink->src;
486 ColorMapContext *s = ctx->priv;
487 AVFilterLink *inlink = ctx->inputs[0];
488 AVFilterLink *source = ctx->inputs[1];
489 AVFilterLink *target = ctx->inputs[2];
490 FFFrameSyncIn *in;
491 int ret;
492
493 outlink->time_base = inlink->time_base;
494 outlink->frame_rate = inlink->frame_rate;
495 outlink->sample_aspect_ratio = inlink->sample_aspect_ratio;
496 outlink->w = inlink->w;
497 outlink->h = inlink->h;
498
499 if ((ret = ff_framesync_init(&s->fs, ctx, 3)) < 0)
500 return ret;
501
502 in = s->fs.in;
503 in[0].time_base = inlink->time_base;
504 in[1].time_base = source->time_base;
505 in[2].time_base = target->time_base;
506 in[0].sync = 1;
507 in[0].before = EXT_STOP;
508 in[0].after = EXT_INFINITY;
509 in[1].sync = 1;
510 in[1].before = EXT_STOP;
511 in[1].after = EXT_INFINITY;
512 in[2].sync = 1;
513 in[2].before = EXT_STOP;
514 in[2].after = EXT_INFINITY;
515 s->fs.opaque = s;
516 s->fs.on_event = process_frame;
517
518 ret = ff_framesync_configure(&s->fs);
519 outlink->time_base = s->fs.time_base;
520
521 return ret;
522 }
523
activate(AVFilterContext * ctx)524 static int activate(AVFilterContext *ctx)
525 {
526 ColorMapContext *s = ctx->priv;
527 return ff_framesync_activate(&s->fs);
528 }
529
uninit(AVFilterContext * ctx)530 static av_cold void uninit(AVFilterContext *ctx)
531 {
532 ColorMapContext *const s = ctx->priv;
533
534 ff_framesync_uninit(&s->fs);
535 }
536
537 static const AVFilterPad inputs[] = {
538 {
539 .name = "default",
540 .type = AVMEDIA_TYPE_VIDEO,
541 },
542 {
543 .name = "source",
544 .type = AVMEDIA_TYPE_VIDEO,
545 },
546 {
547 .name = "target",
548 .type = AVMEDIA_TYPE_VIDEO,
549 },
550 };
551
552 static const AVFilterPad outputs[] = {
553 {
554 .name = "default",
555 .type = AVMEDIA_TYPE_VIDEO,
556 .config_props = config_output,
557 },
558 };
559
560 AVFILTER_DEFINE_CLASS(colormap);
561
562 const AVFilter ff_vf_colormap = {
563 .name = "colormap",
564 .description = NULL_IF_CONFIG_SMALL("Apply custom Color Maps to video stream."),
565 .priv_class = &colormap_class,
566 .priv_size = sizeof(ColorMapContext),
567 .activate = activate,
568 FILTER_INPUTS(inputs),
569 FILTER_OUTPUTS(outputs),
570 FILTER_PIXFMTS(AV_PIX_FMT_GBRPF32, AV_PIX_FMT_GBRAPF32),
571 .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
572 AVFILTER_FLAG_SLICE_THREADS,
573 .process_command = ff_filter_process_command,
574 .uninit = uninit,
575 };
576