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
43static 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
58enum 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
67typedef 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
79static 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    { "smpte240m", "set SMPTE-240M colorspace",  0, AV_OPT_TYPE_CONST, {.i64=COLOR_MODE_SMPTE240M},   .flags=FLAGS, .unit="color_mode" },
86    { NULL }
87};
88
89AVFILTER_DEFINE_CLASS(colormatrix);
90
91#define ma m[0][0]
92#define mb m[0][1]
93#define mc m[0][2]
94#define md m[1][0]
95#define me m[1][1]
96#define mf m[1][2]
97#define mg m[2][0]
98#define mh m[2][1]
99#define mi m[2][2]
100
101#define ima im[0][0]
102#define imb im[0][1]
103#define imc im[0][2]
104#define imd im[1][0]
105#define ime im[1][1]
106#define imf im[1][2]
107#define img im[2][0]
108#define imh im[2][1]
109#define imi im[2][2]
110
111static void inverse3x3(double im[3][3], const double m[3][3])
112{
113    double det = ma * (me * mi - mf * mh) - mb * (md * mi - mf * mg) + mc * (md * mh - me * mg);
114    det = 1.0 / det;
115    ima = det * (me * mi - mf * mh);
116    imb = det * (mc * mh - mb * mi);
117    imc = det * (mb * mf - mc * me);
118    imd = det * (mf * mg - md * mi);
119    ime = det * (ma * mi - mc * mg);
120    imf = det * (mc * md - ma * mf);
121    img = det * (md * mh - me * mg);
122    imh = det * (mb * mg - ma * mh);
123    imi = det * (ma * me - mb * md);
124}
125
126static void solve_coefficients(double cm[3][3], double rgb[3][3], const double yuv[3][3])
127{
128    int i, j;
129    for (i = 0; i < 3; i++)
130        for (j = 0; j < 3; j++)
131            cm[i][j] = yuv[i][0] * rgb[0][j] + yuv[i][1] * rgb[1][j] + yuv[i][2] * rgb[2][j];
132}
133
134static void calc_coefficients(AVFilterContext *ctx)
135{
136    ColorMatrixContext *color = ctx->priv;
137    double rgb_coeffd[4][3][3];
138    double yuv_convertd[16][3][3];
139    int v = 0;
140    int i, j, k;
141
142    for (i = 0; i < 4; i++)
143        inverse3x3(rgb_coeffd[i], yuv_coeff[i]);
144    for (i = 0; i < 4; i++) {
145        for (j = 0; j < 4; j++) {
146            solve_coefficients(yuv_convertd[v], rgb_coeffd[i], yuv_coeff[j]);
147            for (k = 0; k < 3; k++) {
148                color->yuv_convert[v][k][0] = NS(yuv_convertd[v][k][0]);
149                color->yuv_convert[v][k][1] = NS(yuv_convertd[v][k][1]);
150                color->yuv_convert[v][k][2] = NS(yuv_convertd[v][k][2]);
151            }
152            if (color->yuv_convert[v][0][0] != 65536 || color->yuv_convert[v][1][0] != 0 ||
153                color->yuv_convert[v][2][0] != 0) {
154                av_log(ctx, AV_LOG_ERROR, "error calculating conversion coefficients\n");
155            }
156            v++;
157        }
158    }
159}
160
161static const char *color_modes[] = {"bt709", "fcc", "bt601", "smpte240m"};
162
163static av_cold int init(AVFilterContext *ctx)
164{
165    ColorMatrixContext *color = ctx->priv;
166
167    if (color->dest == COLOR_MODE_NONE) {
168        av_log(ctx, AV_LOG_ERROR, "Unspecified destination color space\n");
169        return AVERROR(EINVAL);
170    }
171
172    if (color->source == color->dest) {
173        av_log(ctx, AV_LOG_ERROR, "Source and destination color space must not be identical\n");
174        return AVERROR(EINVAL);
175    }
176
177    return 0;
178}
179
180static void process_frame_uyvy422(ColorMatrixContext *color,
181                                  AVFrame *dst, AVFrame *src)
182{
183    const unsigned char *srcp = src->data[0];
184    const int src_pitch = src->linesize[0];
185    const int height = src->height;
186    const int width = src->width*2;
187    unsigned char *dstp = dst->data[0];
188    const int dst_pitch = dst->linesize[0];
189    const int c2 = color->yuv_convert[color->mode][0][1];
190    const int c3 = color->yuv_convert[color->mode][0][2];
191    const int c4 = color->yuv_convert[color->mode][1][1];
192    const int c5 = color->yuv_convert[color->mode][1][2];
193    const int c6 = color->yuv_convert[color->mode][2][1];
194    const int c7 = color->yuv_convert[color->mode][2][2];
195    int x, y;
196
197    for (y = 0; y < height; y++) {
198        for (x = 0; x < width; x += 4) {
199            const int u = srcp[x + 0] - 128;
200            const int v = srcp[x + 2] - 128;
201            const int uvval = c2 * u + c3 * v + 1081344;
202            dstp[x + 0] = CB((c4 * u + c5 * v + 8421376) >> 16);
203            dstp[x + 1] = CB((65536 * (srcp[x + 1] - 16) + uvval) >> 16);
204            dstp[x + 2] = CB((c6 * u + c7 * v + 8421376) >> 16);
205            dstp[x + 3] = CB((65536 * (srcp[x + 3] - 16) + uvval) >> 16);
206        }
207        srcp += src_pitch;
208        dstp += dst_pitch;
209    }
210}
211
212static void process_frame_yuv422p(ColorMatrixContext *color,
213                                  AVFrame *dst, AVFrame *src)
214{
215    const unsigned char *srcpU = src->data[1];
216    const unsigned char *srcpV = src->data[2];
217    const unsigned char *srcpY = src->data[0];
218    const int src_pitchY  = src->linesize[0];
219    const int src_pitchUV = src->linesize[1];
220    const int height = src->height;
221    const int width = src->width;
222    unsigned char *dstpU = dst->data[1];
223    unsigned char *dstpV = dst->data[2];
224    unsigned char *dstpY = dst->data[0];
225    const int dst_pitchY  = dst->linesize[0];
226    const int dst_pitchUV = dst->linesize[1];
227    const int c2 = color->yuv_convert[color->mode][0][1];
228    const int c3 = color->yuv_convert[color->mode][0][2];
229    const int c4 = color->yuv_convert[color->mode][1][1];
230    const int c5 = color->yuv_convert[color->mode][1][2];
231    const int c6 = color->yuv_convert[color->mode][2][1];
232    const int c7 = color->yuv_convert[color->mode][2][2];
233    int x, y;
234
235    for (y = 0; y < height; y++) {
236        for (x = 0; x < width; x += 2) {
237            const int u = srcpU[x >> 1] - 128;
238            const int v = srcpV[x >> 1] - 128;
239            const int uvval = c2 * u + c3 * v + 1081344;
240            dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
241            dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
242            dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
243            dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
244        }
245        srcpY += src_pitchY;
246        dstpY += dst_pitchY;
247        srcpU += src_pitchUV;
248        srcpV += src_pitchUV;
249        dstpU += dst_pitchUV;
250        dstpV += dst_pitchUV;
251    }
252}
253
254static void process_frame_yuv420p(ColorMatrixContext *color,
255                                  AVFrame *dst, AVFrame *src)
256{
257    const unsigned char *srcpU = src->data[1];
258    const unsigned char *srcpV = src->data[2];
259    const unsigned char *srcpY = src->data[0];
260    const unsigned char *srcpN = src->data[0] + src->linesize[0];
261    const int src_pitchY  = src->linesize[0];
262    const int src_pitchUV = src->linesize[1];
263    const int height = src->height;
264    const int width = src->width;
265    unsigned char *dstpU = dst->data[1];
266    unsigned char *dstpV = dst->data[2];
267    unsigned char *dstpY = dst->data[0];
268    unsigned char *dstpN = dst->data[0] + dst->linesize[0];
269    const int dst_pitchY  = dst->linesize[0];
270    const int dst_pitchUV = dst->linesize[1];
271    const int c2 = color->yuv_convert[color->mode][0][1];
272    const int c3 = color->yuv_convert[color->mode][0][2];
273    const int c4 = color->yuv_convert[color->mode][1][1];
274    const int c5 = color->yuv_convert[color->mode][1][2];
275    const int c6 = color->yuv_convert[color->mode][2][1];
276    const int c7 = color->yuv_convert[color->mode][2][2];
277    int x, y;
278
279    for (y = 0; y < height; y += 2) {
280        for (x = 0; x < width; x += 2) {
281            const int u = srcpU[x >> 1] - 128;
282            const int v = srcpV[x >> 1] - 128;
283            const int uvval = c2 * u + c3 * v + 1081344;
284            dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
285            dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
286            dstpN[x + 0] = CB((65536 * (srcpN[x + 0] - 16) + uvval) >> 16);
287            dstpN[x + 1] = CB((65536 * (srcpN[x + 1] - 16) + uvval) >> 16);
288            dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
289            dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
290        }
291        srcpY += src_pitchY << 1;
292        dstpY += dst_pitchY << 1;
293        srcpN += src_pitchY << 1;
294        dstpN += dst_pitchY << 1;
295        srcpU += src_pitchUV;
296        srcpV += src_pitchUV;
297        dstpU += dst_pitchUV;
298        dstpV += dst_pitchUV;
299    }
300}
301
302static int config_input(AVFilterLink *inlink)
303{
304    AVFilterContext *ctx = inlink->dst;
305    ColorMatrixContext *color = ctx->priv;
306    const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(inlink->format);
307
308    color->hsub = pix_desc->log2_chroma_w;
309    color->vsub = pix_desc->log2_chroma_h;
310
311    av_log(ctx, AV_LOG_VERBOSE, "%s -> %s\n",
312           color_modes[color->source], color_modes[color->dest]);
313
314    return 0;
315}
316
317static int query_formats(AVFilterContext *ctx)
318{
319    static const enum AVPixelFormat pix_fmts[] = {
320        AV_PIX_FMT_YUV422P,
321        AV_PIX_FMT_YUV420P,
322        AV_PIX_FMT_UYVY422,
323        AV_PIX_FMT_NONE
324    };
325
326    ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
327
328    return 0;
329}
330
331static int filter_frame(AVFilterLink *link, AVFrame *in)
332{
333    AVFilterContext *ctx = link->dst;
334    ColorMatrixContext *color = ctx->priv;
335    AVFilterLink *outlink = ctx->outputs[0];
336    AVFrame *out;
337
338    out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
339    if (!out) {
340        av_frame_free(&in);
341        return AVERROR(ENOMEM);
342    }
343    av_frame_copy_props(out, in);
344
345    if (color->source == COLOR_MODE_NONE) {
346        enum AVColorSpace cs = av_frame_get_colorspace(in);
347        enum ColorMode source;
348
349        switch(cs) {
350        case AVCOL_SPC_BT709     : source = COLOR_MODE_BT709     ; break;
351        case AVCOL_SPC_FCC       : source = COLOR_MODE_FCC       ; break;
352        case AVCOL_SPC_SMPTE240M : source = COLOR_MODE_SMPTE240M ; break;
353        case AVCOL_SPC_BT470BG   : source = COLOR_MODE_BT601     ; break;
354        default :
355            av_log(ctx, AV_LOG_ERROR, "Input frame does not specify a supported colorspace, and none has been specified as source either\n");
356            av_frame_free(&out);
357            return AVERROR(EINVAL);
358        }
359        color->mode = source * 4 + color->dest;
360    } else
361        color->mode = color->source * 4 + color->dest;
362
363    switch(color->dest) {
364    case COLOR_MODE_BT709    : av_frame_set_colorspace(out, AVCOL_SPC_BT709)    ; break;
365    case COLOR_MODE_FCC      : av_frame_set_colorspace(out, AVCOL_SPC_FCC)      ; break;
366    case COLOR_MODE_SMPTE240M: av_frame_set_colorspace(out, AVCOL_SPC_SMPTE240M); break;
367    case COLOR_MODE_BT601    : av_frame_set_colorspace(out, AVCOL_SPC_BT470BG)  ; break;
368    }
369
370    calc_coefficients(ctx);
371
372    if (in->format == AV_PIX_FMT_YUV422P)
373        process_frame_yuv422p(color, out, in);
374    else if (in->format == AV_PIX_FMT_YUV420P)
375        process_frame_yuv420p(color, out, in);
376    else
377        process_frame_uyvy422(color, out, in);
378
379    av_frame_free(&in);
380    return ff_filter_frame(outlink, out);
381}
382
383static const AVFilterPad colormatrix_inputs[] = {
384    {
385        .name         = "default",
386        .type         = AVMEDIA_TYPE_VIDEO,
387        .config_props = config_input,
388        .filter_frame = filter_frame,
389    },
390    { NULL }
391};
392
393static const AVFilterPad colormatrix_outputs[] = {
394    {
395        .name = "default",
396        .type = AVMEDIA_TYPE_VIDEO,
397    },
398    { NULL }
399};
400
401AVFilter ff_vf_colormatrix = {
402    .name          = "colormatrix",
403    .description   = NULL_IF_CONFIG_SMALL("Convert color matrix."),
404    .priv_size     = sizeof(ColorMatrixContext),
405    .init          = init,
406    .query_formats = query_formats,
407    .inputs        = colormatrix_inputs,
408    .outputs       = colormatrix_outputs,
409    .priv_class    = &colormatrix_class,
410    .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
411};
412