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