]> git.sesse.net Git - ffmpeg/blob - libavfilter/vf_pseudocolor.c
Merge commit '628ce8b8b6b80cb3985d39e195b71b9d7fad9008'
[ffmpeg] / libavfilter / vf_pseudocolor.c
1 /*
2  * Copyright (c) 2017 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/attributes.h"
22 #include "libavutil/common.h"
23 #include "libavutil/eval.h"
24 #include "libavutil/imgutils.h"
25 #include "libavutil/opt.h"
26 #include "libavutil/pixdesc.h"
27 #include "avfilter.h"
28 #include "formats.h"
29 #include "internal.h"
30 #include "video.h"
31
32 static const char *const var_names[] = {
33     "w",        ///< width of the input video
34     "h",        ///< height of the input video
35     "val",      ///< input value for the pixel
36     "ymin",
37     "umin",
38     "vmin",
39     "amin",
40     "ymax",
41     "umax",
42     "vmax",
43     "amax",
44     NULL
45 };
46
47 enum var_name {
48     VAR_W,
49     VAR_H,
50     VAR_VAL,
51     VAR_YMIN,
52     VAR_UMIN,
53     VAR_VMIN,
54     VAR_AMIN,
55     VAR_YMAX,
56     VAR_UMAX,
57     VAR_VMAX,
58     VAR_AMAX,
59     VAR_VARS_NB
60 };
61
62 typedef struct PseudoColorContext {
63     const AVClass *class;
64     int max;
65     int index;
66     int nb_planes;
67     int color;
68     int linesize[4];
69     int width[4], height[4];
70     double var_values[VAR_VARS_NB];
71     char   *comp_expr_str[4];
72     AVExpr *comp_expr[4];
73     float lut[4][256*256];
74
75     void (*filter[4])(int max, int width, int height,
76                       const uint8_t *index, const uint8_t *src,
77                       uint8_t *dst,
78                       ptrdiff_t ilinesize,
79                       ptrdiff_t slinesize,
80                       ptrdiff_t dlinesize,
81                       float *lut);
82 } PseudoColorContext;
83
84 #define OFFSET(x) offsetof(PseudoColorContext, x)
85 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
86
87 static const AVOption pseudocolor_options[] = {
88     { "c0", "set component #0 expression", OFFSET(comp_expr_str[0]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
89     { "c1", "set component #1 expression", OFFSET(comp_expr_str[1]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
90     { "c2", "set component #2 expression", OFFSET(comp_expr_str[2]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
91     { "c3", "set component #3 expression", OFFSET(comp_expr_str[3]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
92     { "i",  "set component as base",       OFFSET(index),            AV_OPT_TYPE_INT,    {.i64=0}, 0, 3, .flags = FLAGS },
93     { NULL }
94 };
95
96 static const enum AVPixelFormat pix_fmts[] = {
97     AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY16,
98     AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUVA420P,
99     AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVA422P,
100     AV_PIX_FMT_YUV444P, AV_PIX_FMT_GBRP,
101     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_GBRAP,
102     AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUVA422P9,
103     AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUVA420P9,
104     AV_PIX_FMT_YUV444P9, AV_PIX_FMT_YUVA444P9,
105     AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUVA420P10,
106     AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUVA422P10,
107     AV_PIX_FMT_YUV444P10, AV_PIX_FMT_YUVA444P10,
108     AV_PIX_FMT_YUV420P12,
109     AV_PIX_FMT_YUV422P12,
110     AV_PIX_FMT_YUV444P12,
111     AV_PIX_FMT_YUV420P14,
112     AV_PIX_FMT_YUV422P14,
113     AV_PIX_FMT_YUV444P14,
114     AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUVA420P16,
115     AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUVA422P16,
116     AV_PIX_FMT_YUV444P16, AV_PIX_FMT_YUVA444P16,
117     AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRAP10,
118     AV_PIX_FMT_GBRP16, AV_PIX_FMT_GBRAP16,
119     AV_PIX_FMT_NONE
120 };
121
122 static int query_formats(AVFilterContext *ctx)
123 {
124     AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
125     if (!fmts_list)
126         return AVERROR(ENOMEM);
127     return ff_set_common_formats(ctx, fmts_list);
128 }
129
130 static void pseudocolor_filter(int max, int width, int height,
131                                const uint8_t *index,
132                                const uint8_t *src,
133                                uint8_t *dst,
134                                ptrdiff_t ilinesize,
135                                ptrdiff_t slinesize,
136                                ptrdiff_t dlinesize,
137                                float *lut)
138 {
139     int x, y;
140
141     for (y = 0; y < height; y++) {
142         for (x = 0; x < width; x++) {
143             int v = lut[index[x]];
144
145             if (v >= 0 && v <= max) {
146                 dst[x] = v;
147             } else {
148                 dst[x] = src[x];
149             }
150         }
151         index += ilinesize;
152         src += slinesize;
153         dst += dlinesize;
154     }
155 }
156
157 static void pseudocolor_filter_11(int max, int width, int height,
158                                   const uint8_t *index,
159                                   const uint8_t *src,
160                                   uint8_t *dst,
161                                   ptrdiff_t ilinesize,
162                                   ptrdiff_t slinesize,
163                                   ptrdiff_t dlinesize,
164                                   float *lut)
165 {
166     int x, y;
167
168     for (y = 0; y < height; y++) {
169         for (x = 0; x < width; x++) {
170             int v = lut[index[(y << 1) * ilinesize + (x << 1)]];
171
172             if (v >= 0 && v <= max) {
173                 dst[x] = v;
174             } else {
175                 dst[x] = src[x];
176             }
177         }
178         src += slinesize;
179         dst += dlinesize;
180     }
181 }
182
183 static void pseudocolor_filter_11d(int max, int width, int height,
184                                    const uint8_t *index,
185                                    const uint8_t *src,
186                                    uint8_t *dst,
187                                    ptrdiff_t ilinesize,
188                                    ptrdiff_t slinesize,
189                                    ptrdiff_t dlinesize,
190                                    float *lut)
191 {
192     int x, y;
193
194     for (y = 0; y < height; y++) {
195         for (x = 0; x < width; x++) {
196             int v = lut[index[(y >> 1) * ilinesize + (x >> 1)]];
197
198             if (v >= 0 && v <= max) {
199                 dst[x] = v;
200             } else {
201                 dst[x] = src[x];
202             }
203         }
204         src += slinesize;
205         dst += dlinesize;
206     }
207 }
208
209 static void pseudocolor_filter_10(int max, int width, int height,
210                                   const uint8_t *index,
211                                   const uint8_t *src,
212                                   uint8_t *dst,
213                                   ptrdiff_t ilinesize,
214                                   ptrdiff_t slinesize,
215                                   ptrdiff_t dlinesize,
216                                   float *lut)
217 {
218     int x, y;
219
220     for (y = 0; y < height; y++) {
221         for (x = 0; x < width; x++) {
222             int v = lut[index[x << 1]];
223
224             if (v >= 0 && v <= max) {
225                 dst[x] = v;
226             } else {
227                 dst[x] = src[x];
228             }
229         }
230         index += ilinesize;
231         src += slinesize;
232         dst += dlinesize;
233     }
234 }
235
236 static void pseudocolor_filter_10d(int max, int width, int height,
237                                    const uint8_t *index,
238                                    const uint8_t *src,
239                                    uint8_t *dst,
240                                    ptrdiff_t ilinesize,
241                                    ptrdiff_t slinesize,
242                                    ptrdiff_t dlinesize,
243                                    float *lut)
244 {
245     int x, y;
246
247     for (y = 0; y < height; y++) {
248         for (x = 0; x < width; x++) {
249             int v = lut[index[x >> 1]];
250
251             if (v >= 0 && v <= max) {
252                 dst[x] = v;
253             } else {
254                 dst[x] = src[x];
255             }
256         }
257         index += ilinesize;
258         src += slinesize;
259         dst += dlinesize;
260     }
261 }
262
263 static void pseudocolor_filter_16(int max, int width, int height,
264                                   const uint8_t *iindex,
265                                   const uint8_t *ssrc,
266                                   uint8_t *ddst,
267                                   ptrdiff_t ilinesize,
268                                   ptrdiff_t slinesize,
269                                   ptrdiff_t dlinesize,
270                                   float *lut)
271 {
272     const uint16_t *index = (const uint16_t *)iindex;
273     const uint16_t *src = (const uint16_t *)ssrc;
274     uint16_t *dst = (uint16_t *)ddst;
275     int x, y;
276
277     for (y = 0; y < height; y++) {
278         for (x = 0; x < width; x++) {
279             int v = lut[index[x]];
280
281             if (v >= 0 && v <= max) {
282                 dst[x] = v;
283             } else {
284                 dst[x] = src[x];
285             }
286         }
287         index += ilinesize / 2;
288         src += slinesize / 2;
289         dst += dlinesize / 2;
290     }
291 }
292
293 static void pseudocolor_filter_16_10(int max, int width, int height,
294                                      const uint8_t *iindex,
295                                      const uint8_t *ssrc,
296                                      uint8_t *ddst,
297                                      ptrdiff_t ilinesize,
298                                      ptrdiff_t slinesize,
299                                      ptrdiff_t dlinesize,
300                                      float *lut)
301 {
302     const uint16_t *index = (const uint16_t *)iindex;
303     const uint16_t *src = (const uint16_t *)ssrc;
304     uint16_t *dst = (uint16_t *)ddst;
305     int x, y;
306
307     for (y = 0; y < height; y++) {
308         for (x = 0; x < width; x++) {
309             int v = lut[index[x << 1]];
310
311             if (v >= 0 && v <= max) {
312                 dst[x] = v;
313             } else {
314                 dst[x] = src[x];
315             }
316         }
317         index += ilinesize / 2;
318         src += slinesize / 2;
319         dst += dlinesize / 2;
320     }
321 }
322
323 static void pseudocolor_filter_16_10d(int max, int width, int height,
324                                       const uint8_t *iindex,
325                                       const uint8_t *ssrc,
326                                       uint8_t *ddst,
327                                       ptrdiff_t ilinesize,
328                                       ptrdiff_t slinesize,
329                                       ptrdiff_t dlinesize,
330                                       float *lut)
331 {
332     const uint16_t *index = (const uint16_t *)iindex;
333     const uint16_t *src = (const uint16_t *)ssrc;
334     uint16_t *dst = (uint16_t *)ddst;
335     int x, y;
336
337     for (y = 0; y < height; y++) {
338         for (x = 0; x < width; x++) {
339             int v = lut[index[x >> 1]];
340
341             if (v >= 0 && v <= max) {
342                 dst[x] = v;
343             } else {
344                 dst[x] = src[x];
345             }
346         }
347         index += ilinesize / 2;
348         src += slinesize / 2;
349         dst += dlinesize / 2;
350     }
351 }
352
353 static void pseudocolor_filter_16_11(int max, int width, int height,
354                                      const uint8_t *iindex,
355                                      const uint8_t *ssrc,
356                                      uint8_t *ddst,
357                                      ptrdiff_t ilinesize,
358                                      ptrdiff_t slinesize,
359                                      ptrdiff_t dlinesize,
360                                      float *lut)
361 {
362     const uint16_t *index = (const uint16_t *)iindex;
363     const uint16_t *src = (const uint16_t *)ssrc;
364     uint16_t *dst = (uint16_t *)ddst;
365     int x, y;
366
367     ilinesize /= 2;
368     dlinesize /= 2;
369     slinesize /= 2;
370
371     for (y = 0; y < height; y++) {
372         for (x = 0; x < width; x++) {
373             int v = lut[index[(y << 1) * ilinesize + (x << 1)]];
374
375             if (v >= 0 && v <= max) {
376                 dst[x] = v;
377             } else {
378                 dst[x] = src[x];
379             }
380         }
381         src += slinesize;
382         dst += dlinesize;
383     }
384 }
385
386 static void pseudocolor_filter_16_11d(int max, int width, int height,
387                                       const uint8_t *iindex,
388                                       const uint8_t *ssrc,
389                                       uint8_t *ddst,
390                                       ptrdiff_t ilinesize,
391                                       ptrdiff_t slinesize,
392                                       ptrdiff_t dlinesize,
393                                       float *lut)
394 {
395     const uint16_t *index = (const uint16_t *)iindex;
396     const uint16_t *src = (const uint16_t *)ssrc;
397     uint16_t *dst = (uint16_t *)ddst;
398     int x, y;
399
400     ilinesize /= 2;
401     dlinesize /= 2;
402     slinesize /= 2;
403
404     for (y = 0; y < height; y++) {
405         for (x = 0; x < width; x++) {
406             int v = lut[index[(y >> 1) * ilinesize + (x >> 1)]];
407
408             if (v >= 0 && v <= max) {
409                 dst[x] = v;
410             } else {
411                 dst[x] = src[x];
412             }
413         }
414         src += slinesize;
415         dst += dlinesize;
416     }
417 }
418
419 static int config_input(AVFilterLink *inlink)
420 {
421     AVFilterContext *ctx = inlink->dst;
422     PseudoColorContext *s = ctx->priv;
423     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
424     int depth, ret, hsub, vsub, color;
425
426     depth = desc->comp[0].depth;
427     s->max = (1 << depth) - 1;
428     s->nb_planes = av_pix_fmt_count_planes(inlink->format);
429
430     if (s->index >= s->nb_planes) {
431         av_log(ctx, AV_LOG_ERROR, "index out of allowed range\n");
432         return AVERROR(EINVAL);
433     }
434
435     if ((ret = av_image_fill_linesizes(s->linesize, inlink->format, inlink->w)) < 0)
436         return ret;
437
438     hsub = desc->log2_chroma_w;
439     vsub = desc->log2_chroma_h;
440     s->height[1] = s->height[2] = AV_CEIL_RSHIFT(inlink->h, vsub);
441     s->height[0] = s->height[3] = inlink->h;
442     s->width[1]  = s->width[2]  = AV_CEIL_RSHIFT(inlink->w, hsub);
443     s->width[0]  = s->width[3]  = inlink->w;
444
445     s->var_values[VAR_W] = inlink->w;
446     s->var_values[VAR_H] = inlink->h;
447
448     s->var_values[VAR_YMIN] = 16 * (1 << (depth - 8));
449     s->var_values[VAR_UMIN] = 16 * (1 << (depth - 8));
450     s->var_values[VAR_VMIN] = 16 * (1 << (depth - 8));
451     s->var_values[VAR_AMIN] = 0;
452     s->var_values[VAR_YMAX] = 235 * (1 << (depth - 8));
453     s->var_values[VAR_UMAX] = 240 * (1 << (depth - 8));
454     s->var_values[VAR_VMAX] = 240 * (1 << (depth - 8));
455     s->var_values[VAR_AMAX] = s->max;
456
457     for (color = 0; color < s->nb_planes; color++) {
458         double res;
459         int val;
460
461         /* create the parsed expression */
462         av_expr_free(s->comp_expr[color]);
463         s->comp_expr[color] = NULL;
464         ret = av_expr_parse(&s->comp_expr[color], s->comp_expr_str[color],
465                             var_names, NULL, NULL, NULL, NULL, 0, ctx);
466         if (ret < 0) {
467             av_log(ctx, AV_LOG_ERROR,
468                    "Error when parsing the expression '%s' for the component %d and color %d.\n",
469                    s->comp_expr_str[color], color, color);
470             return AVERROR(EINVAL);
471         }
472
473         /* compute the lut */
474         for (val = 0; val < FF_ARRAY_ELEMS(s->lut[color]); val++) {
475             s->var_values[VAR_VAL] = val;
476
477             res = av_expr_eval(s->comp_expr[color], s->var_values, s);
478             if (isnan(res)) {
479                 av_log(ctx, AV_LOG_ERROR,
480                        "Error when evaluating the expression '%s' for the value %d for the component %d.\n",
481                        s->comp_expr_str[color], val, color);
482                 return AVERROR(EINVAL);
483             }
484             s->lut[color][val] = res;
485         }
486     }
487
488     switch (inlink->format) {
489     case AV_PIX_FMT_YUV444P:
490     case AV_PIX_FMT_YUVA444P:
491     case AV_PIX_FMT_GBRP:
492     case AV_PIX_FMT_GBRAP:
493     case AV_PIX_FMT_GRAY8:
494         s->filter[0] = s->filter[1] = s->filter[2] = s->filter[3] = pseudocolor_filter;
495         break;
496     case AV_PIX_FMT_YUV420P:
497     case AV_PIX_FMT_YUVA420P:
498         switch (s->index) {
499         case 0:
500         case 3:
501             s->filter[0] = s->filter[3] = pseudocolor_filter;
502             s->filter[1] = s->filter[2] = pseudocolor_filter_11;
503             break;
504         case 1:
505         case 2:
506             s->filter[0] = s->filter[3] = pseudocolor_filter_11d;
507             s->filter[1] = s->filter[2] = pseudocolor_filter;
508             break;
509         }
510         break;
511     case AV_PIX_FMT_YUV422P:
512     case AV_PIX_FMT_YUVA422P:
513         switch (s->index) {
514         case 0:
515         case 3:
516             s->filter[0] = s->filter[3] = pseudocolor_filter;
517             s->filter[1] = s->filter[2] = pseudocolor_filter_10;
518             break;
519         case 1:
520         case 2:
521             s->filter[0] = s->filter[3] = pseudocolor_filter_10d;
522             s->filter[1] = s->filter[2] = pseudocolor_filter;
523             break;
524         }
525         break;
526     case AV_PIX_FMT_YUV444P9:
527     case AV_PIX_FMT_YUVA444P9:
528     case AV_PIX_FMT_YUV444P10:
529     case AV_PIX_FMT_YUVA444P10:
530     case AV_PIX_FMT_YUV444P12:
531     case AV_PIX_FMT_YUV444P14:
532     case AV_PIX_FMT_YUV444P16:
533     case AV_PIX_FMT_YUVA444P16:
534     case AV_PIX_FMT_GBRP10:
535     case AV_PIX_FMT_GBRAP10:
536     case AV_PIX_FMT_GBRP16:
537     case AV_PIX_FMT_GBRAP16:
538     case AV_PIX_FMT_GRAY10:
539     case AV_PIX_FMT_GRAY16:
540         s->filter[0] = s->filter[1] = s->filter[2] = s->filter[3] = pseudocolor_filter_16;
541         break;
542     case AV_PIX_FMT_YUV422P9:
543     case AV_PIX_FMT_YUVA422P9:
544     case AV_PIX_FMT_YUV422P10:
545     case AV_PIX_FMT_YUVA422P10:
546     case AV_PIX_FMT_YUV422P12:
547     case AV_PIX_FMT_YUV422P14:
548     case AV_PIX_FMT_YUV422P16:
549     case AV_PIX_FMT_YUVA422P16:
550         switch (s->index) {
551         case 0:
552         case 3:
553             s->filter[0] = s->filter[3] = pseudocolor_filter_16;
554             s->filter[1] = s->filter[2] = pseudocolor_filter_16_10;
555             break;
556         case 1:
557         case 2:
558             s->filter[0] = s->filter[3] = pseudocolor_filter_16_10d;
559             s->filter[1] = s->filter[2] = pseudocolor_filter_16;
560             break;
561         }
562         break;
563     case AV_PIX_FMT_YUV420P9:
564     case AV_PIX_FMT_YUVA420P9:
565     case AV_PIX_FMT_YUV420P10:
566     case AV_PIX_FMT_YUVA420P10:
567     case AV_PIX_FMT_YUV420P12:
568     case AV_PIX_FMT_YUV420P14:
569     case AV_PIX_FMT_YUV420P16:
570     case AV_PIX_FMT_YUVA420P16:
571         switch (s->index) {
572         case 0:
573         case 3:
574             s->filter[0] = s->filter[3] = pseudocolor_filter_16;
575             s->filter[1] = s->filter[2] = pseudocolor_filter_16_11;
576             break;
577         case 1:
578         case 2:
579             s->filter[0] = s->filter[3] = pseudocolor_filter_16_11d;
580             s->filter[1] = s->filter[2] = pseudocolor_filter_16;
581             break;
582         }
583         break;
584     }
585
586     return 0;
587 }
588
589 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
590 {
591     AVFilterContext *ctx = inlink->dst;
592     PseudoColorContext *s = ctx->priv;
593     AVFilterLink *outlink = ctx->outputs[0];
594     AVFrame *out;
595     int plane;
596
597     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
598     if (!out) {
599         av_frame_free(&in);
600         return AVERROR(ENOMEM);
601     }
602     av_frame_copy_props(out, in);
603
604     for (plane = 0; plane < s->nb_planes; plane++) {
605         const uint8_t *index = in->data[s->index];
606         const uint8_t *src = in->data[plane];
607         uint8_t *dst = out->data[plane];
608         ptrdiff_t ilinesize = in->linesize[s->index];
609         ptrdiff_t slinesize = in->linesize[plane];
610         ptrdiff_t dlinesize = out->linesize[plane];
611
612         s->filter[plane](s->max, s->width[plane], s->height[plane],
613                          index, src, dst, ilinesize, slinesize,
614                          dlinesize, s->lut[plane]);
615     }
616
617     av_frame_free(&in);
618     return ff_filter_frame(outlink, out);
619 }
620
621 static const AVFilterPad inputs[] = {
622     {
623         .name         = "default",
624         .type         = AVMEDIA_TYPE_VIDEO,
625         .filter_frame = filter_frame,
626         .config_props = config_input,
627     },
628     { NULL }
629 };
630
631 static const AVFilterPad outputs[] = {
632     {
633         .name = "default",
634         .type = AVMEDIA_TYPE_VIDEO,
635     },
636     { NULL }
637 };
638
639 static av_cold void uninit(AVFilterContext *ctx)
640 {
641     PseudoColorContext *s = ctx->priv;
642     int i;
643
644     for (i = 0; i < 4; i++) {
645         av_expr_free(s->comp_expr[i]);
646         s->comp_expr[i] = NULL;
647     }
648 }
649
650 AVFILTER_DEFINE_CLASS(pseudocolor);
651
652 AVFilter ff_vf_pseudocolor = {
653     .name          = "pseudocolor",
654     .description   = NULL_IF_CONFIG_SMALL("Make pseudocolored video frames."),
655     .priv_size     = sizeof(PseudoColorContext),
656     .priv_class    = &pseudocolor_class,
657     .uninit        = uninit,
658     .query_formats = query_formats,
659     .inputs        = inputs,
660     .outputs       = outputs,
661     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
662 };