2 * This file is part of FFmpeg.
4 * FFmpeg is free software; you can redistribute it and/or
5 * modify it under the terms of the GNU Lesser General Public
6 * License as published by the Free Software Foundation; either
7 * version 2.1 of the License, or (at your option) any later version.
9 * FFmpeg is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
12 * Lesser General Public License for more details.
14 * You should have received a copy of the GNU Lesser General Public
15 * License along with FFmpeg; if not, write to the Free Software
16 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
21 * Generate one palette for a whole video stream.
24 #include "libavutil/avassert.h"
25 #include "libavutil/opt.h"
29 /* Reference a color and how much it's used */
35 /* Store a range of colors */
37 uint32_t color; // average color
38 int64_t variance; // overall variance of the box (how much the colors are spread)
39 int start; // index in PaletteGenContext->refs
40 int len; // number of referenced colors
41 int sorted_by; // whether range of colors is sorted by red (0), green (1) or blue (2)
45 struct color_ref *entries;
50 STATS_MODE_ALL_FRAMES,
51 STATS_MODE_DIFF_FRAMES,
56 #define HIST_SIZE (1<<(3*NBITS))
62 int reserve_transparent;
65 AVFrame *prev_frame; // previous frame used for the diff stats_mode
66 struct hist_node histogram[HIST_SIZE]; // histogram/hashtable of the colors
67 struct color_ref **refs; // references of all the colors used in the stream
68 int nb_refs; // number of color references (or number of different colors)
69 struct range_box boxes[256]; // define the segmentation of the colorspace (the final palette)
70 int nb_boxes; // number of boxes (increase will segmenting them)
71 int palette_pushed; // if the palette frame is pushed into the outlink or not
74 #define OFFSET(x) offsetof(PaletteGenContext, x)
75 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
76 static const AVOption palettegen_options[] = {
77 { "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 },
78 { "reserve_transparent", "reserve a palette entry for transparency", OFFSET(reserve_transparent), AV_OPT_TYPE_INT, {.i64=1}, 0, 1, FLAGS },
79 { "stats_mode", "set statistics mode", OFFSET(stats_mode), AV_OPT_TYPE_INT, {.i64=STATS_MODE_ALL_FRAMES}, 0, NB_STATS_MODE, FLAGS, "mode" },
80 { "full", "compute full frame histograms", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_ALL_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
81 { "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" },
85 AVFILTER_DEFINE_CLASS(palettegen);
87 static int query_formats(AVFilterContext *ctx)
89 static const enum AVPixelFormat in_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
90 static const enum AVPixelFormat out_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
91 AVFilterFormats *in = ff_make_format_list(in_fmts);
92 AVFilterFormats *out = ff_make_format_list(out_fmts);
94 return AVERROR(ENOMEM);
95 ff_formats_ref(in, &ctx->inputs[0]->out_formats);
96 ff_formats_ref(out, &ctx->outputs[0]->in_formats);
100 typedef int (*cmp_func)(const void *, const void *);
102 #define DECLARE_CMP_FUNC(name, pos) \
103 static int cmp_##name(const void *pa, const void *pb) \
105 const struct color_ref * const *a = pa; \
106 const struct color_ref * const *b = pb; \
107 return ((*a)->color >> (8 * (2 - (pos))) & 0xff) \
108 - ((*b)->color >> (8 * (2 - (pos))) & 0xff); \
111 DECLARE_CMP_FUNC(r, 0)
112 DECLARE_CMP_FUNC(g, 1)
113 DECLARE_CMP_FUNC(b, 2)
115 static const cmp_func cmp_funcs[] = {cmp_r, cmp_g, cmp_b};
118 * Simple color comparison for sorting the final palette
120 static int cmp_color(const void *a, const void *b)
122 const struct range_box *box1 = a;
123 const struct range_box *box2 = b;
124 return box1->color - box2->color;
127 static av_always_inline int diff(const uint32_t a, const uint32_t b)
129 const uint8_t c1[] = {a >> 16 & 0xff, a >> 8 & 0xff, a & 0xff};
130 const uint8_t c2[] = {b >> 16 & 0xff, b >> 8 & 0xff, b & 0xff};
131 const int dr = c1[0] - c2[0];
132 const int dg = c1[1] - c2[1];
133 const int db = c1[2] - c2[2];
134 return dr*dr + dg*dg + db*db;
138 * Find the next box to split: pick the one with the highest variance
140 static int get_next_box_id_to_split(PaletteGenContext *s)
142 int box_id, i, best_box_id = -1;
143 int64_t max_variance = -1;
145 if (s->nb_boxes == s->max_colors - s->reserve_transparent)
148 for (box_id = 0; box_id < s->nb_boxes; box_id++) {
149 struct range_box *box = &s->boxes[box_id];
151 if (s->boxes[box_id].len >= 2) {
153 if (box->variance == -1) {
154 int64_t variance = 0;
156 for (i = 0; i < box->len; i++) {
157 const struct color_ref *ref = s->refs[box->start + i];
158 variance += diff(ref->color, box->color) * ref->count;
160 box->variance = variance;
162 if (box->variance > max_variance) {
163 best_box_id = box_id;
164 max_variance = box->variance;
174 * Get the 32-bit average color for the range of RGB colors enclosed in the
175 * specified box. Takes into account the weight of each color.
177 static uint32_t get_avg_color(struct color_ref * const *refs,
178 const struct range_box *box)
181 const int n = box->len;
182 uint64_t r = 0, g = 0, b = 0, div = 0;
184 for (i = 0; i < n; i++) {
185 const struct color_ref *ref = refs[box->start + i];
186 r += (ref->color >> 16 & 0xff) * ref->count;
187 g += (ref->color >> 8 & 0xff) * ref->count;
188 b += (ref->color & 0xff) * ref->count;
196 return 0xffU<<24 | r<<16 | g<<8 | b;
200 * Split given box in two at position n. The original box becomes the left part
201 * of the split, and the new index box is the right part.
203 static void split_box(PaletteGenContext *s, struct range_box *box, int n)
205 struct range_box *new_box = &s->boxes[s->nb_boxes++];
206 new_box->start = n + 1;
207 new_box->len = box->start + box->len - new_box->start;
208 new_box->sorted_by = box->sorted_by;
209 box->len -= new_box->len;
211 av_assert0(box->len >= 1);
212 av_assert0(new_box->len >= 1);
214 box->color = get_avg_color(s->refs, box);
215 new_box->color = get_avg_color(s->refs, new_box);
217 new_box->variance = -1;
221 * Write the palette into the output frame.
223 static void write_palette(const PaletteGenContext *s, AVFrame *out)
225 int x, y, box_id = 0;
226 uint32_t *pal = (uint32_t *)out->data[0];
227 const int pal_linesize = out->linesize[0] >> 2;
228 uint32_t last_color = 0;
230 for (y = 0; y < out->height; y++) {
231 for (x = 0; x < out->width; x++) {
232 if (box_id < s->nb_boxes) {
233 pal[x] = s->boxes[box_id++].color;
234 if ((x || y) && pal[x] == last_color)
235 av_log(NULL, AV_LOG_WARNING, "Dupped color: %08X\n", pal[x]);
238 pal[x] = 0xff000000; // pad with black
244 if (s->reserve_transparent) {
245 av_assert0(s->nb_boxes < 256);
246 pal[out->width - pal_linesize - 1] = 0x0000ff00; // add a green transparent color
251 * Crawl the histogram to get all the defined colors, and create a linear list
252 * of them (each color reference entry is a pointer to the value in the
253 * histogram/hash table).
255 static struct color_ref **load_color_refs(const struct hist_node *hist, int nb_refs)
258 struct color_ref **refs = av_malloc_array(nb_refs, sizeof(*refs));
263 for (j = 0; j < HIST_SIZE; j++) {
264 const struct hist_node *node = &hist[j];
266 for (i = 0; i < node->nb_entries; i++)
267 refs[k++] = &node->entries[i];
274 * Main function implementing the Median Cut Algorithm defined by Paul Heckbert
275 * in Color Image Quantization for Frame Buffer Display (1982)
277 static AVFrame *get_palette_frame(AVFilterContext *ctx)
280 PaletteGenContext *s = ctx->priv;
281 AVFilterLink *outlink = ctx->outputs[0];
284 struct range_box *box;
286 /* reference only the used colors from histogram */
287 s->refs = load_color_refs(s->histogram, s->nb_refs);
289 av_log(ctx, AV_LOG_ERROR, "Unable to allocate references for %d different colors\n", s->nb_refs);
293 /* create the palette frame */
294 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
299 /* set first box for 0..nb_refs */
300 box = &s->boxes[box_id];
301 box->len = s->nb_refs;
303 box->color = get_avg_color(s->refs, box);
307 while (box && box->len > 1) {
309 uint64_t median, box_weight = 0;
311 /* compute the box weight (sum all the weights of the colors in the
312 * range) and its boundings */
313 uint8_t min[3] = {0xff, 0xff, 0xff};
314 uint8_t max[3] = {0x00, 0x00, 0x00};
315 for (i = box->start; i < box->start + box->len; i++) {
316 const struct color_ref *ref = s->refs[i];
317 const uint32_t rgb = ref->color;
318 const uint8_t r = rgb >> 16 & 0xff, g = rgb >> 8 & 0xff, b = rgb & 0xff;
319 min[0] = FFMIN(r, min[0]), max[0] = FFMAX(r, max[0]);
320 min[1] = FFMIN(g, min[1]), max[1] = FFMAX(g, max[1]);
321 min[2] = FFMIN(b, min[2]), max[2] = FFMAX(b, max[2]);
322 box_weight += ref->count;
325 /* define the axis to sort by according to the widest range of colors */
326 rr = max[0] - min[0];
327 gr = max[1] - min[1];
328 br = max[2] - min[2];
329 longest = 1; // pick green by default (the color the eye is the most sensitive to)
330 if (br >= rr && br >= gr) longest = 2;
331 if (rr >= gr && rr >= br) longest = 0;
332 if (gr >= rr && gr >= br) longest = 1; // prefer green again
334 av_dlog(ctx, "box #%02X [%6d..%-6d] (%6d) w:%-6"PRIu64" ranges:[%2x %2x %2x] sort by %c (already sorted:%c) ",
335 box_id, box->start, box->start + box->len - 1, box->len, box_weight,
336 rr, gr, br, "rgb"[longest], box->sorted_by == longest ? 'y':'n');
338 /* sort the range by its longest axis if it's not already sorted */
339 if (box->sorted_by != longest) {
340 qsort(&s->refs[box->start], box->len, sizeof(*s->refs), cmp_funcs[longest]);
341 box->sorted_by = longest;
344 /* locate the median where to split */
345 median = (box_weight + 1) >> 1;
347 /* if you have 2 boxes, the maximum is actually #0: you must have at
348 * least 1 color on each side of the split, hence the -2 */
349 for (i = box->start; i < box->start + box->len - 2; i++) {
350 box_weight += s->refs[i]->count;
351 if (box_weight > median)
354 av_dlog(ctx, "split @ i=%-6d with w=%-6"PRIu64" (target=%6"PRIu64")\n", i, box_weight, median);
355 split_box(s, box, i);
357 box_id = get_next_box_id_to_split(s);
358 box = box_id >= 0 ? &s->boxes[box_id] : NULL;
361 av_log(ctx, AV_LOG_DEBUG, "%d%s boxes generated out of %d colors\n",
362 s->nb_boxes, s->reserve_transparent ? "(+1)" : "", s->nb_refs);
364 qsort(s->boxes, s->nb_boxes, sizeof(*s->boxes), cmp_color);
366 write_palette(s, out);
372 * Hashing function for the color.
373 * It keeps the NBITS least significant bit of each component to make it
374 * "random" even if the scene doesn't have much different colors.
376 static inline unsigned color_hash(uint32_t color)
378 const uint8_t r = color >> 16 & ((1<<NBITS)-1);
379 const uint8_t g = color >> 8 & ((1<<NBITS)-1);
380 const uint8_t b = color & ((1<<NBITS)-1);
381 return r<<(NBITS*2) | g<<NBITS | b;
385 * Locate the color in the hash table and increment its counter.
387 static int color_inc(struct hist_node *hist, uint32_t color)
390 const unsigned hash = color_hash(color);
391 struct hist_node *node = &hist[hash];
394 for (i = 0; i < node->nb_entries; i++) {
395 e = &node->entries[i];
396 if (e->color == color) {
402 e = av_dynarray2_add((void**)&node->entries, &node->nb_entries,
403 sizeof(*node->entries), NULL);
405 return AVERROR(ENOMEM);
412 * Update histogram when pixels differ from previous frame.
414 static int update_histogram_diff(struct hist_node *hist,
415 const AVFrame *f1, const AVFrame *f2)
417 int x, y, ret, nb_diff_colors = 0;
419 for (y = 0; y < f1->height; y++) {
420 const uint32_t *p = (const uint32_t *)(f1->data[0] + y*f1->linesize[0]);
421 const uint32_t *q = (const uint32_t *)(f2->data[0] + y*f2->linesize[0]);
423 for (x = 0; x < f2->width; x++) {
426 ret = color_inc(hist, p[x]);
429 nb_diff_colors += ret;
432 return nb_diff_colors;
436 * Simple histogram of the frame.
438 static int update_histogram_frame(struct hist_node *hist, const AVFrame *f)
440 int x, y, ret, nb_diff_colors = 0;
442 for (y = 0; y < f->height; y++) {
443 const uint32_t *p = (const uint32_t *)(f->data[0] + y*f->linesize[0]);
445 for (x = 0; x < f->width; x++) {
446 ret = color_inc(hist, p[x]);
449 nb_diff_colors += ret;
452 return nb_diff_colors;
456 * Update the histogram for each passing frame. No frame will be pushed here.
458 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
460 AVFilterContext *ctx = inlink->dst;
461 PaletteGenContext *s = ctx->priv;
462 const int ret = s->prev_frame ? update_histogram_diff(s->histogram, s->prev_frame, in)
463 : update_histogram_frame(s->histogram, in);
468 if (s->stats_mode == STATS_MODE_DIFF_FRAMES) {
469 av_frame_free(&s->prev_frame);
479 * Returns only one frame at the end containing the full palette.
481 static int request_frame(AVFilterLink *outlink)
483 AVFilterContext *ctx = outlink->src;
484 AVFilterLink *inlink = ctx->inputs[0];
485 PaletteGenContext *s = ctx->priv;
488 r = ff_request_frame(inlink);
489 if (r == AVERROR_EOF && !s->palette_pushed) {
490 r = ff_filter_frame(outlink, get_palette_frame(ctx));
491 s->palette_pushed = 1;
498 * The output is one simple 16x16 squared-pixels palette.
500 static int config_output(AVFilterLink *outlink)
502 outlink->w = outlink->h = 16;
503 outlink->sample_aspect_ratio = av_make_q(1, 1);
504 outlink->flags |= FF_LINK_FLAG_REQUEST_LOOP;
508 static av_cold void uninit(AVFilterContext *ctx)
511 PaletteGenContext *s = ctx->priv;
513 for (i = 0; i < HIST_SIZE; i++)
514 av_freep(&s->histogram[i].entries);
516 av_freep(&s->prev_frame);
519 static const AVFilterPad palettegen_inputs[] = {
522 .type = AVMEDIA_TYPE_VIDEO,
523 .filter_frame = filter_frame,
528 static const AVFilterPad palettegen_outputs[] = {
531 .type = AVMEDIA_TYPE_VIDEO,
532 .config_props = config_output,
533 .request_frame = request_frame,
538 AVFilter ff_vf_palettegen = {
539 .name = "palettegen",
540 .description = NULL_IF_CONFIG_SMALL("Find the optimal palette for a given stream."),
541 .priv_size = sizeof(PaletteGenContext),
543 .query_formats = query_formats,
544 .inputs = palettegen_inputs,
545 .outputs = palettegen_outputs,
546 .priv_class = &palettegen_class,