libavfilter/vf_elbg.c

   1 /*
   2  * Copyright (c) 2013 Stefano Sabatini
   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  * video quantizer filter based on ELBG
  24  */
  25
  26 #include "libavcodec/elbg.h"
  27 #include "libavutil/opt.h"
  28 #include "libavutil/pixdesc.h"
  29 #include "libavutil/random_seed.h"
  30
  31 #include "avfilter.h"
  32 #include "drawutils.h"
  33 #include "internal.h"
  34 #include "video.h"
  35
  36 typedef struct ELBGContext {
  37     const AVClass *class;
  38     AVLFG lfg;
  39     unsigned int lfg_seed;
  40     int max_steps_nb;
  41     int *codeword;
  42     int codeword_length;
  43     int *codeword_closest_codebook_idxs;
  44     int *codebook;
  45     int codebook_length;
  46     const AVPixFmtDescriptor *pix_desc;
  47     uint8_t rgba_map[4];
  48     int pal8;
  49 } ELBGContext;
  50
  51 #define OFFSET(x) offsetof(ELBGContext, x)
  52 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
  53
  54 static const AVOption elbg_options[] = {
  55     { "codebook_length", "set codebook length", OFFSET(codebook_length), AV_OPT_TYPE_INT, { .i64 = 256 }, 1, INT_MAX, FLAGS },
  56     { "l",               "set codebook length", OFFSET(codebook_length), AV_OPT_TYPE_INT, { .i64 = 256 }, 1, INT_MAX, FLAGS },
  57     { "nb_steps", "set max number of steps used to compute the mapping", OFFSET(max_steps_nb), AV_OPT_TYPE_INT, { .i64 = 1 }, 1, INT_MAX, FLAGS },
  58     { "n",        "set max number of steps used to compute the mapping", OFFSET(max_steps_nb), AV_OPT_TYPE_INT, { .i64 = 1 }, 1, INT_MAX, FLAGS },
  59     { "seed", "set the random seed", OFFSET(lfg_seed), AV_OPT_TYPE_INT, {.i64 = -1}, -1, UINT32_MAX, FLAGS },
  60     { "s",    "set the random seed", OFFSET(lfg_seed), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, UINT32_MAX, FLAGS },
  61     { "pal8", "set the pal8 output", OFFSET(pal8), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, FLAGS },
  62     { NULL }
  63 };
  64
  65 AVFILTER_DEFINE_CLASS(elbg);
  66
  67 static av_cold int init(AVFilterContext *ctx)
  68 {
  69     ELBGContext *elbg = ctx->priv;
  70
  71     if (elbg->pal8 && elbg->codebook_length > 256) {
  72         av_log(ctx, AV_LOG_ERROR, "pal8 output allows max 256 codebook length.\n");
  73         return AVERROR(EINVAL);
  74     }
  75
  76     if (elbg->lfg_seed == -1)
  77         elbg->lfg_seed = av_get_random_seed();
  78
  79     av_lfg_init(&elbg->lfg, elbg->lfg_seed);
  80     return 0;
  81 }
  82
  83 static int query_formats(AVFilterContext *ctx)
  84 {
  85     ELBGContext *elbg = ctx->priv;
  86     int ret;
  87
  88     static const enum AVPixelFormat pix_fmts[] = {
  89         AV_PIX_FMT_ARGB, AV_PIX_FMT_RGBA, AV_PIX_FMT_ABGR, AV_PIX_FMT_BGRA,
  90         AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
  91         AV_PIX_FMT_NONE
  92     };
  93     if (!elbg->pal8) {
  94         AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
  95         if (!fmts_list)
  96             return AVERROR(ENOMEM);
  97         return ff_set_common_formats(ctx, fmts_list);
  98     } else {
  99         static const enum AVPixelFormat pal8_fmt[] = {
 100             AV_PIX_FMT_PAL8,
 101             AV_PIX_FMT_NONE
 102         };
 103         if ((ret = ff_formats_ref(ff_make_format_list(pix_fmts), &ctx->inputs[0]->out_formats)) < 0 ||
 104             (ret = ff_formats_ref(ff_make_format_list(pal8_fmt), &ctx->outputs[0]->in_formats)) < 0)
 105             return ret;
 106     }
 107     return 0;
 108 }
 109
 110 #define NB_COMPONENTS 3
 111
 112 static int config_input(AVFilterLink *inlink)
 113 {
 114     AVFilterContext *ctx = inlink->dst;
 115     ELBGContext *elbg = ctx->priv;
 116
 117     elbg->pix_desc = av_pix_fmt_desc_get(inlink->format);
 118     elbg->codeword_length = inlink->w * inlink->h;
 119     elbg->codeword = av_realloc_f(elbg->codeword, elbg->codeword_length,
 120                                   NB_COMPONENTS * sizeof(*elbg->codeword));
 121     if (!elbg->codeword)
 122         return AVERROR(ENOMEM);
 123
 124     elbg->codeword_closest_codebook_idxs =
 125         av_realloc_f(elbg->codeword_closest_codebook_idxs, elbg->codeword_length,
 126                      sizeof(*elbg->codeword_closest_codebook_idxs));
 127     if (!elbg->codeword_closest_codebook_idxs)
 128         return AVERROR(ENOMEM);
 129
 130     elbg->codebook = av_realloc_f(elbg->codebook, elbg->codebook_length,
 131                                   NB_COMPONENTS * sizeof(*elbg->codebook));
 132     if (!elbg->codebook)
 133         return AVERROR(ENOMEM);
 134
 135     ff_fill_rgba_map(elbg->rgba_map, inlink->format);
 136
 137     return 0;
 138 }
 139
 140 #define R 0
 141 #define G 1
 142 #define B 2
 143
 144 static int filter_frame(AVFilterLink *inlink, AVFrame *frame)
 145 {
 146     ELBGContext *elbg = inlink->dst->priv;
 147     int i, j, k;
 148     uint8_t *p, *p0;
 149
 150     const uint8_t r_idx  = elbg->rgba_map[R];
 151     const uint8_t g_idx  = elbg->rgba_map[G];
 152     const uint8_t b_idx  = elbg->rgba_map[B];
 153
 154     /* build the codeword */
 155     p0 = frame->data[0];
 156     k = 0;
 157     for (i = 0; i < inlink->h; i++) {
 158         p = p0;
 159         for (j = 0; j < inlink->w; j++) {
 160             elbg->codeword[k++] = p[r_idx];
 161             elbg->codeword[k++] = p[g_idx];
 162             elbg->codeword[k++] = p[b_idx];
 163             p += elbg->pix_desc->nb_components;
 164         }
 165         p0 += frame->linesize[0];
 166     }
 167
 168     /* compute the codebook */
 169     avpriv_init_elbg(elbg->codeword, NB_COMPONENTS, elbg->codeword_length,
 170                      elbg->codebook, elbg->codebook_length, elbg->max_steps_nb,
 171                      elbg->codeword_closest_codebook_idxs, &elbg->lfg);
 172     avpriv_do_elbg(elbg->codeword, NB_COMPONENTS, elbg->codeword_length,
 173                    elbg->codebook, elbg->codebook_length, elbg->max_steps_nb,
 174                    elbg->codeword_closest_codebook_idxs, &elbg->lfg);
 175
 176     if (elbg->pal8) {
 177         AVFilterLink *outlink = inlink->dst->outputs[0];
 178         AVFrame *out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 179         uint32_t *pal;
 180
 181         if (!out)
 182             return AVERROR(ENOMEM);
 183         out->pts = frame->pts;
 184         av_frame_free(&frame);
 185         pal = (uint32_t *)out->data[1];
 186         p0 = (uint8_t *)out->data[0];
 187
 188         for (i = 0; i < elbg->codebook_length; i++) {
 189             pal[i] =  0xFFU                 << 24  |
 190                      (elbg->codebook[i*3  ] << 16) |
 191                      (elbg->codebook[i*3+1] <<  8) |
 192                       elbg->codebook[i*3+2];
 193         }
 194
 195         k = 0;
 196         for (i = 0; i < inlink->h; i++) {
 197             p = p0;
 198             for (j = 0; j < inlink->w; j++, p++) {
 199                 p[0] = elbg->codeword_closest_codebook_idxs[k++];
 200             }
 201             p0 += out->linesize[0];
 202         }
 203
 204         return ff_filter_frame(outlink, out);
 205     }
 206
 207     /* fill the output with the codebook values */
 208     p0 = frame->data[0];
 209
 210     k = 0;
 211     for (i = 0; i < inlink->h; i++) {
 212         p = p0;
 213         for (j = 0; j < inlink->w; j++) {
 214             int cb_idx = NB_COMPONENTS * elbg->codeword_closest_codebook_idxs[k++];
 215             p[r_idx] = elbg->codebook[cb_idx];
 216             p[g_idx] = elbg->codebook[cb_idx+1];
 217             p[b_idx] = elbg->codebook[cb_idx+2];
 218             p += elbg->pix_desc->nb_components;
 219         }
 220         p0 += frame->linesize[0];
 221     }
 222
 223     return ff_filter_frame(inlink->dst->outputs[0], frame);
 224 }
 225
 226 static av_cold void uninit(AVFilterContext *ctx)
 227 {
 228     ELBGContext *elbg = ctx->priv;
 229
 230     av_freep(&elbg->codebook);
 231     av_freep(&elbg->codeword);
 232     av_freep(&elbg->codeword_closest_codebook_idxs);
 233 }
 234
 235 static const AVFilterPad elbg_inputs[] = {
 236     {
 237         .name           = "default",
 238         .type           = AVMEDIA_TYPE_VIDEO,
 239         .config_props   = config_input,
 240         .filter_frame   = filter_frame,
 241         .needs_writable = 1,
 242     },
 243     { NULL }
 244 };
 245
 246 static const AVFilterPad elbg_outputs[] = {
 247     {
 248         .name = "default",
 249         .type = AVMEDIA_TYPE_VIDEO,
 250     },
 251     { NULL }
 252 };
 253
 254 AVFilter ff_vf_elbg = {
 255     .name          = "elbg",
 256     .description   = NULL_IF_CONFIG_SMALL("Apply posterize effect, using the ELBG algorithm."),
 257     .priv_size     = sizeof(ELBGContext),
 258     .priv_class    = &elbg_class,
 259     .query_formats = query_formats,
 260     .init          = init,
 261     .uninit        = uninit,
 262     .inputs        = elbg_inputs,
 263     .outputs       = elbg_outputs,
 264 };