libavfilter/vf_palettegen.c

   1 /*
   2  * Copyright (c) 2015 Stupeflix
   3  * Copyright (c) 2022 Clément Bœsch <u pkh me>
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 /**
  23  * @file
  24  * Generate one palette for a whole video stream.
  25  */
  26
  27 #include "libavutil/avassert.h"
  28 #include "libavutil/internal.h"
  29 #include "libavutil/opt.h"
  30 #include "libavutil/intreadwrite.h"
  31 #include "avfilter.h"
  32 #include "formats.h"
  33 #include "internal.h"
  34 #include "palette.h"
  35 #include "video.h"
  36
  37 /* Reference a color and how much it's used */
  38 struct color_ref {
  39     uint32_t color;
  40     struct Lab lab;
  41     int64_t count;
  42 };
  43
  44 /* Store a range of colors */
  45 struct range_box {
  46     uint32_t color;     // average color
  47     struct Lab avg;     // average color in perceptual OkLab space
  48     int major_axis;     // best axis candidate for cutting the box
  49     int64_t weight;     // sum of all the weights of the colors
  50     int64_t cut_score;  // how likely the box is to be cut down (higher implying more likely)
  51     int start;          // index in PaletteGenContext->refs
  52     int len;            // number of referenced colors
  53     int sorted_by;      // whether range of colors is sorted by red (0), green (1) or blue (2)
  54 };
  55
  56 struct hist_node {
  57     struct color_ref *entries;
  58     int nb_entries;
  59 };
  60
  61 enum {
  62     STATS_MODE_ALL_FRAMES,
  63     STATS_MODE_DIFF_FRAMES,
  64     STATS_MODE_SINGLE_FRAMES,
  65     NB_STATS_MODE
  66 };
  67
  68 #define HIST_SIZE (1<<15)
  69
  70 typedef struct PaletteGenContext {
  71     const AVClass *class;
  72
  73     int max_colors;
  74     int reserve_transparent;
  75     int stats_mode;
  76
  77     AVFrame *prev_frame;                    // previous frame used for the diff stats_mode
  78     struct hist_node histogram[HIST_SIZE];  // histogram/hashtable of the colors
  79     struct color_ref **refs;                // references of all the colors used in the stream
  80     int nb_refs;                            // number of color references (or number of different colors)
  81     struct range_box boxes[256];            // define the segmentation of the colorspace (the final palette)
  82     int nb_boxes;                           // number of boxes (increase will segmenting them)
  83     int palette_pushed;                     // if the palette frame is pushed into the outlink or not
  84     uint8_t transparency_color[4];          // background color for transparency
  85 } PaletteGenContext;
  86
  87 #define OFFSET(x) offsetof(PaletteGenContext, x)
  88 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
  89 static const AVOption palettegen_options[] = {
  90     { "max_colors", "set the maximum number of colors to use in the palette", OFFSET(max_colors), AV_OPT_TYPE_INT, {.i64=256}, 2, 256, FLAGS },
  91     { "reserve_transparent", "reserve a palette entry for transparency", OFFSET(reserve_transparent), AV_OPT_TYPE_BOOL, {.i64=1}, 0, 1, FLAGS },
  92     { "transparency_color", "set a background color for transparency", OFFSET(transparency_color), AV_OPT_TYPE_COLOR, {.str="lime"}, 0, 0, FLAGS },
  93     { "stats_mode", "set statistics mode", OFFSET(stats_mode), AV_OPT_TYPE_INT, {.i64=STATS_MODE_ALL_FRAMES}, 0, NB_STATS_MODE-1, FLAGS, "mode" },
  94         { "full", "compute full frame histograms", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_ALL_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
  95         { "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" },
  96         { "single", "compute new histogram for each frame", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_SINGLE_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
  97     { NULL }
  98 };
  99
 100 AVFILTER_DEFINE_CLASS(palettegen);
 101
 102 static int query_formats(AVFilterContext *ctx)
 103 {
 104     static const enum AVPixelFormat in_fmts[]  = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
 105     static const enum AVPixelFormat out_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
 106     int ret;
 107
 108     if ((ret = ff_formats_ref(ff_make_format_list(in_fmts) , &ctx->inputs[0]->outcfg.formats)) < 0)
 109         return ret;
 110     if ((ret = ff_formats_ref(ff_make_format_list(out_fmts), &ctx->outputs[0]->incfg.formats)) < 0)
 111         return ret;
 112     return 0;
 113 }
 114
 115 typedef int (*cmp_func)(const void *, const void *);
 116
 117 #define DECLARE_CMP_FUNC(k0, k1, k2)                        \
 118 static int cmp_##k0##k1##k2(const void *pa, const void *pb) \
 119 {                                                           \
 120     const struct color_ref * const *a = pa;                 \
 121     const struct color_ref * const *b = pb;                 \
 122     const int c0 = FFDIFFSIGN((*a)->lab.k0, (*b)->lab.k0);  \
 123     const int c1 = FFDIFFSIGN((*a)->lab.k1, (*b)->lab.k1);  \
 124     const int c2 = FFDIFFSIGN((*a)->lab.k2, (*b)->lab.k2);  \
 125     return c0 ? c0 : c1 ? c1 : c2;                          \
 126 }
 127
 128 DECLARE_CMP_FUNC(L, a, b)
 129 DECLARE_CMP_FUNC(L, b, a)
 130 DECLARE_CMP_FUNC(a, L, b)
 131 DECLARE_CMP_FUNC(a, b, L)
 132 DECLARE_CMP_FUNC(b, L, a)
 133 DECLARE_CMP_FUNC(b, a, L)
 134
 135 enum { ID_XYZ, ID_XZY, ID_ZXY, ID_YXZ, ID_ZYX, ID_YZX };
 136 static const char * const sortstr[] = { "Lab", "Lba", "bLa", "aLb", "baL", "abL" };
 137
 138 static const cmp_func cmp_funcs[] = {
 139     [ID_XYZ] = cmp_Lab,
 140     [ID_XZY] = cmp_Lba,
 141     [ID_ZXY] = cmp_bLa,
 142     [ID_YXZ] = cmp_aLb,
 143     [ID_ZYX] = cmp_baL,
 144     [ID_YZX] = cmp_abL,
 145 };
 146
 147 /*
 148  * Return an identifier for the order of x, y, z (from higher to lower),
 149  * preferring x over y and y over z in case of equality.
 150  */
 151 static int sort3id(int64_t x, int64_t y, int64_t z)
 152 {
 153     if (x >= y) {
 154         if (y >= z) return ID_XYZ;
 155         if (x >= z) return ID_XZY;
 156         return ID_ZXY;
 157     }
 158     if (x >= z) return ID_YXZ;
 159     if (y >= z) return ID_YZX;
 160     return ID_ZYX;
 161 }
 162
 163 /**
 164  * Simple color comparison for sorting the final palette
 165  */
 166 static int cmp_color(const void *a, const void *b)
 167 {
 168     const struct range_box *box1 = a;
 169     const struct range_box *box2 = b;
 170     return FFDIFFSIGN(box1->color, box2->color);
 171 }
 172
 173 static void compute_box_stats(PaletteGenContext *s, struct range_box *box)
 174 {
 175     int64_t er2[3] = {0};
 176
 177     /* Compute average color */
 178     int64_t sL = 0, sa = 0, sb = 0;
 179     box->weight = 0;
 180     for (int i = box->start; i < box->start + box->len; i++) {
 181         const struct color_ref *ref = s->refs[i];
 182         sL += ref->lab.L * ref->count;
 183         sa += ref->lab.a * ref->count;
 184         sb += ref->lab.b * ref->count;
 185         box->weight += ref->count;
 186     }
 187     box->avg.L = sL / box->weight;
 188     box->avg.a = sa / box->weight;
 189     box->avg.b = sb / box->weight;
 190
 191     /* Compute squared error of each color channel */
 192     for (int i = box->start; i < box->start + box->len; i++) {
 193         const struct color_ref *ref = s->refs[i];
 194         const int64_t dL = ref->lab.L - box->avg.L;
 195         const int64_t da = ref->lab.a - box->avg.a;
 196         const int64_t db = ref->lab.b - box->avg.b;
 197         er2[0] += dL * dL * ref->count;
 198         er2[1] += da * da * ref->count;
 199         er2[2] += db * db * ref->count;
 200     }
 201
 202     /* Define the best axis candidate for cutting the box */
 203     box->major_axis = sort3id(er2[0], er2[1], er2[2]);
 204
 205     /* The box that has the axis with the biggest error amongst all boxes will but cut down */
 206     box->cut_score = FFMAX3(er2[0], er2[1], er2[2]);
 207 }
 208
 209 /**
 210  * Find the next box to split: pick the one with the highest cut score
 211  */
 212 static int get_next_box_id_to_split(PaletteGenContext *s)
 213 {
 214     int best_box_id = -1;
 215     int64_t max_score = -1;
 216
 217     if (s->nb_boxes == s->max_colors - s->reserve_transparent)
 218         return -1;
 219
 220     for (int box_id = 0; box_id < s->nb_boxes; box_id++) {
 221         const struct range_box *box = &s->boxes[box_id];
 222         if (s->boxes[box_id].len >= 2 && box->cut_score > max_score) {
 223             best_box_id = box_id;
 224             max_score = box->cut_score;
 225         }
 226     }
 227     return best_box_id;
 228 }
 229
 230 /**
 231  * Split given box in two at position n. The original box becomes the left part
 232  * of the split, and the new index box is the right part.
 233  */
 234 static void split_box(PaletteGenContext *s, struct range_box *box, int n)
 235 {
 236     struct range_box *new_box = &s->boxes[s->nb_boxes++];
 237     new_box->start     = n + 1;
 238     new_box->len       = box->start + box->len - new_box->start;
 239     new_box->sorted_by = box->sorted_by;
 240     box->len -= new_box->len;
 241
 242     av_assert0(box->len     >= 1);
 243     av_assert0(new_box->len >= 1);
 244
 245     compute_box_stats(s, box);
 246     compute_box_stats(s, new_box);
 247 }
 248
 249 /**
 250  * Write the palette into the output frame.
 251  */
 252 static void write_palette(AVFilterContext *ctx, AVFrame *out)
 253 {
 254     const PaletteGenContext *s = ctx->priv;
 255     int box_id = 0;
 256     uint32_t *pal = (uint32_t *)out->data[0];
 257     const int pal_linesize = out->linesize[0] >> 2;
 258     uint32_t last_color = 0;
 259
 260     for (int y = 0; y < out->height; y++) {
 261         for (int x = 0; x < out->width; x++) {
 262             if (box_id < s->nb_boxes) {
 263                 pal[x] = s->boxes[box_id++].color;
 264                 if ((x || y) && pal[x] == last_color)
 265                     av_log(ctx, AV_LOG_WARNING, "Duped color: %08"PRIX32"\n", pal[x]);
 266                 last_color = pal[x];
 267             } else {
 268                 pal[x] = last_color; // pad with last color
 269             }
 270         }
 271         pal += pal_linesize;
 272     }
 273
 274     if (s->reserve_transparent) {
 275         av_assert0(s->nb_boxes < 256);
 276         pal[out->width - pal_linesize - 1] = AV_RB32(&s->transparency_color) >> 8;
 277     }
 278 }
 279
 280 /**
 281  * Crawl the histogram to get all the defined colors, and create a linear list
 282  * of them (each color reference entry is a pointer to the value in the
 283  * histogram/hash table).
 284  */
 285 static struct color_ref **load_color_refs(const struct hist_node *hist, int nb_refs)
 286 {
 287     int k = 0;
 288     struct color_ref **refs = av_malloc_array(nb_refs, sizeof(*refs));
 289
 290     if (!refs)
 291         return NULL;
 292
 293     for (int j = 0; j < HIST_SIZE; j++) {
 294         const struct hist_node *node = &hist[j];
 295
 296         for (int i = 0; i < node->nb_entries; i++)
 297             refs[k++] = &node->entries[i];
 298     }
 299
 300     return refs;
 301 }
 302
 303 static double set_colorquant_ratio_meta(AVFrame *out, int nb_out, int nb_in)
 304 {
 305     char buf[32];
 306     const double ratio = (double)nb_out / nb_in;
 307     snprintf(buf, sizeof(buf), "%f", ratio);
 308     av_dict_set(&out->metadata, "lavfi.color_quant_ratio", buf, 0);
 309     return ratio;
 310 }
 311
 312 /**
 313  * Main function implementing the Median Cut Algorithm defined by Paul Heckbert
 314  * in Color Image Quantization for Frame Buffer Display (1982)
 315  */
 316 static AVFrame *get_palette_frame(AVFilterContext *ctx)
 317 {
 318     AVFrame *out;
 319     PaletteGenContext *s = ctx->priv;
 320     AVFilterLink *outlink = ctx->outputs[0];
 321     double ratio;
 322     int box_id = 0;
 323     struct range_box *box;
 324
 325     /* reference only the used colors from histogram */
 326     s->refs = load_color_refs(s->histogram, s->nb_refs);
 327     if (!s->refs) {
 328         av_log(ctx, AV_LOG_ERROR, "Unable to allocate references for %d different colors\n", s->nb_refs);
 329         return NULL;
 330     }
 331
 332     /* create the palette frame */
 333     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 334     if (!out)
 335         return NULL;
 336     out->pts = 0;
 337
 338     /* set first box for 0..nb_refs */
 339     box = &s->boxes[box_id];
 340     box->len = s->nb_refs;
 341     box->sorted_by = -1;
 342     compute_box_stats(s, box);
 343     s->nb_boxes = 1;
 344
 345     while (box && box->len > 1) {
 346         int i;
 347         int64_t median, weight;
 348
 349         ff_dlog(ctx, "box #%02X [%6d..%-6d] (%6d) w:%-6"PRIu64" sort by %s (already sorted:%c) ",
 350                 box_id, box->start, box->start + box->len - 1, box->len, box->weight,
 351                 sortstr[box->major_axis], box->sorted_by == box->major_axis ? 'y':'n');
 352
 353         /* sort the range by its major axis if it's not already sorted */
 354         if (box->sorted_by != box->major_axis) {
 355             cmp_func cmpf = cmp_funcs[box->major_axis];
 356             qsort(&s->refs[box->start], box->len, sizeof(struct color_ref *), cmpf);
 357             box->sorted_by = box->major_axis;
 358         }
 359
 360         /* locate the median where to split */
 361         median = (box->weight + 1) >> 1;
 362         weight = 0;
 363         /* if you have 2 boxes, the maximum is actually #0: you must have at
 364          * least 1 color on each side of the split, hence the -2 */
 365         for (i = box->start; i < box->start + box->len - 2; i++) {
 366             weight += s->refs[i]->count;
 367             if (weight > median)
 368                 break;
 369         }
 370         ff_dlog(ctx, "split @ i=%-6d with w=%-6"PRIu64" (target=%6"PRIu64")\n", i, weight, median);
 371         split_box(s, box, i);
 372
 373         box_id = get_next_box_id_to_split(s);
 374         box = box_id >= 0 ? &s->boxes[box_id] : NULL;
 375     }
 376
 377     ratio = set_colorquant_ratio_meta(out, s->nb_boxes, s->nb_refs);
 378     av_log(ctx, AV_LOG_INFO, "%d%s colors generated out of %d colors; ratio=%f\n",
 379            s->nb_boxes, s->reserve_transparent ? "(+1)" : "", s->nb_refs, ratio);
 380
 381     for (int i = 0; i < s->nb_boxes; i++)
 382         s->boxes[i].color = 0xffU<<24 | ff_oklab_int_to_srgb_u8(s->boxes[i].avg);
 383
 384     qsort(s->boxes, s->nb_boxes, sizeof(*s->boxes), cmp_color);
 385
 386     write_palette(ctx, out);
 387
 388     return out;
 389 }
 390
 391 /**
 392  * Locate the color in the hash table and increment its counter.
 393  */
 394 static int color_inc(struct hist_node *hist, uint32_t color)
 395 {
 396     const uint32_t hash = ff_lowbias32(color) & (HIST_SIZE - 1);
 397     struct hist_node *node = &hist[hash];
 398     struct color_ref *e;
 399
 400     for (int i = 0; i < node->nb_entries; i++) {
 401         e = &node->entries[i];
 402         if (e->color == color) {
 403             e->count++;
 404             return 0;
 405         }
 406     }
 407
 408     e = av_dynarray2_add((void**)&node->entries, &node->nb_entries,
 409                          sizeof(*node->entries), NULL);
 410     if (!e)
 411         return AVERROR(ENOMEM);
 412     e->color = color;
 413     e->lab = ff_srgb_u8_to_oklab_int(color);
 414     e->count = 1;
 415     return 1;
 416 }
 417
 418 /**
 419  * Update histogram when pixels differ from previous frame.
 420  */
 421 static int update_histogram_diff(struct hist_node *hist,
 422                                  const AVFrame *f1, const AVFrame *f2)
 423 {
 424     int x, y, ret, nb_diff_colors = 0;
 425
 426     for (y = 0; y < f1->height; y++) {
 427         const uint32_t *p = (const uint32_t *)(f1->data[0] + y*f1->linesize[0]);
 428         const uint32_t *q = (const uint32_t *)(f2->data[0] + y*f2->linesize[0]);
 429
 430         for (x = 0; x < f1->width; x++) {
 431             if (p[x] == q[x])
 432                 continue;
 433             ret = color_inc(hist, p[x]);
 434             if (ret < 0)
 435                 return ret;
 436             nb_diff_colors += ret;
 437         }
 438     }
 439     return nb_diff_colors;
 440 }
 441
 442 /**
 443  * Simple histogram of the frame.
 444  */
 445 static int update_histogram_frame(struct hist_node *hist, const AVFrame *f)
 446 {
 447     int x, y, ret, nb_diff_colors = 0;
 448
 449     for (y = 0; y < f->height; y++) {
 450         const uint32_t *p = (const uint32_t *)(f->data[0] + y*f->linesize[0]);
 451
 452         for (x = 0; x < f->width; x++) {
 453             ret = color_inc(hist, p[x]);
 454             if (ret < 0)
 455                 return ret;
 456             nb_diff_colors += ret;
 457         }
 458     }
 459     return nb_diff_colors;
 460 }
 461
 462 /**
 463  * Update the histogram for each passing frame. No frame will be pushed here.
 464  */
 465 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 466 {
 467     AVFilterContext *ctx = inlink->dst;
 468     PaletteGenContext *s = ctx->priv;
 469     int ret;
 470
 471     if (in->color_trc != AVCOL_TRC_UNSPECIFIED && in->color_trc != AVCOL_TRC_IEC61966_2_1)
 472         av_log(ctx, AV_LOG_WARNING, "The input frame is not in sRGB, colors may be off\n");
 473
 474     ret = s->prev_frame ? update_histogram_diff(s->histogram, s->prev_frame, in)
 475                         : update_histogram_frame(s->histogram, in);
 476     if (ret > 0)
 477         s->nb_refs += ret;
 478
 479     if (s->stats_mode == STATS_MODE_DIFF_FRAMES) {
 480         av_frame_free(&s->prev_frame);
 481         s->prev_frame = in;
 482     } else if (s->stats_mode == STATS_MODE_SINGLE_FRAMES && s->nb_refs > 0) {
 483         AVFrame *out;
 484         int i;
 485
 486         out = get_palette_frame(ctx);
 487         out->pts = in->pts;
 488         av_frame_free(&in);
 489         ret = ff_filter_frame(ctx->outputs[0], out);
 490         for (i = 0; i < HIST_SIZE; i++)
 491             av_freep(&s->histogram[i].entries);
 492         av_freep(&s->refs);
 493         s->nb_refs = 0;
 494         s->nb_boxes = 0;
 495         memset(s->boxes, 0, sizeof(s->boxes));
 496         memset(s->histogram, 0, sizeof(s->histogram));
 497     } else {
 498         av_frame_free(&in);
 499     }
 500
 501     return ret;
 502 }
 503
 504 /**
 505  * Returns only one frame at the end containing the full palette.
 506  */
 507 static int request_frame(AVFilterLink *outlink)
 508 {
 509     AVFilterContext *ctx = outlink->src;
 510     AVFilterLink *inlink = ctx->inputs[0];
 511     PaletteGenContext *s = ctx->priv;
 512     int r;
 513
 514     r = ff_request_frame(inlink);
 515     if (r == AVERROR_EOF && !s->palette_pushed && s->nb_refs && s->stats_mode != STATS_MODE_SINGLE_FRAMES) {
 516         r = ff_filter_frame(outlink, get_palette_frame(ctx));
 517         s->palette_pushed = 1;
 518         return r;
 519     }
 520     return r;
 521 }
 522
 523 /**
 524  * The output is one simple 16x16 squared-pixels palette.
 525  */
 526 static int config_output(AVFilterLink *outlink)
 527 {
 528     outlink->w = outlink->h = 16;
 529     outlink->sample_aspect_ratio = av_make_q(1, 1);
 530     return 0;
 531 }
 532
 533 static int init(AVFilterContext *ctx)
 534 {
 535     PaletteGenContext* s = ctx->priv;
 536
 537     if (s->max_colors - s->reserve_transparent < 2) {
 538         av_log(ctx, AV_LOG_ERROR, "max_colors=2 is only allowed without reserving a transparent color slot\n");
 539         return AVERROR(EINVAL);
 540     }
 541
 542     return 0;
 543 }
 544
 545 static av_cold void uninit(AVFilterContext *ctx)
 546 {
 547     int i;
 548     PaletteGenContext *s = ctx->priv;
 549
 550     for (i = 0; i < HIST_SIZE; i++)
 551         av_freep(&s->histogram[i].entries);
 552     av_freep(&s->refs);
 553     av_frame_free(&s->prev_frame);
 554 }
 555
 556 static const AVFilterPad palettegen_inputs[] = {
 557     {
 558         .name         = "default",
 559         .type         = AVMEDIA_TYPE_VIDEO,
 560         .filter_frame = filter_frame,
 561     },
 562 };
 563
 564 static const AVFilterPad palettegen_outputs[] = {
 565     {
 566         .name          = "default",
 567         .type          = AVMEDIA_TYPE_VIDEO,
 568         .config_props  = config_output,
 569         .request_frame = request_frame,
 570     },
 571 };
 572
 573 const AVFilter ff_vf_palettegen = {
 574     .name          = "palettegen",
 575     .description   = NULL_IF_CONFIG_SMALL("Find the optimal palette for a given stream."),
 576     .priv_size     = sizeof(PaletteGenContext),
 577     .init          = init,
 578     .uninit        = uninit,
 579     FILTER_INPUTS(palettegen_inputs),
 580     FILTER_OUTPUTS(palettegen_outputs),
 581     FILTER_QUERY_FUNC(query_formats),
 582     .priv_class    = &palettegen_class,
 583 };