Merge commit 'a5e8c41c28f907d98d2a739db08f7aef4cbfcf3a'
[ffmpeg.git] / 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  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 <strings.h>
32 #include <float.h>
33 #include "avfilter.h"
34 #include "formats.h"
35 #include "video.h"
36 #include "libavutil/pixdesc.h"
37 #include "libavutil/avstring.h"
38
39 #define NS(n) n < 0 ? (int)(n*65536.0-0.5+DBL_EPSILON) : (int)(n*65536.0+0.5)
40 #define CB(n) av_clip_uint8(n)
41
42 static const double yuv_coeff[4][3][3] = {
43     { { +0.7152, +0.0722, +0.2126 }, // Rec.709 (0)
44       { -0.3850, +0.5000, -0.1150 },
45       { -0.4540, -0.0460, +0.5000 } },
46     { { +0.5900, +0.1100, +0.3000 }, // FCC (1)
47       { -0.3310, +0.5000, -0.1690 },
48       { -0.4210, -0.0790, +0.5000 } },
49     { { +0.5870, +0.1140, +0.2990 }, // Rec.601 (ITU-R BT.470-2/SMPTE 170M) (2)
50       { -0.3313, +0.5000, -0.1687 },
51       { -0.4187, -0.0813, +0.5000 } },
52     { { +0.7010, +0.0870, +0.2120 }, // SMPTE 240M (3)
53       { -0.3840, +0.5000, -0.1160 },
54       { -0.4450, -0.0550, +0.5000 } },
55 };
56
57 typedef struct {
58     int yuv_convert[16][3][3];
59     int interlaced;
60     int source, dest, mode;
61     char src[256];
62     char dst[256];
63     int hsub, vsub;
64 } ColorMatrixContext;
65
66 #define ma m[0][0]
67 #define mb m[0][1]
68 #define mc m[0][2]
69 #define md m[1][0]
70 #define me m[1][1]
71 #define mf m[1][2]
72 #define mg m[2][0]
73 #define mh m[2][1]
74 #define mi m[2][2]
75
76 #define ima im[0][0]
77 #define imb im[0][1]
78 #define imc im[0][2]
79 #define imd im[1][0]
80 #define ime im[1][1]
81 #define imf im[1][2]
82 #define img im[2][0]
83 #define imh im[2][1]
84 #define imi im[2][2]
85
86 static void inverse3x3(double im[3][3], const double m[3][3])
87 {
88     double det = ma * (me * mi - mf * mh) - mb * (md * mi - mf * mg) + mc * (md * mh - me * mg);
89     det = 1.0 / det;
90     ima = det * (me * mi - mf * mh);
91     imb = det * (mc * mh - mb * mi);
92     imc = det * (mb * mf - mc * me);
93     imd = det * (mf * mg - md * mi);
94     ime = det * (ma * mi - mc * mg);
95     imf = det * (mc * md - ma * mf);
96     img = det * (md * mh - me * mg);
97     imh = det * (mb * mg - ma * mh);
98     imi = det * (ma * me - mb * md);
99 }
100
101 static void solve_coefficients(double cm[3][3], double rgb[3][3], const double yuv[3][3])
102 {
103     int i, j;
104     for (i = 0; i < 3; i++)
105         for (j = 0; j < 3; j++)
106             cm[i][j] = yuv[i][0] * rgb[0][j] + yuv[i][1] * rgb[1][j] + yuv[i][2] * rgb[2][j];
107 }
108
109 static void calc_coefficients(AVFilterContext *ctx)
110 {
111     ColorMatrixContext *color = ctx->priv;
112     double rgb_coeffd[4][3][3];
113     double yuv_convertd[16][3][3];
114     int v = 0;
115     int i, j, k;
116
117     for (i = 0; i < 4; i++)
118         inverse3x3(rgb_coeffd[i], yuv_coeff[i]);
119     for (i = 0; i < 4; i++) {
120         for (j = 0; j < 4; j++) {
121             solve_coefficients(yuv_convertd[v], rgb_coeffd[i], yuv_coeff[j]);
122             for (k = 0; k < 3; k++) {
123                 color->yuv_convert[v][k][0] = NS(yuv_convertd[v][k][0]);
124                 color->yuv_convert[v][k][1] = NS(yuv_convertd[v][k][1]);
125                 color->yuv_convert[v][k][2] = NS(yuv_convertd[v][k][2]);
126             }
127             if (color->yuv_convert[v][0][0] != 65536 || color->yuv_convert[v][1][0] != 0 ||
128                 color->yuv_convert[v][2][0] != 0) {
129                 av_log(ctx, AV_LOG_ERROR, "error calculating conversion coefficients\n");
130             }
131             v++;
132         }
133     }
134 }
135
136 static const char *color_modes[] = {"bt709", "FCC", "bt601", "smpte240m"};
137
138 static int get_color_mode_index(const char *name)
139 {
140     int i;
141
142     for (i = 0; i < FF_ARRAY_ELEMS(color_modes); i++)
143         if (!av_strcasecmp(color_modes[i], name))
144             return i;
145     return -1;
146 }
147
148 static av_cold int init(AVFilterContext *ctx, const char *args)
149 {
150     ColorMatrixContext *color = ctx->priv;
151
152     if (!args)
153         goto usage;
154     if (sscanf(args, "%255[^:]:%255[^:]", color->src, color->dst) != 2) {
155     usage:
156         av_log(ctx, AV_LOG_ERROR, "usage: <src>:<dst>\n");
157         av_log(ctx, AV_LOG_ERROR, "possible options: bt709,bt601,smpte240m,fcc\n");
158         return -1;
159     }
160
161     color->source = get_color_mode_index(color->src);
162     if (color->source < 0) {
163         av_log(ctx, AV_LOG_ERROR, "unknown color space %s\n", color->src);
164         return AVERROR(EINVAL);
165     }
166
167     color->dest = get_color_mode_index(color->dst);
168     if (color->dest < 0) {
169         av_log(ctx, AV_LOG_ERROR, "unknown color space %s\n", color->dst);
170         return AVERROR(EINVAL);
171     }
172
173     if (color->source == color->dest) {
174         av_log(ctx, AV_LOG_ERROR, "source and destination color space are identical\n");
175         return AVERROR(EINVAL);
176     }
177
178     color->mode = color->source * 4 + color->dest;
179
180     calc_coefficients(ctx);
181
182     return 0;
183 }
184
185 static void process_frame_uyvy422(ColorMatrixContext *color,
186                                   AVFilterBufferRef *dst, AVFilterBufferRef *src)
187 {
188     const unsigned char *srcp = src->data[0];
189     const int src_pitch = src->linesize[0];
190     const int height = src->video->h;
191     const int width = src->video->w*2;
192     unsigned char *dstp = dst->data[0];
193     const int dst_pitch = dst->linesize[0];
194     const int c2 = color->yuv_convert[color->mode][0][1];
195     const int c3 = color->yuv_convert[color->mode][0][2];
196     const int c4 = color->yuv_convert[color->mode][1][1];
197     const int c5 = color->yuv_convert[color->mode][1][2];
198     const int c6 = color->yuv_convert[color->mode][2][1];
199     const int c7 = color->yuv_convert[color->mode][2][2];
200     int x, y;
201
202     for (y = 0; y < height; y++) {
203         for (x = 0; x < width; x += 4) {
204             const int u = srcp[x + 0] - 128;
205             const int v = srcp[x + 2] - 128;
206             const int uvval = c2 * u + c3 * v + 1081344;
207             dstp[x + 0] = CB((c4 * u + c5 * v + 8421376) >> 16);
208             dstp[x + 1] = CB((65536 * (srcp[x + 1] - 16) + uvval) >> 16);
209             dstp[x + 2] = CB((c6 * u + c7 * v + 8421376) >> 16);
210             dstp[x + 3] = CB((65536 * (srcp[x + 3] - 16) + uvval) >> 16);
211         }
212         srcp += src_pitch;
213         dstp += dst_pitch;
214     }
215 }
216
217 static void process_frame_yuv422p(ColorMatrixContext *color,
218                                   AVFilterBufferRef *dst, AVFilterBufferRef *src)
219 {
220     const unsigned char *srcpU = src->data[1];
221     const unsigned char *srcpV = src->data[2];
222     const unsigned char *srcpY = src->data[0];
223     const int src_pitchY  = src->linesize[0];
224     const int src_pitchUV = src->linesize[1];
225     const int height = src->video->h;
226     const int width = src->video->w;
227     unsigned char *dstpU = dst->data[1];
228     unsigned char *dstpV = dst->data[2];
229     unsigned char *dstpY = dst->data[0];
230     const int dst_pitchY  = dst->linesize[0];
231     const int dst_pitchUV = dst->linesize[1];
232     const int c2 = color->yuv_convert[color->mode][0][1];
233     const int c3 = color->yuv_convert[color->mode][0][2];
234     const int c4 = color->yuv_convert[color->mode][1][1];
235     const int c5 = color->yuv_convert[color->mode][1][2];
236     const int c6 = color->yuv_convert[color->mode][2][1];
237     const int c7 = color->yuv_convert[color->mode][2][2];
238     int x, y;
239
240     for (y = 0; y < height; y++) {
241         for (x = 0; x < width; x += 2) {
242             const int u = srcpU[x >> 1] - 128;
243             const int v = srcpV[x >> 1] - 128;
244             const int uvval = c2 * u + c3 * v + 1081344;
245             dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
246             dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
247             dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
248             dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
249         }
250         srcpY += src_pitchY;
251         dstpY += dst_pitchY;
252         srcpU += src_pitchUV;
253         srcpV += src_pitchUV;
254         dstpU += dst_pitchUV;
255         dstpV += dst_pitchUV;
256     }
257 }
258
259 static void process_frame_yuv420p(ColorMatrixContext *color,
260                                   AVFilterBufferRef *dst, AVFilterBufferRef *src)
261 {
262     const unsigned char *srcpU = src->data[1];
263     const unsigned char *srcpV = src->data[2];
264     const unsigned char *srcpY = src->data[0];
265     const unsigned char *srcpN = src->data[0] + src->linesize[0];
266     const int src_pitchY  = src->linesize[0];
267     const int src_pitchUV = src->linesize[1];
268     const int height = src->video->h;
269     const int width = src->video->w;
270     unsigned char *dstpU = dst->data[1];
271     unsigned char *dstpV = dst->data[2];
272     unsigned char *dstpY = dst->data[0];
273     unsigned char *dstpN = dst->data[0] + dst->linesize[0];
274     const int dst_pitchY  = dst->linesize[0];
275     const int dst_pitchUV = dst->linesize[1];
276     const int c2 = color->yuv_convert[color->mode][0][1];
277     const int c3 = color->yuv_convert[color->mode][0][2];
278     const int c4 = color->yuv_convert[color->mode][1][1];
279     const int c5 = color->yuv_convert[color->mode][1][2];
280     const int c6 = color->yuv_convert[color->mode][2][1];
281     const int c7 = color->yuv_convert[color->mode][2][2];
282     int x, y;
283
284     for (y = 0; y < height; y += 2) {
285         for (x = 0; x < width; x += 2) {
286             const int u = srcpU[x >> 1] - 128;
287             const int v = srcpV[x >> 1] - 128;
288             const int uvval = c2 * u + c3 * v + 1081344;
289             dstpY[x + 0] = CB((65536 * (srcpY[x + 0] - 16) + uvval) >> 16);
290             dstpY[x + 1] = CB((65536 * (srcpY[x + 1] - 16) + uvval) >> 16);
291             dstpN[x + 0] = CB((65536 * (srcpN[x + 0] - 16) + uvval) >> 16);
292             dstpN[x + 1] = CB((65536 * (srcpN[x + 1] - 16) + uvval) >> 16);
293             dstpU[x >> 1] = CB((c4 * u + c5 * v + 8421376) >> 16);
294             dstpV[x >> 1] = CB((c6 * u + c7 * v + 8421376) >> 16);
295         }
296         srcpY += src_pitchY << 1;
297         dstpY += dst_pitchY << 1;
298         srcpN += src_pitchY << 1;
299         dstpN += dst_pitchY << 1;
300         srcpU += src_pitchUV;
301         srcpV += src_pitchUV;
302         dstpU += dst_pitchUV;
303         dstpV += dst_pitchUV;
304     }
305 }
306
307 static int config_input(AVFilterLink *inlink)
308 {
309     AVFilterContext *ctx = inlink->dst;
310     ColorMatrixContext *color = ctx->priv;
311     const AVPixFmtDescriptor *pix_desc = &av_pix_fmt_descriptors[inlink->format];
312
313     color->hsub = pix_desc->log2_chroma_w;
314     color->vsub = pix_desc->log2_chroma_h;
315
316     av_log(ctx, AV_LOG_INFO, "%s -> %s\n", color->src, color->dst);
317
318     return 0;
319 }
320
321 static int query_formats(AVFilterContext *ctx)
322 {
323     static const enum PixelFormat pix_fmts[] = {
324         PIX_FMT_YUV422P,
325         PIX_FMT_YUV420P,
326         PIX_FMT_UYVY422,
327         PIX_FMT_NONE
328     };
329
330     ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
331
332     return 0;
333 }
334
335 static AVFilterBufferRef *get_video_buffer(AVFilterLink *inlink, int perms, int w, int h)
336 {
337     AVFilterBufferRef *picref =
338         ff_get_video_buffer(inlink->dst->outputs[0], perms, w, h);
339     return picref;
340 }
341
342 static void start_frame(AVFilterLink *link, AVFilterBufferRef *picref)
343 {
344     AVFilterBufferRef *outpicref = avfilter_ref_buffer(picref, ~0);
345
346     link->dst->outputs[0]->out_buf = outpicref;
347
348     ff_start_frame(link->dst->outputs[0], outpicref);
349 }
350
351 static void end_frame(AVFilterLink *link)
352 {
353     AVFilterContext *ctx = link->dst;
354     ColorMatrixContext *color = ctx->priv;
355     AVFilterBufferRef *out = link->dst->outputs[0]->out_buf;
356
357     if (link->cur_buf->format == PIX_FMT_YUV422P)
358         process_frame_yuv422p(color, out, link->cur_buf);
359     else if (link->cur_buf->format == PIX_FMT_YUV420P)
360         process_frame_yuv420p(color, out, link->cur_buf);
361     else
362         process_frame_uyvy422(color, out, link->cur_buf);
363
364     ff_draw_slice(ctx->outputs[0], 0, link->dst->outputs[0]->h, 1);
365     ff_end_frame(ctx->outputs[0]);
366     avfilter_unref_buffer(link->cur_buf);
367 }
368
369 static void null_draw_slice(AVFilterLink *link, int y, int h, int slice_dir) { }
370
371 AVFilter avfilter_vf_colormatrix = {
372     .name          = "colormatrix",
373     .description   = NULL_IF_CONFIG_SMALL("Color matrix conversion"),
374
375     .priv_size     = sizeof(ColorMatrixContext),
376     .init          = init,
377     .query_formats = query_formats,
378
379     .inputs    = (AVFilterPad[]) {{ .name             = "default",
380                                     .type             = AVMEDIA_TYPE_VIDEO,
381                                     .config_props     = config_input,
382                                     .start_frame      = start_frame,
383                                     .get_video_buffer = get_video_buffer,
384                                     .draw_slice       = null_draw_slice,
385                                     .end_frame        = end_frame, },
386                                   { .name = NULL }},
387
388     .outputs   = (AVFilterPad[]) {{ .name             = "default",
389                                     .type             = AVMEDIA_TYPE_VIDEO, },
390                                   { .name = NULL }},
391 };