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