generic-library/ffmpeg
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity
f099cb2f8a
ffmpeg/libavfilter/vf_palettegen.c
Ganesh Ajjanagadde 92e483f8ed all: use FFDIFFSIGN to resolve possible undefined behavior in comparators 
FFDIFFSIGN was created explicitly for this purpose, since the common
return a - b idiom is unsafe regarding overflow on signed integers. It
optimizes to branchless code on common compilers.

FFDIFFSIGN also has the subjective benefit of being easier to read due
to lack of ternary operators.

Tested with FATE.

Things not covered by this are unsigned integers, for which overflows
are well defined, and also places where overflow is clearly impossible,
e.g an instance where the a - b was being done on 24 bit values.

Reviewed-by: Michael Niedermayer <michael@niedermayer.cc>
Reviewed-by: Clément Bœsch <u@pkh.me>
Signed-off-by: Ganesh Ajjanagadde <gajjanagadde@gmail.com>
2015-11-03 16:28:30 -05:00

568 lines
18 KiB
C
Raw Blame History

/*
* Copyright (c) 2015 Stupeflix
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
/**
* @file
* Generate one palette for a whole video stream.
*/
#include "libavutil/avassert.h"
#include "libavutil/internal.h"
#include "libavutil/opt.h"
#include "libavutil/qsort.h"
#include "avfilter.h"
#include "internal.h"
/* Reference a color and how much it's used */
struct color_ref {
uint32_t color;
uint64_t count;
};
/* Store a range of colors */
struct range_box {
uint32_t color; // average color
int64_t variance; // overall variance of the box (how much the colors are spread)
int start; // index in PaletteGenContext->refs
int len; // number of referenced colors
int sorted_by; // whether range of colors is sorted by red (0), green (1) or blue (2)
};
struct hist_node {
struct color_ref *entries;
int nb_entries;
};
enum {
STATS_MODE_ALL_FRAMES,
STATS_MODE_DIFF_FRAMES,
NB_STATS_MODE
};
#define NBITS 5
#define HIST_SIZE (1<<(3*NBITS))
typedef struct {
const AVClass *class;
int max_colors;
int reserve_transparent;
int stats_mode;
AVFrame *prev_frame; // previous frame used for the diff stats_mode
struct hist_node histogram[HIST_SIZE]; // histogram/hashtable of the colors
struct color_ref **refs; // references of all the colors used in the stream
int nb_refs; // number of color references (or number of different colors)
struct range_box boxes[256]; // define the segmentation of the colorspace (the final palette)
int nb_boxes; // number of boxes (increase will segmenting them)
int palette_pushed; // if the palette frame is pushed into the outlink or not
} PaletteGenContext;
#define OFFSET(x) offsetof(PaletteGenContext, x)
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
static const AVOption palettegen_options[] = {
{ "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 },
{ "reserve_transparent", "reserve a palette entry for transparency", OFFSET(reserve_transparent), AV_OPT_TYPE_BOOL, {.i64=1}, 0, 1, FLAGS },
{ "stats_mode", "set statistics mode", OFFSET(stats_mode), AV_OPT_TYPE_INT, {.i64=STATS_MODE_ALL_FRAMES}, 0, NB_STATS_MODE, FLAGS, "mode" },
{ "full", "compute full frame histograms", 0, AV_OPT_TYPE_CONST, {.i64=STATS_MODE_ALL_FRAMES}, INT_MIN, INT_MAX, FLAGS, "mode" },
{ "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" },
{ NULL }
};
AVFILTER_DEFINE_CLASS(palettegen);
static int query_formats(AVFilterContext *ctx)
{
static const enum AVPixelFormat in_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
static const enum AVPixelFormat out_fmts[] = {AV_PIX_FMT_RGB32, AV_PIX_FMT_NONE};
int ret;
AVFilterFormats *in = ff_make_format_list(in_fmts);
AVFilterFormats *out = ff_make_format_list(out_fmts);
if (!in || !out) {
av_freep(&in);
av_freep(&out);
return AVERROR(ENOMEM);
}
if ((ret = ff_formats_ref(in , &ctx->inputs[0]->out_formats)) < 0 ||
(ret = ff_formats_ref(out, &ctx->outputs[0]->in_formats)) < 0)
return ret;
return 0;
}
typedef int (*cmp_func)(const void *, const void *);
#define DECLARE_CMP_FUNC(name, pos) \
static int cmp_##name(const void *pa, const void *pb) \
{ \
const struct color_ref * const *a = pa; \
const struct color_ref * const *b = pb; \
return ((*a)->color >> (8 * (2 - (pos))) & 0xff) \
- ((*b)->color >> (8 * (2 - (pos))) & 0xff); \
}
DECLARE_CMP_FUNC(r, 0)
DECLARE_CMP_FUNC(g, 1)
DECLARE_CMP_FUNC(b, 2)
static const cmp_func cmp_funcs[] = {cmp_r, cmp_g, cmp_b};
/**
* Simple color comparison for sorting the final palette
*/
static int cmp_color(const void *a, const void *b)
{
const struct range_box *box1 = a;
const struct range_box *box2 = b;
return FFDIFFSIGN(box1->color , box2->color);
}
static av_always_inline int diff(const uint32_t a, const uint32_t b)
{
const uint8_t c1[] = {a >> 16 & 0xff, a >> 8 & 0xff, a & 0xff};
const uint8_t c2[] = {b >> 16 & 0xff, b >> 8 & 0xff, b & 0xff};
const int dr = c1[0] - c2[0];
const int dg = c1[1] - c2[1];
const int db = c1[2] - c2[2];
return dr*dr + dg*dg + db*db;
}
/**
* Find the next box to split: pick the one with the highest variance
*/
static int get_next_box_id_to_split(PaletteGenContext *s)
{
int box_id, i, best_box_id = -1;
int64_t max_variance = -1;
if (s->nb_boxes == s->max_colors - s->reserve_transparent)
return -1;
for (box_id = 0; box_id < s->nb_boxes; box_id++) {
struct range_box *box = &s->boxes[box_id];
if (s->boxes[box_id].len >= 2) {
if (box->variance == -1) {
int64_t variance = 0;
for (i = 0; i < box->len; i++) {
const struct color_ref *ref = s->refs[box->start + i];
variance += diff(ref->color, box->color) * ref->count;
}
box->variance = variance;
}
if (box->variance > max_variance) {
best_box_id = box_id;
max_variance = box->variance;
}
} else {
box->variance = -1;
}
}
return best_box_id;
}
/**
* Get the 32-bit average color for the range of RGB colors enclosed in the
* specified box. Takes into account the weight of each color.
*/
static uint32_t get_avg_color(struct color_ref * const *refs,
const struct range_box *box)
{
int i;
const int n = box->len;
uint64_t r = 0, g = 0, b = 0, div = 0;
for (i = 0; i < n; i++) {
const struct color_ref *ref = refs[box->start + i];
r += (ref->color >> 16 & 0xff) * ref->count;
g += (ref->color >> 8 & 0xff) * ref->count;
b += (ref->color & 0xff) * ref->count;
div += ref->count;
}
r = r / div;
g = g / div;
b = b / div;
return 0xffU<<24 | r<<16 | g<<8 | b;
}
/**
* Split given box in two at position n. The original box becomes the left part
* of the split, and the new index box is the right part.
*/
static void split_box(PaletteGenContext *s, struct range_box *box, int n)
{
struct range_box *new_box = &s->boxes[s->nb_boxes++];
new_box->start = n + 1;
new_box->len = box->start + box->len - new_box->start;
new_box->sorted_by = box->sorted_by;
box->len -= new_box->len;
av_assert0(box->len >= 1);
av_assert0(new_box->len >= 1);
box->color = get_avg_color(s->refs, box);
new_box->color = get_avg_color(s->refs, new_box);
box->variance = -1;
new_box->variance = -1;
}
/**
* Write the palette into the output frame.
*/
static void write_palette(AVFilterContext *ctx, AVFrame *out)
{
const PaletteGenContext *s = ctx->priv;
int x, y, box_id = 0;
uint32_t *pal = (uint32_t *)out->data[0];
const int pal_linesize = out->linesize[0] >> 2;
uint32_t last_color = 0;
for (y = 0; y < out->height; y++) {
for (x = 0; x < out->width; x++) {
if (box_id < s->nb_boxes) {
pal[x] = s->boxes[box_id++].color;
if ((x || y) && pal[x] == last_color)
av_log(ctx, AV_LOG_WARNING, "Dupped color: %08X\n", pal[x]);
last_color = pal[x];
} else {
pal[x] = 0xff000000; // pad with black
}
}
pal += pal_linesize;
}
if (s->reserve_transparent) {
av_assert0(s->nb_boxes < 256);
pal[out->width - pal_linesize - 1] = 0x0000ff00; // add a green transparent color
}
}
/**
* Crawl the histogram to get all the defined colors, and create a linear list
* of them (each color reference entry is a pointer to the value in the
* histogram/hash table).
*/
static struct color_ref **load_color_refs(const struct hist_node *hist, int nb_refs)
{
int i, j, k = 0;
struct color_ref **refs = av_malloc_array(nb_refs, sizeof(*refs));
if (!refs)
return NULL;
for (j = 0; j < HIST_SIZE; j++) {
const struct hist_node *node = &hist[j];
for (i = 0; i < node->nb_entries; i++)
refs[k++] = &node->entries[i];
}
return refs;
}
static double set_colorquant_ratio_meta(AVFrame *out, int nb_out, int nb_in)
{
char buf[32];
const double ratio = (double)nb_out / nb_in;
snprintf(buf, sizeof(buf), "%f", ratio);
av_dict_set(&out->metadata, "lavfi.color_quant_ratio", buf, 0);
return ratio;
}
/**
* Main function implementing the Median Cut Algorithm defined by Paul Heckbert
* in Color Image Quantization for Frame Buffer Display (1982)
*/
static AVFrame *get_palette_frame(AVFilterContext *ctx)
{
AVFrame *out;
PaletteGenContext *s = ctx->priv;
AVFilterLink *outlink = ctx->outputs[0];
double ratio;
int box_id = 0;
struct range_box *box;
/* reference only the used colors from histogram */
s->refs = load_color_refs(s->histogram, s->nb_refs);
if (!s->refs) {
av_log(ctx, AV_LOG_ERROR, "Unable to allocate references for %d different colors\n", s->nb_refs);
return NULL;
}
/* create the palette frame */
out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
if (!out)
return NULL;
out->pts = 0;
/* set first box for 0..nb_refs */
box = &s->boxes[box_id];
box->len = s->nb_refs;
box->sorted_by = -1;
box->color = get_avg_color(s->refs, box);
box->variance = -1;
s->nb_boxes = 1;
while (box && box->len > 1) {
int i, rr, gr, br, longest;
uint64_t median, box_weight = 0;
/* compute the box weight (sum all the weights of the colors in the
* range) and its boundings */
uint8_t min[3] = {0xff, 0xff, 0xff};
uint8_t max[3] = {0x00, 0x00, 0x00};
for (i = box->start; i < box->start + box->len; i++) {
const struct color_ref *ref = s->refs[i];
const uint32_t rgb = ref->color;
const uint8_t r = rgb >> 16 & 0xff, g = rgb >> 8 & 0xff, b = rgb & 0xff;
min[0] = FFMIN(r, min[0]), max[0] = FFMAX(r, max[0]);
min[1] = FFMIN(g, min[1]), max[1] = FFMAX(g, max[1]);
min[2] = FFMIN(b, min[2]), max[2] = FFMAX(b, max[2]);
box_weight += ref->count;
}
/* define the axis to sort by according to the widest range of colors */
rr = max[0] - min[0];
gr = max[1] - min[1];
br = max[2] - min[2];
longest = 1; // pick green by default (the color the eye is the most sensitive to)
if (br >= rr && br >= gr) longest = 2;
if (rr >= gr && rr >= br) longest = 0;
if (gr >= rr && gr >= br) longest = 1; // prefer green again
ff_dlog(ctx, "box #%02X [%6d..%-6d] (%6d) w:%-6"PRIu64" ranges:[%2x %2x %2x] sort by %c (already sorted:%c) ",
box_id, box->start, box->start + box->len - 1, box->len, box_weight,
rr, gr, br, "rgb"[longest], box->sorted_by == longest ? 'y':'n');
/* sort the range by its longest axis if it's not already sorted */
if (box->sorted_by != longest) {
cmp_func cmpf = cmp_funcs[longest];
AV_QSORT(&s->refs[box->start], box->len, const struct color_ref *, cmpf);
box->sorted_by = longest;
}
/* locate the median where to split */
median = (box_weight + 1) >> 1;
box_weight = 0;
/* if you have 2 boxes, the maximum is actually #0: you must have at
* least 1 color on each side of the split, hence the -2 */
for (i = box->start; i < box->start + box->len - 2; i++) {
box_weight += s->refs[i]->count;
if (box_weight > median)
break;
}
ff_dlog(ctx, "split @ i=%-6d with w=%-6"PRIu64" (target=%6"PRIu64")\n", i, box_weight, median);
split_box(s, box, i);
box_id = get_next_box_id_to_split(s);
box = box_id >= 0 ? &s->boxes[box_id] : NULL;
}
ratio = set_colorquant_ratio_meta(out, s->nb_boxes, s->nb_refs);
av_log(ctx, AV_LOG_INFO, "%d%s colors generated out of %d colors; ratio=%f\n",
s->nb_boxes, s->reserve_transparent ? "(+1)" : "", s->nb_refs, ratio);
qsort(s->boxes, s->nb_boxes, sizeof(*s->boxes), cmp_color);
write_palette(ctx, out);
return out;
}
/**
* Hashing function for the color.
* It keeps the NBITS least significant bit of each component to make it
* "random" even if the scene doesn't have much different colors.
*/
static inline unsigned color_hash(uint32_t color)
{
const uint8_t r = color >> 16 & ((1<<NBITS)-1);
const uint8_t g = color >> 8 & ((1<<NBITS)-1);
const uint8_t b = color & ((1<<NBITS)-1);
return r<<(NBITS*2) | g<<NBITS | b;
}
/**
* Locate the color in the hash table and increment its counter.
*/
static int color_inc(struct hist_node *hist, uint32_t color)
{
int i;
const unsigned hash = color_hash(color);
struct hist_node *node = &hist[hash];
struct color_ref *e;
for (i = 0; i < node->nb_entries; i++) {
e = &node->entries[i];
if (e->color == color) {
e->count++;
return 0;
}
}
e = av_dynarray2_add((void**)&node->entries, &node->nb_entries,
sizeof(*node->entries), NULL);
if (!e)
return AVERROR(ENOMEM);
e->color = color;
e->count = 1;
return 1;
}
/**
* Update histogram when pixels differ from previous frame.
*/
static int update_histogram_diff(struct hist_node *hist,
const AVFrame *f1, const AVFrame *f2)
{
int x, y, ret, nb_diff_colors = 0;
for (y = 0; y < f1->height; y++) {
const uint32_t *p = (const uint32_t *)(f1->data[0] + y*f1->linesize[0]);
const uint32_t *q = (const uint32_t *)(f2->data[0] + y*f2->linesize[0]);
for (x = 0; x < f1->width; x++) {
if (p[x] == q[x])
continue;
ret = color_inc(hist, p[x]);
if (ret < 0)
return ret;
nb_diff_colors += ret;
}
}
return nb_diff_colors;
}
/**
* Simple histogram of the frame.
*/
static int update_histogram_frame(struct hist_node *hist, const AVFrame *f)
{
int x, y, ret, nb_diff_colors = 0;
for (y = 0; y < f->height; y++) {
const uint32_t *p = (const uint32_t *)(f->data[0] + y*f->linesize[0]);
for (x = 0; x < f->width; x++) {
ret = color_inc(hist, p[x]);
if (ret < 0)
return ret;
nb_diff_colors += ret;
}
}
return nb_diff_colors;
}
/**
* Update the histogram for each passing frame. No frame will be pushed here.
*/
static int filter_frame(AVFilterLink *inlink, AVFrame *in)
{
AVFilterContext *ctx = inlink->dst;
PaletteGenContext *s = ctx->priv;
const int ret = s->prev_frame ? update_histogram_diff(s->histogram, s->prev_frame, in)
: update_histogram_frame(s->histogram, in);
if (ret > 0)
s->nb_refs += ret;
if (s->stats_mode == STATS_MODE_DIFF_FRAMES) {
av_frame_free(&s->prev_frame);
s->prev_frame = in;
} else {
av_frame_free(&in);
}
return ret;
}
/**
* Returns only one frame at the end containing the full palette.
*/
static int request_frame(AVFilterLink *outlink)
{
AVFilterContext *ctx = outlink->src;
AVFilterLink *inlink = ctx->inputs[0];
PaletteGenContext *s = ctx->priv;
int r;
r = ff_request_frame(inlink);
if (r == AVERROR_EOF && !s->palette_pushed && s->nb_refs) {
r = ff_filter_frame(outlink, get_palette_frame(ctx));
s->palette_pushed = 1;
return r;
}
return r;
}
/**
* The output is one simple 16x16 squared-pixels palette.
*/
static int config_output(AVFilterLink *outlink)
{
outlink->w = outlink->h = 16;
outlink->sample_aspect_ratio = av_make_q(1, 1);
return 0;
}
static av_cold void uninit(AVFilterContext *ctx)
{
int i;
PaletteGenContext *s = ctx->priv;
for (i = 0; i < HIST_SIZE; i++)
av_freep(&s->histogram[i].entries);
av_freep(&s->refs);
av_frame_free(&s->prev_frame);
}
static const AVFilterPad palettegen_inputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
.filter_frame = filter_frame,
},
{ NULL }
};
static const AVFilterPad palettegen_outputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
.config_props = config_output,
.request_frame = request_frame,
},
{ NULL }
};
AVFilter ff_vf_palettegen = {
.name = "palettegen",
.description = NULL_IF_CONFIG_SMALL("Find the optimal palette for a given stream."),
.priv_size = sizeof(PaletteGenContext),
.uninit = uninit,
.query_formats = query_formats,
.inputs = palettegen_inputs,
.outputs = palettegen_outputs,
.priv_class = &palettegen_class,
};
Reference in New Issue View Git Blame Copy Permalink