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

   1 /*
   2  * ColorMatrix v2.2 for Avisynth 2.5.x
   3  *
   4  * Copyright (C) 2006-2007 Kevin Stone
   5  *
   6  * ColorMatrix 1.x is Copyright (C) Wilbert Dijkhof
   7  *
   8  * This program is free software; you can redistribute it and/or modify it
   9  * under the terms of the GNU General Public License as published by the
  10  * Free Software Foundation; either version 2 of the License, or (at your
  11  * option) any later version.
  12  *
  13  * This program is distributed in the hope that it will be useful, but
  14  * OUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  15  * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
  16  * License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License
  19  * along with this program; if not, write to the Free Software Foundation,
  20  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  21  */
  22
  23 /**
  24  * @file
  25  * ColorMatrix 2.0 is based on the original ColorMatrix filter by Wilbert
  26  * Dijkhof.  It adds the ability to convert between any of: Rec.709, FCC,
  27  * Rec.601, and SMPTE 240M. It also makes pre and post clipping optional,
  28  * adds an option to use scaled or non-scaled coefficients, and more...
  29  */
  30
  31 #include <float.h>
  32 #include "avfilter.h"
  33 #include "formats.h"
  34 #include "internal.h"
  35 #include "video.h"
  36 #include "libavutil/opt.h"
  37 #include "libavutil/pixdesc.h"
  38 #include "libavutil/avstring.h"
  39
  40 #define NS(n) ((n) < 0 ? (int)((n)*65536.0-0.5+DBL_EPSILON) : (int)((n)*65536.0+0.5))
  41 #define CB(n) av_clip_uint8(n)
  42
  43 static const double yuv_coeff[4][3][3] = {
  44     { { +0.7152, +0.0722, +0.2126 }, // Rec.709 (0)
  45       { -0.3850, +0.5000, -0.1150 },
  46       { -0.4540, -0.0460, +0.5000 } },
  47     { { +0.5900, +0.1100, +0.3000 }, // FCC (1)
  48       { -0.3310, +0.5000, -0.1690 },
  49       { -0.4210, -0.0790, +0.5000 } },
  50     { { +0.5870, +0.1140, +0.2990 }, // Rec.601 (ITU-R BT.470-2/SMPTE 170M) (2)
  51       { -0.3313, +0.5000, -0.1687 },
  52       { -0.4187, -0.0813, +0.5000 } },
  53     { { +0.7010, +0.0870, +0.2120 }, // SMPTE 240M (3)
  54       { -0.3840, +0.5000, -0.1160 },
  55       { -0.4450, -0.0550, +0.5000 } },
  56 };
  57
  58 enum ColorMode {
  59     COLOR_MODE_NONE = -1,
  60     COLOR_MODE_BT709,
  61     COLOR_MODE_FCC,
  62     COLOR_MODE_BT601,
  63     COLOR_MODE_SMPTE240M,
  64     COLOR_MODE_COUNT
  65 };
  66
  67 typedef struct {
  68     const AVClass *class;
  69     int yuv_convert[16][3][3];
  70     int interlaced;
  71     enum ColorMode source, dest;
  72     int mode;
  73     int hsub, vsub;
  74 } ColorMatrixContext;
  75
  76 #define OFFSET(x) offsetof(ColorMatrixContext, x)
  77 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
  78
  79 static const AVOption colormatrix_options[] = {
  80     { "src", "set source color matrix",      OFFSET(source), AV_OPT_TYPE_INT, {.i64=COLOR_MODE_NONE}, COLOR_MODE_NONE, COLOR_MODE_COUNT-1, .flags=FLAGS, .unit="color_mode" },
  81     { "dst", "set destination color matrix", OFFSET(dest),   AV_OPT_TYPE_INT, {.i64=COLOR_MODE_NONE}, COLOR_MODE_NONE, COLOR_MODE_COUNT-1, .flags=FLAGS, .unit="color_mode" },
  82     { "bt709",     "set BT.709 colorspace",      0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_BT709},       .flags=FLAGS, .unit="color_mode" },
  83     { "fcc",       "set FCC colorspace   ",      0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_FCC},         .flags=FLAGS, .unit="color_mode" },
  84     { "bt601",     "set BT.601 colorspace",      0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_BT601},       .flags=FLAGS, .unit="color_mode" },
  85     { "bt470",     "set BT.470 colorspace",      0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_BT601},       .flags=FLAGS, .unit="color_mode" },
  86     { "smpte170m", "set SMTPE-170M colorspace",  0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_BT601},       .flags=FLAGS, .unit="color_mode" },
  87     { "smpte240m", "set SMPTE-240M colorspace",  0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_SMPTE240M},   .flags=FLAGS, .unit="color_mode" },
  88     { NULL }
  89 };
  90
  91 AVFILTER_DEFINE_CLASS(colormatrix);
  92
  93 #define ma m[0][0]
  94 #define mb m[0][1]
  95 #define mc m[0][2]
  96 #define md m[1][0]
  97 #define me m[1][1]
  98 #define mf m[1][2]
  99 #define mg m[2][0]
 100 #define mh m[2][1]
 101 #define mi m[2][2]
 102
 103 #define ima im[0][0]
 104 #define imb im[0][1]
 105 #define imc im[0][2]
 106 #define imd im[1][0]
 107 #define ime im[1][1]
 108 #define imf im[1][2]
 109 #define img im[2][0]
 110 #define imh im[2][1]
 111 #define imi im[2][2]
 112
 113 static void inverse3x3(double im[3][3], const double m[3][3])
 114 {
 115     double det = ma * (me * mi - mf * mh) - mb * (md * mi - mf * mg) + mc * (md * mh - me * mg);
 116     det = 1.0 / det;
 117     ima = det * (me * mi - mf * mh);
 118     imb = det * (mc * mh - mb * mi);
 119     imc = det * (mb * mf - mc * me);
 120     imd = det * (mf * mg - md * mi);
 121     ime = det * (ma * mi - mc * mg);
 122     imf = det * (mc * md - ma * mf);
 123     img = det * (md * mh - me * mg);
 124     imh = det * (mb * mg - ma * mh);
 125     imi = det * (ma * me - mb * md);
 126 }
 127
 128 static void solve_coefficients(double cm[3][3], double rgb[3][3], const double yuv[3][3])
 129 {
 130     int i, j;
 131     for (i = 0; i < 3; i++)
 132         for (j = 0; j < 3; j++)
 133             cm[i][j] = yuv[i][0] * rgb[0][j] + yuv[i][1] * rgb[1][j] + yuv[i][2] * rgb[2][j];
 134 }
 135
 136 static void calc_coefficients(AVFilterContext *ctx)
 137 {
 138     ColorMatrixContext *color = ctx->priv;
 139     double rgb_coeffd[4][3][3];
 140     double yuv_convertd[16][3][3];
 141     int v = 0;
 142     int i, j, k;
 143
 144     for (i = 0; i < 4; i++)
 145         inverse3x3(rgb_coeffd[i], yuv_coeff[i]);
 146     for (i = 0; i < 4; i++) {
 147         for (j = 0; j < 4; j++) {
 148             solve_coefficients(yuv_convertd[v], rgb_coeffd[i], yuv_coeff[j]);
 149             for (k = 0; k < 3; k++) {
 150                 color->yuv_convert[v][k][0] = NS(yuv_convertd[v][k][0]);
 151                 color->yuv_convert[v][k][1] = NS(yuv_convertd[v][k][1]);
 152                 color->yuv_convert[v][k][2] = NS(yuv_convertd[v][k][2]);
 153             }
 154             if (color->yuv_convert[v][0][0] != 65536 || color->yuv_convert[v][1][0] != 0 ||
 155                 color->yuv_convert[v][2][0] != 0) {
 156                 av_log(ctx, AV_LOG_ERROR, "error calculating conversion coefficients\n");
 157             }
 158             v++;
 159         }
 160     }
 161 }
 162
 163 static const char * const color_modes[] = {"bt709", "fcc", "bt601", "smpte240m"};
 164
 165 static av_cold int init(AVFilterContext *ctx)
 166 {
 167     ColorMatrixContext *color = ctx->priv;
 168
 169     if (color->dest == COLOR_MODE_NONE) {
 170         av_log(ctx, AV_LOG_ERROR, "Unspecified destination color space\n");
 171         return AVERROR(EINVAL);
 172     }
 173
 174     if (color->source == color->dest) {
 175         av_log(ctx, AV_LOG_ERROR, "Source and destination color space must not be identical\n");
 176         return AVERROR(EINVAL);
 177     }
 178
 179     return 0;
 180 }
 181
 182 static void process_frame_uyvy422(ColorMatrixContext *color,
 183                                   AVFrame *dst, AVFrame *src)
 184 {
 185     const unsigned char *srcp = src->data[0];
 186     const int src_pitch = src->linesize[0];
 187     const int height = src->height;
 188     const int width = src->width*2;
 189     unsigned char *dstp = dst->data[0];
 190     const int dst_pitch = dst->linesize[0];
 191     const int c2 = color->yuv_convert[color->mode][0][1];
 192     const int c3 = color->yuv_convert[color->mode][0][2];
 193     const int c4 = color->yuv_convert[color->mode][1][1];
 194     const int c5 = color->yuv_convert[color->mode][1][2];
 195     const int c6 = color->yuv_convert[color->mode][2][1];
 196     const int c7 = color->yuv_convert[color->mode][2][2];
 197     int x, y;
 198
 199     for (y = 0; y < height; y++) {
 200         for (x = 0; x < width; x += 4) {
 201             const int u = srcp[x + 0] - 128;
 202             const int v = srcp[x + 2] - 128;
 203             const int uvval = c2 * u + c3 * v + 1081344;
 204             dstp[x + 0] = CB((c4 * u + c5 * v + 8421376) >> 16);
 205             dstp[x + 1] = CB((65536 * (srcp[x + 1] - 16) + uvval) >> 16);
 206             dstp[x + 2] = CB((c6 * u + c7 * v + 8421376) >> 16);
 207             dstp[x + 3] = CB((65536 * (srcp[x + 3] - 16) + uvval) >> 16);
 208         }
 209         srcp += src_pitch;
 210         dstp += dst_pitch;
 211     }
 212 }
 213
 214 static void process_frame_yuv422p(ColorMatrixContext *color,
 215                                   AVFrame *dst, AVFrame *src)
 216 {
 217     const unsigned char *srcpU = src->data[1];
 218     const unsigned char *srcpV = src->data[2];
 219     const unsigned char *srcpY = src->data[0];
 220     const int src_pitchY  = src->linesize[0];
 221     const int src_pitchUV = src->linesize[1];
 222     const int height = src->height;
 223     const int width = src->width;
 224     unsigned char *dstpU = dst->data[1];
 225     unsigned char *dstpV = dst->data[2];
 226     unsigned char *dstpY = dst->data[0];
 227     const int dst_pitchY  = dst->linesize[0];
 228     const int dst_pitchUV = dst->linesize[1];
 229     const int c2 = color->yuv_convert[color->mode][0][1];
 230     const int c3 = color->yuv_convert[color->mode][0][2];
 231     const int c4 = color->yuv_convert[color->mode][1][1];
 232     const int c5 = color->yuv_convert[color->mode][1][2];
 233     const int c6 = color->yuv_convert[color->mode][2][1];
 234     const int c7 = color->yuv_convert[color->mode][2][2];
 235     int x, y;
 236
 237     for (y = 0; y < height; y++) {
 238         for (x = 0; x < width; x += 2) {
 239             const int u = srcpU[x >> 1] - 128;
 240             const int v = srcpV[x >> 1] - 128;
 241             const int uvval = c2 * u + c3 * v + 1081344;
 242             dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
 243             dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
 244             dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
 245             dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
 246         }
 247         srcpY += src_pitchY;
 248         dstpY += dst_pitchY;
 249         srcpU += src_pitchUV;
 250         srcpV += src_pitchUV;
 251         dstpU += dst_pitchUV;
 252         dstpV += dst_pitchUV;
 253     }
 254 }
 255
 256 static void process_frame_yuv420p(ColorMatrixContext *color,
 257                                   AVFrame *dst, AVFrame *src)
 258 {
 259     const unsigned char *srcpU = src->data[1];
 260     const unsigned char *srcpV = src->data[2];
 261     const unsigned char *srcpY = src->data[0];
 262     const unsigned char *srcpN = src->data[0] + src->linesize[0];
 263     const int src_pitchY  = src->linesize[0];
 264     const int src_pitchUV = src->linesize[1];
 265     const int height = src->height;
 266     const int width = src->width;
 267     unsigned char *dstpU = dst->data[1];
 268     unsigned char *dstpV = dst->data[2];
 269     unsigned char *dstpY = dst->data[0];
 270     unsigned char *dstpN = dst->data[0] + dst->linesize[0];
 271     const int dst_pitchY  = dst->linesize[0];
 272     const int dst_pitchUV = dst->linesize[1];
 273     const int c2 = color->yuv_convert[color->mode][0][1];
 274     const int c3 = color->yuv_convert[color->mode][0][2];
 275     const int c4 = color->yuv_convert[color->mode][1][1];
 276     const int c5 = color->yuv_convert[color->mode][1][2];
 277     const int c6 = color->yuv_convert[color->mode][2][1];
 278     const int c7 = color->yuv_convert[color->mode][2][2];
 279     int x, y;
 280
 281     for (y = 0; y < height; y += 2) {
 282         for (x = 0; x < width; x += 2) {
 283             const int u = srcpU[x >> 1] - 128;
 284             const int v = srcpV[x >> 1] - 128;
 285             const int uvval = c2 * u + c3 * v + 1081344;
 286             dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
 287             dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
 288             dstpN[x + 0] = CB((65536 * (srcpN[x + 0] - 16) + uvval) >> 16);
 289             dstpN[x + 1] = CB((65536 * (srcpN[x + 1] - 16) + uvval) >> 16);
 290             dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
 291             dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
 292         }
 293         srcpY += src_pitchY << 1;
 294         dstpY += dst_pitchY << 1;
 295         srcpN += src_pitchY << 1;
 296         dstpN += dst_pitchY << 1;
 297         srcpU += src_pitchUV;
 298         srcpV += src_pitchUV;
 299         dstpU += dst_pitchUV;
 300         dstpV += dst_pitchUV;
 301     }
 302 }
 303
 304 static int config_input(AVFilterLink *inlink)
 305 {
 306     AVFilterContext *ctx = inlink->dst;
 307     ColorMatrixContext *color = ctx->priv;
 308     const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(inlink->format);
 309
 310     color->hsub = pix_desc->log2_chroma_w;
 311     color->vsub = pix_desc->log2_chroma_h;
 312
 313     av_log(ctx, AV_LOG_VERBOSE, "%s -> %s\n",
 314            color_modes[color->source], color_modes[color->dest]);
 315
 316     return 0;
 317 }
 318
 319 static int query_formats(AVFilterContext *ctx)
 320 {
 321     static const enum AVPixelFormat pix_fmts[] = {
 322         AV_PIX_FMT_YUV422P,
 323         AV_PIX_FMT_YUV420P,
 324         AV_PIX_FMT_UYVY422,
 325         AV_PIX_FMT_NONE
 326     };
 327
 328     ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
 329
 330     return 0;
 331 }
 332
 333 static int filter_frame(AVFilterLink *link, AVFrame *in)
 334 {
 335     AVFilterContext *ctx = link->dst;
 336     ColorMatrixContext *color = ctx->priv;
 337     AVFilterLink *outlink = ctx->outputs[0];
 338     AVFrame *out;
 339
 340     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 341     if (!out) {
 342         av_frame_free(&in);
 343         return AVERROR(ENOMEM);
 344     }
 345     av_frame_copy_props(out, in);
 346
 347     if (color->source == COLOR_MODE_NONE) {
 348         enum AVColorSpace cs = av_frame_get_colorspace(in);
 349         enum ColorMode source;
 350
 351         switch(cs) {
 352         case AVCOL_SPC_BT709     : source = COLOR_MODE_BT709     ; break;
 353         case AVCOL_SPC_FCC       : source = COLOR_MODE_FCC       ; break;
 354         case AVCOL_SPC_SMPTE240M : source = COLOR_MODE_SMPTE240M ; break;
 355         case AVCOL_SPC_BT470BG   : source = COLOR_MODE_BT601     ; break;
 356         case AVCOL_SPC_SMPTE170M : source = COLOR_MODE_BT601     ; break;
 357         default :
 358             av_log(ctx, AV_LOG_ERROR, "Input frame does not specify a supported colorspace, and none has been specified as source either\n");
 359             av_frame_free(&out);
 360             return AVERROR(EINVAL);
 361         }
 362         color->mode = source * 4 + color->dest;
 363     } else
 364         color->mode = color->source * 4 + color->dest;
 365
 366     switch(color->dest) {
 367     case COLOR_MODE_BT709    : av_frame_set_colorspace(out, AVCOL_SPC_BT709)    ; break;
 368     case COLOR_MODE_FCC      : av_frame_set_colorspace(out, AVCOL_SPC_FCC)      ; break;
 369     case COLOR_MODE_SMPTE240M: av_frame_set_colorspace(out, AVCOL_SPC_SMPTE240M); break;
 370     case COLOR_MODE_BT601    : av_frame_set_colorspace(out, AVCOL_SPC_BT470BG)  ; break;
 371     }
 372
 373     calc_coefficients(ctx);
 374
 375     if (in->format == AV_PIX_FMT_YUV422P)
 376         process_frame_yuv422p(color, out, in);
 377     else if (in->format == AV_PIX_FMT_YUV420P)
 378         process_frame_yuv420p(color, out, in);
 379     else
 380         process_frame_uyvy422(color, out, in);
 381
 382     av_frame_free(&in);
 383     return ff_filter_frame(outlink, out);
 384 }
 385
 386 static const AVFilterPad colormatrix_inputs[] = {
 387     {
 388         .name         = "default",
 389         .type         = AVMEDIA_TYPE_VIDEO,
 390         .config_props = config_input,
 391         .filter_frame = filter_frame,
 392     },
 393     { NULL }
 394 };
 395
 396 static const AVFilterPad colormatrix_outputs[] = {
 397     {
 398         .name = "default",
 399         .type = AVMEDIA_TYPE_VIDEO,
 400     },
 401     { NULL }
 402 };
 403
 404 AVFilter ff_vf_colormatrix = {
 405     .name          = "colormatrix",
 406     .description   = NULL_IF_CONFIG_SMALL("Convert color matrix."),
 407     .priv_size     = sizeof(ColorMatrixContext),
 408     .init          = init,
 409     .query_formats = query_formats,
 410     .inputs        = colormatrix_inputs,
 411     .outputs       = colormatrix_outputs,
 412     .priv_class    = &colormatrix_class,
 413     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
 414 };