git.sesse.net Git - ffmpeg/blob - 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
  87     static const enum AVPixelFormat pix_fmts[] = {
  88         AV_PIX_FMT_ARGB, AV_PIX_FMT_RGBA, AV_PIX_FMT_ABGR, AV_PIX_FMT_BGRA,
  89         AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
  90         AV_PIX_FMT_NONE
  91     };
  92     if (!elbg->pal8) {
  93         AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
  94         if (!fmts_list)
  95             return AVERROR(ENOMEM);
  96         return ff_set_common_formats(ctx, fmts_list);
  97     } else {
  98         static const enum AVPixelFormat pal8_fmt[] = {
  99             AV_PIX_FMT_PAL8,
 100             AV_PIX_FMT_NONE
 101         };
 102         ff_formats_ref(ff_make_format_list(pix_fmts), &ctx->inputs[0]->out_formats);
 103         ff_formats_ref(ff_make_format_list(pal8_fmt), &ctx->outputs[0]->in_formats);
 104     }
 105     return 0;
 106 }
 107
 108 #define NB_COMPONENTS 3
 109
 110 static int config_input(AVFilterLink *inlink)
 111 {
 112     AVFilterContext *ctx = inlink->dst;
 113     ELBGContext *elbg = ctx->priv;
 114
 115     elbg->pix_desc = av_pix_fmt_desc_get(inlink->format);
 116     elbg->codeword_length = inlink->w * inlink->h;
 117     elbg->codeword = av_realloc_f(elbg->codeword, elbg->codeword_length,
 118                                   NB_COMPONENTS * sizeof(*elbg->codeword));
 119     if (!elbg->codeword)
 120         return AVERROR(ENOMEM);
 121
 122     elbg->codeword_closest_codebook_idxs =
 123         av_realloc_f(elbg->codeword_closest_codebook_idxs, elbg->codeword_length,
 124                      sizeof(*elbg->codeword_closest_codebook_idxs));
 125     if (!elbg->codeword_closest_codebook_idxs)
 126         return AVERROR(ENOMEM);
 127
 128     elbg->codebook = av_realloc_f(elbg->codebook, elbg->codebook_length,
 129                                   NB_COMPONENTS * sizeof(*elbg->codebook));
 130     if (!elbg->codebook)
 131         return AVERROR(ENOMEM);
 132
 133     ff_fill_rgba_map(elbg->rgba_map, inlink->format);
 134
 135     return 0;
 136 }
 137
 138 #define R 0
 139 #define G 1
 140 #define B 2
 141
 142 static int filter_frame(AVFilterLink *inlink, AVFrame *frame)
 143 {
 144     ELBGContext *elbg = inlink->dst->priv;
 145     int i, j, k;
 146     uint8_t *p, *p0;
 147
 148     const uint8_t r_idx  = elbg->rgba_map[R];
 149     const uint8_t g_idx  = elbg->rgba_map[G];
 150     const uint8_t b_idx  = elbg->rgba_map[B];
 151
 152     /* build the codeword */
 153     p0 = frame->data[0];
 154     k = 0;
 155     for (i = 0; i < inlink->h; i++) {
 156         p = p0;
 157         for (j = 0; j < inlink->w; j++) {
 158             elbg->codeword[k++] = p[r_idx];
 159             elbg->codeword[k++] = p[g_idx];
 160             elbg->codeword[k++] = p[b_idx];
 161             p += elbg->pix_desc->nb_components;
 162         }
 163         p0 += frame->linesize[0];
 164     }
 165
 166     /* compute the codebook */
 167     avpriv_init_elbg(elbg->codeword, NB_COMPONENTS, elbg->codeword_length,
 168                      elbg->codebook, elbg->codebook_length, elbg->max_steps_nb,
 169                      elbg->codeword_closest_codebook_idxs, &elbg->lfg);
 170     avpriv_do_elbg(elbg->codeword, NB_COMPONENTS, elbg->codeword_length,
 171                    elbg->codebook, elbg->codebook_length, elbg->max_steps_nb,
 172                    elbg->codeword_closest_codebook_idxs, &elbg->lfg);
 173
 174     if (elbg->pal8) {
 175         AVFilterLink *outlink = inlink->dst->outputs[0];
 176         AVFrame *out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 177         uint32_t *pal;
 178
 179         if (!out)
 180             return AVERROR(ENOMEM);
 181         out->pts = frame->pts;
 182         av_frame_free(&frame);
 183         pal = (uint32_t *)out->data[1];
 184         p0 = (uint8_t *)out->data[0];
 185
 186         for (i = 0; i < elbg->codebook_length; i++) {
 187             pal[i] = (elbg->codebook[i*3  ] << 16) |
 188                      (elbg->codebook[i*3+1] <<  8) |
 189                       elbg->codebook[i*3+2];
 190         }
 191
 192         k = 0;
 193         for (i = 0; i < inlink->h; i++) {
 194             p = p0;
 195             for (j = 0; j < inlink->w; j++, p++) {
 196                 p[0] = elbg->codeword_closest_codebook_idxs[k++];
 197             }
 198             p0 += out->linesize[0];
 199         }
 200
 201         return ff_filter_frame(outlink, out);
 202     }
 203
 204     /* fill the output with the codebook values */
 205     p0 = frame->data[0];
 206
 207     k = 0;
 208     for (i = 0; i < inlink->h; i++) {
 209         p = p0;
 210         for (j = 0; j < inlink->w; j++) {
 211             int cb_idx = NB_COMPONENTS * elbg->codeword_closest_codebook_idxs[k++];
 212             p[r_idx] = elbg->codebook[cb_idx];
 213             p[g_idx] = elbg->codebook[cb_idx+1];
 214             p[b_idx] = elbg->codebook[cb_idx+2];
 215             p += elbg->pix_desc->nb_components;
 216         }
 217         p0 += frame->linesize[0];
 218     }
 219
 220     return ff_filter_frame(inlink->dst->outputs[0], frame);
 221 }
 222
 223 static av_cold void uninit(AVFilterContext *ctx)
 224 {
 225     ELBGContext *elbg = ctx->priv;
 226
 227     av_freep(&elbg->codebook);
 228     av_freep(&elbg->codeword);
 229     av_freep(&elbg->codeword_closest_codebook_idxs);
 230 }
 231
 232 static const AVFilterPad elbg_inputs[] = {
 233     {
 234         .name           = "default",
 235         .type           = AVMEDIA_TYPE_VIDEO,
 236         .config_props   = config_input,
 237         .filter_frame   = filter_frame,
 238         .needs_writable = 1,
 239     },
 240     { NULL }
 241 };
 242
 243 static const AVFilterPad elbg_outputs[] = {
 244     {
 245         .name = "default",
 246         .type = AVMEDIA_TYPE_VIDEO,
 247     },
 248     { NULL }
 249 };
 250
 251 AVFilter ff_vf_elbg = {
 252     .name          = "elbg",
 253     .description   = NULL_IF_CONFIG_SMALL("Apply posterize effect, using the ELBG algorithm."),
 254     .priv_size     = sizeof(ELBGContext),
 255     .priv_class    = &elbg_class,
 256     .query_formats = query_formats,
 257     .init          = init,
 258     .uninit        = uninit,
 259     .inputs        = elbg_inputs,
 260     .outputs       = elbg_outputs,
 261 };