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