summaryrefslogtreecommitdiff
path: root/libavfilter/vf_palettegen.c (plain)
blob: b470079ccc8f0945c560925266d96c36878962ce
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 "avfilter.h"
31#include "internal.h"
32
33/* Reference a color and how much it's used */
34struct color_ref {
35 uint32_t color;
36 uint64_t count;
37};
38
39/* Store a range of colors */
40struct range_box {
41 uint32_t color; // average color
42 int64_t variance; // overall variance of the box (how much the colors are spread)
43 int start; // index in PaletteGenContext->refs
44 int len; // number of referenced colors
45 int sorted_by; // whether range of colors is sorted by red (0), green (1) or blue (2)
46};
47
48struct hist_node {
49 struct color_ref *entries;
50 int nb_entries;
51};
52
53enum {
54 STATS_MODE_ALL_FRAMES,
55 STATS_MODE_DIFF_FRAMES,
56 STATS_MODE_SINGLE_FRAMES,
57 NB_STATS_MODE
58};
59
60#define NBITS 5
61#define HIST_SIZE (1<<(3*NBITS))
62
63typedef struct {
64 const AVClass *class;
65
66 int max_colors;
67 int reserve_transparent;
68 int stats_mode;
69
70 AVFrame *prev_frame; // previous frame used for the diff stats_mode
71 struct hist_node histogram[HIST_SIZE]; // histogram/hashtable of the colors
72 struct color_ref **refs; // references of all the colors used in the stream
73 int nb_refs; // number of color references (or number of different colors)
74 struct range_box boxes[256]; // define the segmentation of the colorspace (the final palette)
75 int nb_boxes; // number of boxes (increase will segmenting them)
76 int palette_pushed; // if the palette frame is pushed into the outlink or not
77} PaletteGenContext;
78
79#define OFFSET(x) offsetof(PaletteGenContext, x)
80#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
81static const AVOption palettegen_options[] = {
82 { "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 },
83 { "reserve_transparent", "reserve a palette entry for transparency", OFFSET(reserve_transparent), AV_OPT_TYPE_BOOL, {.i64=1}, 0, 1, FLAGS },
84 { "stats_mode", "set statistics mode", OFFSET(stats_mode), AV_OPT_TYPE_INT, {.i64=STATS_MODE_ALL_FRAMES}, 0, NB_STATS_MODE-1, FLAGS, "mode" },
85 { "full", "compute full frame histograms", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_ALL_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
86 { "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" },
87 { "single", "compute new histogram for each frame", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_SINGLE_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
88 { NULL }
89};
90
91AVFILTER_DEFINE_CLASS(palettegen);
92
93static int query_formats(AVFilterContext *ctx)
94{
95 static const enum AVPixelFormat in_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
96 static const enum AVPixelFormat out_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
97 int ret;
98
99 if ((ret = ff_formats_ref(ff_make_format_list(in_fmts) , &ctx->inputs[0]->out_formats)) < 0)
100 return ret;
101 if ((ret = ff_formats_ref(ff_make_format_list(out_fmts), &ctx->outputs[0]->in_formats)) < 0)
102 return ret;
103 return 0;
104}
105
106typedef int (*cmp_func)(const void *, const void *);
107
108#define DECLARE_CMP_FUNC(name, pos) \
109static 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
117DECLARE_CMP_FUNC(r, 0)
118DECLARE_CMP_FUNC(g, 1)
119DECLARE_CMP_FUNC(b, 2)
120
121static const cmp_func cmp_funcs[] = {cmp_r, cmp_g, cmp_b};
122
123/**
124 * Simple color comparison for sorting the final palette
125 */
126static 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 FFDIFFSIGN(box1->color , box2->color);
131}
132
133static 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 */
146static int get_next_box_id_to_split(PaletteGenContext *s)
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 */
183static 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 */
209static 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 */
229static void write_palette(AVFilterContext *ctx, AVFrame *out)
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: %08"PRIX32"\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 */
262static 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
280static 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 */
293static AVFrame *get_palette_frame(AVFilterContext *ctx)
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 ff_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 ff_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 */
394static 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 */
405static 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 */
432static 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 */
456static 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 */
476static int filter_frame(AVFilterLink *inlink, AVFrame *in)
477{
478 AVFilterContext *ctx = inlink->dst;
479 PaletteGenContext *s = ctx->priv;
480 int ret = s->prev_frame ? update_histogram_diff(s->histogram, s->prev_frame, in)
481 : update_histogram_frame(s->histogram, in);
482
483 if (ret > 0)
484 s->nb_refs += ret;
485
486 if (s->stats_mode == STATS_MODE_DIFF_FRAMES) {
487 av_frame_free(&s->prev_frame);
488 s->prev_frame = in;
489 } else if (s->stats_mode == STATS_MODE_SINGLE_FRAMES) {
490 AVFrame *out;
491 int i;
492
493 out = get_palette_frame(ctx);
494 out->pts = in->pts;
495 av_frame_free(&in);
496 ret = ff_filter_frame(ctx->outputs[0], out);
497 for (i = 0; i < HIST_SIZE; i++)
498 av_freep(&s->histogram[i].entries);
499 av_freep(&s->refs);
500 s->nb_refs = 0;
501 s->nb_boxes = 0;
502 memset(s->boxes, 0, sizeof(s->boxes));
503 memset(s->histogram, 0, sizeof(s->histogram));
504 } else {
505 av_frame_free(&in);
506 }
507
508 return ret;
509}
510
511/**
512 * Returns only one frame at the end containing the full palette.
513 */
514static int request_frame(AVFilterLink *outlink)
515{
516 AVFilterContext *ctx = outlink->src;
517 AVFilterLink *inlink = ctx->inputs[0];
518 PaletteGenContext *s = ctx->priv;
519 int r;
520
521 r = ff_request_frame(inlink);
522 if (r == AVERROR_EOF && !s->palette_pushed && s->nb_refs && s->stats_mode != STATS_MODE_SINGLE_FRAMES) {
523 r = ff_filter_frame(outlink, get_palette_frame(ctx));
524 s->palette_pushed = 1;
525 return r;
526 }
527 return r;
528}
529
530/**
531 * The output is one simple 16x16 squared-pixels palette.
532 */
533static int config_output(AVFilterLink *outlink)
534{
535 outlink->w = outlink->h = 16;
536 outlink->sample_aspect_ratio = av_make_q(1, 1);
537 return 0;
538}
539
540static av_cold void uninit(AVFilterContext *ctx)
541{
542 int i;
543 PaletteGenContext *s = ctx->priv;
544
545 for (i = 0; i < HIST_SIZE; i++)
546 av_freep(&s->histogram[i].entries);
547 av_freep(&s->refs);
548 av_frame_free(&s->prev_frame);
549}
550
551static const AVFilterPad palettegen_inputs[] = {
552 {
553 .name = "default",
554 .type = AVMEDIA_TYPE_VIDEO,
555 .filter_frame = filter_frame,
556 },
557 { NULL }
558};
559
560static const AVFilterPad palettegen_outputs[] = {
561 {
562 .name = "default",
563 .type = AVMEDIA_TYPE_VIDEO,
564 .config_props = config_output,
565 .request_frame = request_frame,
566 },
567 { NULL }
568};
569
570AVFilter ff_vf_palettegen = {
571 .name = "palettegen",
572 .description = NULL_IF_CONFIG_SMALL("Find the optimal palette for a given stream."),
573 .priv_size = sizeof(PaletteGenContext),
574 .uninit = uninit,
575 .query_formats = query_formats,
576 .inputs = palettegen_inputs,
577 .outputs = palettegen_outputs,
578 .priv_class = &palettegen_class,
579};
580
generated by cgit at 2024-05-17 06:08:39 (GMT)