git.sesse.net Git - ffmpeg/blob - libavfilter/vf_chromakey.c

   1 /*
   2  * Copyright (c) 2015 Timo Rothenpieler <timo@rothenpieler.org>
   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 #include "libavutil/opt.h"
  22 #include "libavutil/imgutils.h"
  23 #include "libavutil/intreadwrite.h"
  24 #include "avfilter.h"
  25 #include "formats.h"
  26 #include "internal.h"
  27 #include "video.h"
  28
  29 typedef struct ChromakeyContext {
  30     const AVClass *class;
  31
  32     uint8_t chromakey_rgba[4];
  33     uint16_t chromakey_uv[2];
  34
  35     float similarity;
  36     float blend;
  37
  38     int is_yuv;
  39     int depth;
  40     int mid;
  41     int max;
  42
  43     int hsub_log2;
  44     int vsub_log2;
  45
  46     int (*do_slice)(AVFilterContext *ctx, void *arg,
  47                     int jobnr, int nb_jobs);
  48 } ChromakeyContext;
  49
  50 static uint8_t do_chromakey_pixel(ChromakeyContext *ctx, uint8_t u[9], uint8_t v[9])
  51 {
  52     double diff = 0.0;
  53     int du, dv, i;
  54
  55     for (i = 0; i < 9; ++i) {
  56         du = (int)u[i] - ctx->chromakey_uv[0];
  57         dv = (int)v[i] - ctx->chromakey_uv[1];
  58
  59         diff += sqrt((du * du + dv * dv) / (255.0 * 255.0 * 2));
  60     }
  61
  62     diff /= 9.0;
  63
  64     if (ctx->blend > 0.0001) {
  65         return av_clipd((diff - ctx->similarity) / ctx->blend, 0.0, 1.0) * 255.0;
  66     } else {
  67         return (diff > ctx->similarity) ? 255 : 0;
  68     }
  69 }
  70
  71 static uint16_t do_chromakey_pixel16(ChromakeyContext *ctx, uint16_t u[9], uint16_t v[9])
  72 {
  73     double max = ctx->max;
  74     double diff = 0.0;
  75     int du, dv, i;
  76
  77     for (i = 0; i < 9; ++i) {
  78         du = (int)u[i] - ctx->chromakey_uv[0];
  79         dv = (int)v[i] - ctx->chromakey_uv[1];
  80
  81         diff += sqrt((du * du + dv * dv) / (max * max * 2));
  82     }
  83
  84     diff /= 9.0;
  85
  86     if (ctx->blend > 0.0001) {
  87         return av_clipd((diff - ctx->similarity) / ctx->blend, 0.0, 1.0) * max;
  88     } else {
  89         return (diff > ctx->similarity) ? max : 0;
  90     }
  91 }
  92
  93 static av_always_inline void get_pixel_uv(AVFrame *frame, int hsub_log2, int vsub_log2, int x, int y, uint8_t *u, uint8_t *v)
  94 {
  95     if (x < 0 || x >= frame->width || y < 0 || y >= frame->height)
  96         return;
  97
  98     x >>= hsub_log2;
  99     y >>= vsub_log2;
 100
 101     *u = frame->data[1][frame->linesize[1] * y + x];
 102     *v = frame->data[2][frame->linesize[2] * y + x];
 103 }
 104
 105 static av_always_inline void get_pixel16_uv(AVFrame *frame, int hsub_log2, int vsub_log2, int x, int y, uint16_t *u, uint16_t *v)
 106 {
 107     if (x < 0 || x >= frame->width || y < 0 || y >= frame->height)
 108         return;
 109
 110     x >>= hsub_log2;
 111     y >>= vsub_log2;
 112
 113     *u = AV_RN16(&frame->data[1][frame->linesize[1] * y + 2 * x]);
 114     *v = AV_RN16(&frame->data[2][frame->linesize[2] * y + 2 * x]);
 115 }
 116
 117 static int do_chromakey_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs)
 118 {
 119     AVFrame *frame = arg;
 120
 121     const int slice_start = (frame->height * jobnr) / nb_jobs;
 122     const int slice_end = (frame->height * (jobnr + 1)) / nb_jobs;
 123
 124     ChromakeyContext *ctx = avctx->priv;
 125
 126     int x, y, xo, yo;
 127     uint8_t u[9], v[9];
 128
 129     memset(u, ctx->chromakey_uv[0], sizeof(u));
 130     memset(v, ctx->chromakey_uv[1], sizeof(v));
 131
 132     for (y = slice_start; y < slice_end; ++y) {
 133         for (x = 0; x < frame->width; ++x) {
 134             for (yo = 0; yo < 3; ++yo) {
 135                 for (xo = 0; xo < 3; ++xo) {
 136                     get_pixel_uv(frame, ctx->hsub_log2, ctx->vsub_log2, x + xo - 1, y + yo - 1, &u[yo * 3 + xo], &v[yo * 3 + xo]);
 137                 }
 138             }
 139
 140             frame->data[3][frame->linesize[3] * y + x] = do_chromakey_pixel(ctx, u, v);
 141         }
 142     }
 143
 144     return 0;
 145 }
 146
 147 static int do_chromakey16_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs)
 148 {
 149     AVFrame *frame = arg;
 150
 151     const int slice_start = (frame->height * jobnr) / nb_jobs;
 152     const int slice_end = (frame->height * (jobnr + 1)) / nb_jobs;
 153
 154     ChromakeyContext *ctx = avctx->priv;
 155
 156     int x, y, xo, yo;
 157     uint16_t u[9], v[9];
 158
 159     for (int i = 0; i < 9; i++) {
 160         u[i] = ctx->chromakey_uv[0];
 161         v[i] = ctx->chromakey_uv[1];
 162     }
 163
 164     for (y = slice_start; y < slice_end; ++y) {
 165         for (x = 0; x < frame->width; ++x) {
 166             uint16_t *dst = (uint16_t *)(frame->data[3] + frame->linesize[3] * y);
 167
 168             for (yo = 0; yo < 3; ++yo) {
 169                 for (xo = 0; xo < 3; ++xo) {
 170                     get_pixel16_uv(frame, ctx->hsub_log2, ctx->vsub_log2, x + xo - 1, y + yo - 1, &u[yo * 3 + xo], &v[yo * 3 + xo]);
 171                 }
 172             }
 173
 174             dst[x] = do_chromakey_pixel16(ctx, u, v);
 175         }
 176     }
 177
 178     return 0;
 179 }
 180
 181 static int do_chromahold_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs)
 182 {
 183     ChromakeyContext *ctx = avctx->priv;
 184     AVFrame *frame = arg;
 185     const int slice_start = ((frame->height >> ctx->vsub_log2) * jobnr) / nb_jobs;
 186     const int slice_end = ((frame->height >> ctx->vsub_log2) * (jobnr + 1)) / nb_jobs;
 187
 188     int x, y, alpha;
 189
 190     for (y = slice_start; y < slice_end; ++y) {
 191         for (x = 0; x < frame->width >> ctx->hsub_log2; ++x) {
 192             int u = frame->data[1][frame->linesize[1] * y + x];
 193             int v = frame->data[2][frame->linesize[2] * y + x];
 194             double diff;
 195             int du, dv;
 196
 197             du = u - ctx->chromakey_uv[0];
 198             dv = v - ctx->chromakey_uv[1];
 199
 200             diff = sqrt((du * du + dv * dv) / (255.0 * 255.0));
 201
 202             alpha = diff > ctx->similarity;
 203             if (ctx->blend > 0.0001) {
 204                 double f = 1. - av_clipd((diff - ctx->similarity) / ctx->blend, 0.0, 1.0);
 205
 206                 frame->data[1][frame->linesize[1] * y + x] = 128 + (u - 128) * f;
 207                 frame->data[2][frame->linesize[2] * y + x] = 128 + (v - 128) * f;
 208             } else if (alpha) {
 209                 frame->data[1][frame->linesize[1] * y + x] = 128;
 210                 frame->data[2][frame->linesize[2] * y + x] = 128;
 211             }
 212         }
 213     }
 214
 215     return 0;
 216 }
 217
 218 static int do_chromahold16_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs)
 219 {
 220     ChromakeyContext *ctx = avctx->priv;
 221     AVFrame *frame = arg;
 222     const int slice_start = ((frame->height >> ctx->vsub_log2) * jobnr) / nb_jobs;
 223     const int slice_end = ((frame->height >> ctx->vsub_log2) * (jobnr + 1)) / nb_jobs;
 224     const int mid = ctx->mid;
 225     double max = ctx->max;
 226
 227     int x, y, alpha;
 228
 229     for (y = slice_start; y < slice_end; ++y) {
 230         for (x = 0; x < frame->width >> ctx->hsub_log2; ++x) {
 231             int u = AV_RN16(&frame->data[1][frame->linesize[1] * y + 2 * x]);
 232             int v = AV_RN16(&frame->data[2][frame->linesize[2] * y + 2 * x]);
 233             double diff;
 234             int du, dv;
 235
 236             du = u - ctx->chromakey_uv[0];
 237             dv = v - ctx->chromakey_uv[1];
 238
 239             diff = sqrt((du * du + dv * dv) / (max * max));
 240
 241             alpha = diff > ctx->similarity;
 242             if (ctx->blend > 0.0001) {
 243                 double f = 1. - av_clipd((diff - ctx->similarity) / ctx->blend, 0.0, 1.0);
 244
 245                 AV_WN16(&frame->data[1][frame->linesize[1] * y + 2 * x], mid + (u - mid) * f);
 246                 AV_WN16(&frame->data[2][frame->linesize[2] * y + 2 * x], mid + (v - mid) * f);
 247             } else if (alpha) {
 248                 AV_WN16(&frame->data[1][frame->linesize[1] * y + 2 * x], mid);
 249                 AV_WN16(&frame->data[2][frame->linesize[2] * y + 2 * x], mid);
 250             }
 251         }
 252     }
 253
 254     return 0;
 255 }
 256
 257 static int filter_frame(AVFilterLink *link, AVFrame *frame)
 258 {
 259     AVFilterContext *avctx = link->dst;
 260     ChromakeyContext *ctx = avctx->priv;
 261     int res;
 262
 263     if (res = avctx->internal->execute(avctx, ctx->do_slice, frame, NULL, FFMIN(frame->height, ff_filter_get_nb_threads(avctx))))
 264         return res;
 265
 266     return ff_filter_frame(avctx->outputs[0], frame);
 267 }
 268
 269 #define FIXNUM(x) lrint((x) * (1 << 10))
 270 #define RGB_TO_U(rgb) (((- FIXNUM(0.16874) * rgb[0] - FIXNUM(0.33126) * rgb[1] + FIXNUM(0.50000) * rgb[2] + (1 << 9) - 1) >> 10) + 128)
 271 #define RGB_TO_V(rgb) (((  FIXNUM(0.50000) * rgb[0] - FIXNUM(0.41869) * rgb[1] - FIXNUM(0.08131) * rgb[2] + (1 << 9) - 1) >> 10) + 128)
 272
 273 static av_cold int config_output(AVFilterLink *outlink)
 274 {
 275     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(outlink->format);
 276     AVFilterContext *avctx = outlink->src;
 277     ChromakeyContext *ctx = avctx->priv;
 278     int factor;
 279
 280     ctx->depth = desc->comp[0].depth;
 281     ctx->mid = 1 << (ctx->depth - 1);
 282     ctx->max = (1 << ctx->depth) - 1;
 283
 284     factor = 1 << (ctx->depth - 8);
 285
 286     if (ctx->is_yuv) {
 287         ctx->chromakey_uv[0] = ctx->chromakey_rgba[1] * factor;
 288         ctx->chromakey_uv[1] = ctx->chromakey_rgba[2] * factor;
 289     } else {
 290         ctx->chromakey_uv[0] = RGB_TO_U(ctx->chromakey_rgba) * factor;
 291         ctx->chromakey_uv[1] = RGB_TO_V(ctx->chromakey_rgba) * factor;
 292     }
 293
 294     if (!strcmp(avctx->filter->name, "chromakey")) {
 295         ctx->do_slice = ctx->depth <= 8 ? do_chromakey_slice : do_chromakey16_slice;
 296     } else {
 297         ctx->do_slice = ctx->depth <= 8 ? do_chromahold_slice: do_chromahold16_slice;
 298     }
 299
 300     return 0;
 301 }
 302
 303 static av_cold int query_formats(AVFilterContext *avctx)
 304 {
 305     static const enum AVPixelFormat pixel_fmts[] = {
 306         AV_PIX_FMT_YUVA420P,
 307         AV_PIX_FMT_YUVA422P,
 308         AV_PIX_FMT_YUVA444P,
 309         AV_PIX_FMT_YUVA420P9,  AV_PIX_FMT_YUVA422P9,  AV_PIX_FMT_YUVA444P9,
 310         AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10,
 311         AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12,
 312         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
 313         AV_PIX_FMT_NONE
 314     };
 315
 316     static const enum AVPixelFormat hold_pixel_fmts[] = {
 317         AV_PIX_FMT_YUV420P,
 318         AV_PIX_FMT_YUV422P,
 319         AV_PIX_FMT_YUV444P,
 320         AV_PIX_FMT_YUVA420P,
 321         AV_PIX_FMT_YUVA422P,
 322         AV_PIX_FMT_YUVA444P,
 323         AV_PIX_FMT_YUV420P9,   AV_PIX_FMT_YUV422P9,   AV_PIX_FMT_YUV444P9,
 324         AV_PIX_FMT_YUV420P10,  AV_PIX_FMT_YUV422P10,  AV_PIX_FMT_YUV444P10,
 325         AV_PIX_FMT_YUV444P12,  AV_PIX_FMT_YUV422P12,  AV_PIX_FMT_YUV420P12,
 326         AV_PIX_FMT_YUV444P14,  AV_PIX_FMT_YUV422P14,  AV_PIX_FMT_YUV420P14,
 327         AV_PIX_FMT_YUV420P16,  AV_PIX_FMT_YUV422P16,  AV_PIX_FMT_YUV444P16,
 328         AV_PIX_FMT_YUVA420P9,  AV_PIX_FMT_YUVA422P9,  AV_PIX_FMT_YUVA444P9,
 329         AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10,
 330         AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12,
 331         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
 332         AV_PIX_FMT_NONE
 333     };
 334
 335     AVFilterFormats *formats = NULL;
 336
 337     formats = ff_make_format_list(!strcmp(avctx->filter->name, "chromahold") ? hold_pixel_fmts : pixel_fmts);
 338     if (!formats)
 339         return AVERROR(ENOMEM);
 340
 341     return ff_set_common_formats(avctx, formats);
 342 }
 343
 344 static av_cold int config_input(AVFilterLink *inlink)
 345 {
 346     AVFilterContext *avctx = inlink->dst;
 347     ChromakeyContext *ctx = avctx->priv;
 348     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
 349
 350     ctx->hsub_log2 = desc->log2_chroma_w;
 351     ctx->vsub_log2 = desc->log2_chroma_h;
 352
 353     return 0;
 354 }
 355
 356 static int process_command(AVFilterContext *ctx, const char *cmd, const char *args,
 357                            char *res, int res_len, int flags)
 358 {
 359     int ret;
 360
 361     ret = ff_filter_process_command(ctx, cmd, args, res, res_len, flags);
 362     if (ret < 0)
 363         return ret;
 364
 365     return config_output(ctx->outputs[0]);
 366 }
 367
 368 static const AVFilterPad chromakey_inputs[] = {
 369     {
 370         .name           = "default",
 371         .type           = AVMEDIA_TYPE_VIDEO,
 372         .needs_writable = 1,
 373         .filter_frame   = filter_frame,
 374         .config_props   = config_input,
 375     },
 376     { NULL }
 377 };
 378
 379 static const AVFilterPad chromakey_outputs[] = {
 380     {
 381         .name           = "default",
 382         .type           = AVMEDIA_TYPE_VIDEO,
 383         .config_props   = config_output,
 384     },
 385     { NULL }
 386 };
 387
 388 #define OFFSET(x) offsetof(ChromakeyContext, x)
 389 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
 390
 391 static const AVOption chromakey_options[] = {
 392     { "color", "set the chromakey key color", OFFSET(chromakey_rgba), AV_OPT_TYPE_COLOR, { .str = "black" }, 0, 0, FLAGS },
 393     { "similarity", "set the chromakey similarity value", OFFSET(similarity), AV_OPT_TYPE_FLOAT, { .dbl = 0.01 }, 0.01, 1.0, FLAGS },
 394     { "blend", "set the chromakey key blend value", OFFSET(blend), AV_OPT_TYPE_FLOAT, { .dbl = 0.0 }, 0.0, 1.0, FLAGS },
 395     { "yuv", "color parameter is in yuv instead of rgb", OFFSET(is_yuv), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, FLAGS },
 396     { NULL }
 397 };
 398
 399 AVFILTER_DEFINE_CLASS(chromakey);
 400
 401 const AVFilter ff_vf_chromakey = {
 402     .name          = "chromakey",
 403     .description   = NULL_IF_CONFIG_SMALL("Turns a certain color into transparency. Operates on YUV colors."),
 404     .priv_size     = sizeof(ChromakeyContext),
 405     .priv_class    = &chromakey_class,
 406     .query_formats = query_formats,
 407     .inputs        = chromakey_inputs,
 408     .outputs       = chromakey_outputs,
 409     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
 410     .process_command = process_command,
 411 };
 412
 413 static const AVOption chromahold_options[] = {
 414     { "color", "set the chromahold key color", OFFSET(chromakey_rgba), AV_OPT_TYPE_COLOR, { .str = "black" }, 0, 0, FLAGS },
 415     { "similarity", "set the chromahold similarity value", OFFSET(similarity), AV_OPT_TYPE_FLOAT, { .dbl = 0.01 }, 0.01, 1.0, FLAGS },
 416     { "blend", "set the chromahold blend value", OFFSET(blend), AV_OPT_TYPE_FLOAT, { .dbl = 0.0 }, 0.0, 1.0, FLAGS },
 417     { "yuv", "color parameter is in yuv instead of rgb", OFFSET(is_yuv), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, FLAGS },
 418     { NULL }
 419 };
 420
 421 static const AVFilterPad chromahold_inputs[] = {
 422     {
 423         .name           = "default",
 424         .type           = AVMEDIA_TYPE_VIDEO,
 425         .needs_writable = 1,
 426         .filter_frame   = filter_frame,
 427         .config_props   = config_input,
 428     },
 429     { NULL }
 430 };
 431
 432 static const AVFilterPad chromahold_outputs[] = {
 433     {
 434         .name           = "default",
 435         .type           = AVMEDIA_TYPE_VIDEO,
 436         .config_props   = config_output,
 437     },
 438     { NULL }
 439 };
 440
 441 AVFILTER_DEFINE_CLASS(chromahold);
 442
 443 const AVFilter ff_vf_chromahold = {
 444     .name          = "chromahold",
 445     .description   = NULL_IF_CONFIG_SMALL("Turns a certain color range into gray."),
 446     .priv_size     = sizeof(ChromakeyContext),
 447     .priv_class    = &chromahold_class,
 448     .query_formats = query_formats,
 449     .inputs        = chromahold_inputs,
 450     .outputs       = chromahold_outputs,
 451     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
 452     .process_command = process_command,
 453 };