lavfi/tests/filtfmts: fix the build warning.
[ffmpeg.git] / libavfilter / vf_palettegen.c
1 /*
2  * Copyright (c) 2015 Stupeflix
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  * Generate one palette for a whole video stream.
24  */
25
26 #include "libavutil/avassert.h"
27 #include "libavutil/internal.h"
28 #include "libavutil/opt.h"
29 #include "libavutil/qsort.h"
30 #include "libavutil/intreadwrite.h"
31 #include "avfilter.h"
32 #include "internal.h"
33
34 /* Reference a color and how much it's used */
35 struct color_ref {
36     uint32_t color;
37     uint64_t count;
38 };
39
40 /* Store a range of colors */
41 struct range_box {
42     uint32_t color;     // average color
43     int64_t variance;   // overall variance of the box (how much the colors are spread)
44     int start;          // index in PaletteGenContext->refs
45     int len;            // number of referenced colors
46     int sorted_by;      // whether range of colors is sorted by red (0), green (1) or blue (2)
47 };
48
49 struct hist_node {
50     struct color_ref *entries;
51     int nb_entries;
52 };
53
54 enum {
55     STATS_MODE_ALL_FRAMES,
56     STATS_MODE_DIFF_FRAMES,
57     STATS_MODE_SINGLE_FRAMES,
58     NB_STATS_MODE
59 };
60
61 #define NBITS 5
62 #define HIST_SIZE (1<<(3*NBITS))
63
64 typedef struct PaletteGenContext {
65     const AVClass *class;
66
67     int max_colors;
68     int reserve_transparent;
69     int stats_mode;
70
71     AVFrame *prev_frame;                    // previous frame used for the diff stats_mode
72     struct hist_node histogram[HIST_SIZE];  // histogram/hashtable of the colors
73     struct color_ref **refs;                // references of all the colors used in the stream
74     int nb_refs;                            // number of color references (or number of different colors)
75     struct range_box boxes[256];            // define the segmentation of the colorspace (the final palette)
76     int nb_boxes;                           // number of boxes (increase will segmenting them)
77     int palette_pushed;                     // if the palette frame is pushed into the outlink or not
78     uint8_t transparency_color[4];          // background color for transparency
79 } PaletteGenContext;
80
81 #define OFFSET(x) offsetof(PaletteGenContext, x)
82 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
83 static const AVOption palettegen_options[] = {
84     { "max_colors", "set the maximum number of colors to use in the palette", OFFSET(max_colors), AV_OPT_TYPE_INT, {.i64=256}, 4, 256, FLAGS },
85     { "reserve_transparent", "reserve a palette entry for transparency", OFFSET(reserve_transparent), AV_OPT_TYPE_BOOL, {.i64=1}, 0, 1, FLAGS },
86     { "transparency_color", "set a background color for transparency", OFFSET(transparency_color), AV_OPT_TYPE_COLOR, {.str="lime"}, CHAR_MIN, CHAR_MAX, FLAGS },
87     { "stats_mode", "set statistics mode", OFFSET(stats_mode), AV_OPT_TYPE_INT, {.i64=STATS_MODE_ALL_FRAMES}, 0, NB_STATS_MODE-1, FLAGS, "mode" },
88         { "full", "compute full frame histograms", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_ALL_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
89         { "diff", "compute histograms only for the part that differs from previous frame", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_DIFF_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
90         { "single", "compute new histogram for each frame", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_SINGLE_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
91     { NULL }
92 };
93
94 AVFILTER_DEFINE_CLASS(palettegen);
95
96 static int query_formats(AVFilterContext *ctx)
97 {
98     static const enum AVPixelFormat in_fmts[]  = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
99     static const enum AVPixelFormat out_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
100     int ret;
101
102     if ((ret = ff_formats_ref(ff_make_format_list(in_fmts) , &ctx->inputs[0]->out_formats)) < 0)
103         return ret;
104     if ((ret = ff_formats_ref(ff_make_format_list(out_fmts), &ctx->outputs[0]->in_formats)) < 0)
105         return ret;
106     return 0;
107 }
108
109 typedef int (*cmp_func)(const void *, const void *);
110
111 #define DECLARE_CMP_FUNC(name, pos)                     \
112 static int cmp_##name(const void *pa, const void *pb)   \
113 {                                                       \
114     const struct color_ref * const *a = pa;             \
115     const struct color_ref * const *b = pb;             \
116     return   ((*a)->color >> (8 * (2 - (pos))) & 0xff)  \
117            - ((*b)->color >> (8 * (2 - (pos))) & 0xff); \
118 }
119
120 DECLARE_CMP_FUNC(r, 0)
121 DECLARE_CMP_FUNC(g, 1)
122 DECLARE_CMP_FUNC(b, 2)
123
124 static const cmp_func cmp_funcs[] = {cmp_r, cmp_g, cmp_b};
125
126 /**
127  * Simple color comparison for sorting the final palette
128  */
129 static int cmp_color(const void *a, const void *b)
130 {
131     const struct range_box *box1 = a;
132     const struct range_box *box2 = b;
133     return FFDIFFSIGN(box1->color , box2->color);
134 }
135
136 static av_always_inline int diff(const uint32_t a, const uint32_t b)
137 {
138     const uint8_t c1[] = {a >> 16 & 0xff, a >> 8 & 0xff, a & 0xff};
139     const uint8_t c2[] = {b >> 16 & 0xff, b >> 8 & 0xff, b & 0xff};
140     const int dr = c1[0] - c2[0];
141     const int dg = c1[1] - c2[1];
142     const int db = c1[2] - c2[2];
143     return dr*dr + dg*dg + db*db;
144 }
145
146 /**
147  * Find the next box to split: pick the one with the highest variance
148  */
149 static int get_next_box_id_to_split(PaletteGenContext *s)
150 {
151     int box_id, i, best_box_id = -1;
152     int64_t max_variance = -1;
153
154     if (s->nb_boxes == s->max_colors - s->reserve_transparent)
155         return -1;
156
157     for (box_id = 0; box_id < s->nb_boxes; box_id++) {
158         struct range_box *box = &s->boxes[box_id];
159
160         if (s->boxes[box_id].len >= 2) {
161
162             if (box->variance == -1) {
163                 int64_t variance = 0;
164
165                 for (i = 0; i < box->len; i++) {
166                     const struct color_ref *ref = s->refs[box->start + i];
167                     variance += diff(ref->color, box->color) * ref->count;
168                 }
169                 box->variance = variance;
170             }
171             if (box->variance > max_variance) {
172                 best_box_id = box_id;
173                 max_variance = box->variance;
174             }
175         } else {
176             box->variance = -1;
177         }
178     }
179     return best_box_id;
180 }
181
182 /**
183  * Get the 32-bit average color for the range of RGB colors enclosed in the
184  * specified box. Takes into account the weight of each color.
185  */
186 static uint32_t get_avg_color(struct color_ref * const *refs,
187                               const struct range_box *box)
188 {
189     int i;
190     const int n = box->len;
191     uint64_t r = 0, g = 0, b = 0, div = 0;
192
193     for (i = 0; i < n; i++) {
194         const struct color_ref *ref = refs[box->start + i];
195         r += (ref->color >> 16 & 0xff) * ref->count;
196         g += (ref->color >>  8 & 0xff) * ref->count;
197         b += (ref->color       & 0xff) * ref->count;
198         div += ref->count;
199     }
200
201     r = r / div;
202     g = g / div;
203     b = b / div;
204
205     return 0xffU<<24 | r<<16 | g<<8 | b;
206 }
207
208 /**
209  * Split given box in two at position n. The original box becomes the left part
210  * of the split, and the new index box is the right part.
211  */
212 static void split_box(PaletteGenContext *s, struct range_box *box, int n)
213 {
214     struct range_box *new_box = &s->boxes[s->nb_boxes++];
215     new_box->start     = n + 1;
216     new_box->len       = box->start + box->len - new_box->start;
217     new_box->sorted_by = box->sorted_by;
218     box->len -= new_box->len;
219
220     av_assert0(box->len     >= 1);
221     av_assert0(new_box->len >= 1);
222
223     box->color     = get_avg_color(s->refs, box);
224     new_box->color = get_avg_color(s->refs, new_box);
225     box->variance     = -1;
226     new_box->variance = -1;
227 }
228
229 /**
230  * Write the palette into the output frame.
231  */
232 static void write_palette(AVFilterContext *ctx, AVFrame *out)
233 {
234     const PaletteGenContext *s = ctx->priv;
235     int x, y, box_id = 0;
236     uint32_t *pal = (uint32_t *)out->data[0];
237     const int pal_linesize = out->linesize[0] >> 2;
238     uint32_t last_color = 0;
239
240     for (y = 0; y < out->height; y++) {
241         for (x = 0; x < out->width; x++) {
242             if (box_id < s->nb_boxes) {
243                 pal[x] = s->boxes[box_id++].color;
244                 if ((x || y) && pal[x] == last_color)
245                     av_log(ctx, AV_LOG_WARNING, "Dupped color: %08"PRIX32"\n", pal[x]);
246                 last_color = pal[x];
247             } else {
248                 pal[x] = 0xff000000; // pad with black
249             }
250         }
251         pal += pal_linesize;
252     }
253
254     if (s->reserve_transparent) {
255         av_assert0(s->nb_boxes < 256);
256         pal[out->width - pal_linesize - 1] = AV_RB32(&s->transparency_color) >> 8;
257     }
258 }
259
260 /**
261  * Crawl the histogram to get all the defined colors, and create a linear list
262  * of them (each color reference entry is a pointer to the value in the
263  * histogram/hash table).
264  */
265 static struct color_ref **load_color_refs(const struct hist_node *hist, int nb_refs)
266 {
267     int i, j, k = 0;
268     struct color_ref **refs = av_malloc_array(nb_refs, sizeof(*refs));
269
270     if (!refs)
271         return NULL;
272
273     for (j = 0; j < HIST_SIZE; j++) {
274         const struct hist_node *node = &hist[j];
275
276         for (i = 0; i < node->nb_entries; i++)
277             refs[k++] = &node->entries[i];
278     }
279
280     return refs;
281 }
282
283 static double set_colorquant_ratio_meta(AVFrame *out, int nb_out, int nb_in)
284 {
285     char buf[32];
286     const double ratio = (double)nb_out / nb_in;
287     snprintf(buf, sizeof(buf), "%f", ratio);
288     av_dict_set(&out->metadata, "lavfi.color_quant_ratio", buf, 0);
289     return ratio;
290 }
291
292 /**
293  * Main function implementing the Median Cut Algorithm defined by Paul Heckbert
294  * in Color Image Quantization for Frame Buffer Display (1982)
295  */
296 static AVFrame *get_palette_frame(AVFilterContext *ctx)
297 {
298     AVFrame *out;
299     PaletteGenContext *s = ctx->priv;
300     AVFilterLink *outlink = ctx->outputs[0];
301     double ratio;
302     int box_id = 0;
303     struct range_box *box;
304
305     /* reference only the used colors from histogram */
306     s->refs = load_color_refs(s->histogram, s->nb_refs);
307     if (!s->refs) {
308         av_log(ctx, AV_LOG_ERROR, "Unable to allocate references for %d different colors\n", s->nb_refs);
309         return NULL;
310     }
311
312     /* create the palette frame */
313     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
314     if (!out)
315         return NULL;
316     out->pts = 0;
317
318     /* set first box for 0..nb_refs */
319     box = &s->boxes[box_id];
320     box->len = s->nb_refs;
321     box->sorted_by = -1;
322     box->color = get_avg_color(s->refs, box);
323     box->variance = -1;
324     s->nb_boxes = 1;
325
326     while (box && box->len > 1) {
327         int i, rr, gr, br, longest;
328         uint64_t median, box_weight = 0;
329
330         /* compute the box weight (sum all the weights of the colors in the
331          * range) and its boundings */
332         uint8_t min[3] = {0xff, 0xff, 0xff};
333         uint8_t max[3] = {0x00, 0x00, 0x00};
334         for (i = box->start; i < box->start + box->len; i++) {
335             const struct color_ref *ref = s->refs[i];
336             const uint32_t rgb = ref->color;
337             const uint8_t r = rgb >> 16 & 0xff, g = rgb >> 8 & 0xff, b = rgb & 0xff;
338             min[0] = FFMIN(r, min[0]), max[0] = FFMAX(r, max[0]);
339             min[1] = FFMIN(g, min[1]), max[1] = FFMAX(g, max[1]);
340             min[2] = FFMIN(b, min[2]), max[2] = FFMAX(b, max[2]);
341             box_weight += ref->count;
342         }
343
344         /* define the axis to sort by according to the widest range of colors */
345         rr = max[0] - min[0];
346         gr = max[1] - min[1];
347         br = max[2] - min[2];
348         longest = 1; // pick green by default (the color the eye is the most sensitive to)
349         if (br >= rr && br >= gr) longest = 2;
350         if (rr >= gr && rr >= br) longest = 0;
351         if (gr >= rr && gr >= br) longest = 1; // prefer green again
352
353         ff_dlog(ctx, "box #%02X [%6d..%-6d] (%6d) w:%-6"PRIu64" ranges:[%2x %2x %2x] sort by %c (already sorted:%c) ",
354                 box_id, box->start, box->start + box->len - 1, box->len, box_weight,
355                 rr, gr, br, "rgb"[longest], box->sorted_by == longest ? 'y':'n');
356
357         /* sort the range by its longest axis if it's not already sorted */
358         if (box->sorted_by != longest) {
359             cmp_func cmpf = cmp_funcs[longest];
360             AV_QSORT(&s->refs[box->start], box->len, const struct color_ref *, cmpf);
361             box->sorted_by = longest;
362         }
363
364         /* locate the median where to split */
365         median = (box_weight + 1) >> 1;
366         box_weight = 0;
367         /* if you have 2 boxes, the maximum is actually #0: you must have at
368          * least 1 color on each side of the split, hence the -2 */
369         for (i = box->start; i < box->start + box->len - 2; i++) {
370             box_weight += s->refs[i]->count;
371             if (box_weight > median)
372                 break;
373         }
374         ff_dlog(ctx, "split @ i=%-6d with w=%-6"PRIu64" (target=%6"PRIu64")\n", i, box_weight, median);
375         split_box(s, box, i);
376
377         box_id = get_next_box_id_to_split(s);
378         box = box_id >= 0 ? &s->boxes[box_id] : NULL;
379     }
380
381     ratio = set_colorquant_ratio_meta(out, s->nb_boxes, s->nb_refs);
382     av_log(ctx, AV_LOG_INFO, "%d%s colors generated out of %d colors; ratio=%f\n",
383            s->nb_boxes, s->reserve_transparent ? "(+1)" : "", s->nb_refs, ratio);
384
385     qsort(s->boxes, s->nb_boxes, sizeof(*s->boxes), cmp_color);
386
387     write_palette(ctx, out);
388
389     return out;
390 }
391
392 /**
393  * Hashing function for the color.
394  * It keeps the NBITS least significant bit of each component to make it
395  * "random" even if the scene doesn't have much different colors.
396  */
397 static inline unsigned color_hash(uint32_t color)
398 {
399     const uint8_t r = color >> 16 & ((1<<NBITS)-1);
400     const uint8_t g = color >>  8 & ((1<<NBITS)-1);
401     const uint8_t b = color       & ((1<<NBITS)-1);
402     return r<<(NBITS*2) | g<<NBITS | b;
403 }
404
405 /**
406  * Locate the color in the hash table and increment its counter.
407  */
408 static int color_inc(struct hist_node *hist, uint32_t color)
409 {
410     int i;
411     const unsigned hash = color_hash(color);
412     struct hist_node *node = &hist[hash];
413     struct color_ref *e;
414
415     for (i = 0; i < node->nb_entries; i++) {
416         e = &node->entries[i];
417         if (e->color == color) {
418             e->count++;
419             return 0;
420         }
421     }
422
423     e = av_dynarray2_add((void**)&node->entries, &node->nb_entries,
424                          sizeof(*node->entries), NULL);
425     if (!e)
426         return AVERROR(ENOMEM);
427     e->color = color;
428     e->count = 1;
429     return 1;
430 }
431
432 /**
433  * Update histogram when pixels differ from previous frame.
434  */
435 static int update_histogram_diff(struct hist_node *hist,
436                                  const AVFrame *f1, const AVFrame *f2)
437 {
438     int x, y, ret, nb_diff_colors = 0;
439
440     for (y = 0; y < f1->height; y++) {
441         const uint32_t *p = (const uint32_t *)(f1->data[0] + y*f1->linesize[0]);
442         const uint32_t *q = (const uint32_t *)(f2->data[0] + y*f2->linesize[0]);
443
444         for (x = 0; x < f1->width; x++) {
445             if (p[x] == q[x])
446                 continue;
447             ret = color_inc(hist, p[x]);
448             if (ret < 0)
449                 return ret;
450             nb_diff_colors += ret;
451         }
452     }
453     return nb_diff_colors;
454 }
455
456 /**
457  * Simple histogram of the frame.
458  */
459 static int update_histogram_frame(struct hist_node *hist, const AVFrame *f)
460 {
461     int x, y, ret, nb_diff_colors = 0;
462
463     for (y = 0; y < f->height; y++) {
464         const uint32_t *p = (const uint32_t *)(f->data[0] + y*f->linesize[0]);
465
466         for (x = 0; x < f->width; x++) {
467             ret = color_inc(hist, p[x]);
468             if (ret < 0)
469                 return ret;
470             nb_diff_colors += ret;
471         }
472     }
473     return nb_diff_colors;
474 }
475
476 /**
477  * Update the histogram for each passing frame. No frame will be pushed here.
478  */
479 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
480 {
481     AVFilterContext *ctx = inlink->dst;
482     PaletteGenContext *s = ctx->priv;
483     int ret = s->prev_frame ? update_histogram_diff(s->histogram, s->prev_frame, in)
484                             : update_histogram_frame(s->histogram, in);
485
486     if (ret > 0)
487         s->nb_refs += ret;
488
489     if (s->stats_mode == STATS_MODE_DIFF_FRAMES) {
490         av_frame_free(&s->prev_frame);
491         s->prev_frame = in;
492     } else if (s->stats_mode == STATS_MODE_SINGLE_FRAMES) {
493         AVFrame *out;
494         int i;
495
496         out = get_palette_frame(ctx);
497         out->pts = in->pts;
498         av_frame_free(&in);
499         ret = ff_filter_frame(ctx->outputs[0], out);
500         for (i = 0; i < HIST_SIZE; i++)
501             av_freep(&s->histogram[i].entries);
502         av_freep(&s->refs);
503         s->nb_refs = 0;
504         s->nb_boxes = 0;
505         memset(s->boxes, 0, sizeof(s->boxes));
506         memset(s->histogram, 0, sizeof(s->histogram));
507     } else {
508         av_frame_free(&in);
509     }
510
511     return ret;
512 }
513
514 /**
515  * Returns only one frame at the end containing the full palette.
516  */
517 static int request_frame(AVFilterLink *outlink)
518 {
519     AVFilterContext *ctx = outlink->src;
520     AVFilterLink *inlink = ctx->inputs[0];
521     PaletteGenContext *s = ctx->priv;
522     int r;
523
524     r = ff_request_frame(inlink);
525     if (r == AVERROR_EOF && !s->palette_pushed && s->nb_refs && s->stats_mode != STATS_MODE_SINGLE_FRAMES) {
526         r = ff_filter_frame(outlink, get_palette_frame(ctx));
527         s->palette_pushed = 1;
528         return r;
529     }
530     return r;
531 }
532
533 /**
534  * The output is one simple 16x16 squared-pixels palette.
535  */
536 static int config_output(AVFilterLink *outlink)
537 {
538     outlink->w = outlink->h = 16;
539     outlink->sample_aspect_ratio = av_make_q(1, 1);
540     return 0;
541 }
542
543 static av_cold void uninit(AVFilterContext *ctx)
544 {
545     int i;
546     PaletteGenContext *s = ctx->priv;
547
548     for (i = 0; i < HIST_SIZE; i++)
549         av_freep(&s->histogram[i].entries);
550     av_freep(&s->refs);
551     av_frame_free(&s->prev_frame);
552 }
553
554 static const AVFilterPad palettegen_inputs[] = {
555     {
556         .name         = "default",
557         .type         = AVMEDIA_TYPE_VIDEO,
558         .filter_frame = filter_frame,
559     },
560     { NULL }
561 };
562
563 static const AVFilterPad palettegen_outputs[] = {
564     {
565         .name          = "default",
566         .type          = AVMEDIA_TYPE_VIDEO,
567         .config_props  = config_output,
568         .request_frame = request_frame,
569     },
570     { NULL }
571 };
572
573 AVFilter ff_vf_palettegen = {
574     .name          = "palettegen",
575     .description   = NULL_IF_CONFIG_SMALL("Find the optimal palette for a given stream."),
576     .priv_size     = sizeof(PaletteGenContext),
577     .uninit        = uninit,
578     .query_formats = query_formats,
579     .inputs        = palettegen_inputs,
580     .outputs       = palettegen_outputs,
581     .priv_class    = &palettegen_class,
582 };