avfilter/vf_vectorscope: add more pixel formats
[ffmpeg.git] / libavfilter / vf_vectorscope.c
1 /*
2  * Copyright (c) 2015 Paul B Mahol
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/avassert.h"
22 #include "libavutil/opt.h"
23 #include "libavutil/parseutils.h"
24 #include "libavutil/pixdesc.h"
25 #include "avfilter.h"
26 #include "formats.h"
27 #include "internal.h"
28 #include "video.h"
29
30 enum VectorscopeMode {
31     GRAY,
32     COLOR,
33     COLOR2,
34     COLOR3,
35     MODE_NB
36 };
37
38 typedef struct VectorscopeContext {
39     const AVClass *class;
40     int mode;
41     const uint8_t *bg_color;
42     int planewidth[4];
43     int planeheight[4];
44     int x, y, pd;
45     int is_yuv;
46 } VectorscopeContext;
47
48 #define OFFSET(x) offsetof(VectorscopeContext, x)
49 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
50
51 static const AVOption vectorscope_options[] = {
52     { "mode", "set vectorscope mode", OFFSET(mode), AV_OPT_TYPE_INT, {.i64=0}, 0, MODE_NB-1, FLAGS, "mode"},
53     { "gray",   0, 0, AV_OPT_TYPE_CONST, {.i64=GRAY},   0, 0, FLAGS, "mode" },
54     { "color",  0, 0, AV_OPT_TYPE_CONST, {.i64=COLOR},  0, 0, FLAGS, "mode" },
55     { "color2", 0, 0, AV_OPT_TYPE_CONST, {.i64=COLOR2}, 0, 0, FLAGS, "mode" },
56     { "color3", 0, 0, AV_OPT_TYPE_CONST, {.i64=COLOR3}, 0, 0, FLAGS, "mode" },
57     { "x", "set color component on X axis", OFFSET(x), AV_OPT_TYPE_INT, {.i64=1}, 0, 2, FLAGS},
58     { "y", "set color component on Y axis", OFFSET(y), AV_OPT_TYPE_INT, {.i64=2}, 0, 2, FLAGS},
59     { NULL }
60 };
61
62 AVFILTER_DEFINE_CLASS(vectorscope);
63
64 static const enum AVPixelFormat out_yuv_pix_fmts[] = {
65     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P,
66     AV_PIX_FMT_NONE
67 };
68
69 static const enum AVPixelFormat out_rgb_pix_fmts[] = {
70     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
71     AV_PIX_FMT_NONE
72 };
73
74 static const enum AVPixelFormat in1_pix_fmts[] = {
75     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
76     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
77     AV_PIX_FMT_NONE
78 };
79
80 static const enum AVPixelFormat in2_pix_fmts[] = {
81     AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUVJ420P,
82     AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVJ422P,
83     AV_PIX_FMT_YUV411P,  AV_PIX_FMT_YUVJ411P,
84     AV_PIX_FMT_YUV440P,  AV_PIX_FMT_YUV410P,
85     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
86     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
87     AV_PIX_FMT_NONE
88 };
89
90 static int query_formats(AVFilterContext *ctx)
91 {
92     VectorscopeContext *s = ctx->priv;
93     const enum AVPixelFormat *out_pix_fmts;
94     const AVPixFmtDescriptor *desc;
95     AVFilterFormats *avff;
96     int rgb, i;
97
98     if (!ctx->inputs[0]->in_formats ||
99         !ctx->inputs[0]->in_formats->nb_formats) {
100         return AVERROR(EAGAIN);
101     }
102
103     if (!ctx->inputs[0]->out_formats) {
104         const enum AVPixelFormat *in_pix_fmts;
105
106         if ((s->x == 1 && s->y == 2) || (s->x == 2 && s->y == 1))
107             in_pix_fmts = in2_pix_fmts;
108         else
109             in_pix_fmts = in1_pix_fmts;
110         ff_formats_ref(ff_make_format_list(in_pix_fmts), &ctx->inputs[0]->out_formats);
111     }
112
113     avff = ctx->inputs[0]->in_formats;
114     desc = av_pix_fmt_desc_get(avff->formats[0]);
115     rgb = desc->flags & AV_PIX_FMT_FLAG_RGB;
116     for (i = 1; i < avff->nb_formats; i++) {
117         desc = av_pix_fmt_desc_get(avff->formats[i]);
118         if (rgb != desc->flags & AV_PIX_FMT_FLAG_RGB)
119             return AVERROR(EAGAIN);
120     }
121
122     if (rgb)
123         out_pix_fmts = out_rgb_pix_fmts;
124     else
125         out_pix_fmts = out_yuv_pix_fmts;
126     ff_formats_ref(ff_make_format_list(out_pix_fmts), &ctx->outputs[0]->in_formats);
127
128     return 0;
129 }
130
131 static const uint8_t black_yuva_color[4] = { 0, 127, 127, 0 };
132 static const uint8_t black_gbrp_color[4] = { 0, 0, 0, 0 };
133
134 static int config_input(AVFilterLink *inlink)
135 {
136     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
137     VectorscopeContext *s = inlink->dst->priv;
138
139     if (s->mode == GRAY)
140         s->pd = 0;
141     else {
142         if ((s->x == 1 && s->y == 2) || (s->x == 2 && s->y == 1))
143             s->pd = 0;
144         else if ((s->x == 0 && s->y == 2) || (s->x == 2 && s->y == 0))
145             s->pd = 1;
146         else if ((s->x == 0 && s->y == 1) || (s->x == 1 && s->y == 0))
147             s->pd = 2;
148     }
149
150     switch (inlink->format) {
151     case AV_PIX_FMT_GBRAP:
152     case AV_PIX_FMT_GBRP:
153         s->bg_color = black_gbrp_color;
154         break;
155     default:
156         s->bg_color = black_yuva_color;
157     }
158
159     s->planeheight[1] = s->planeheight[2] = FF_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
160     s->planeheight[0] = s->planeheight[3] = inlink->h;
161     s->planewidth[1]  = s->planewidth[2]  = FF_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
162     s->planewidth[0]  = s->planewidth[3]  = inlink->w;
163
164     return 0;
165 }
166
167 static int config_output(AVFilterLink *outlink)
168 {
169     VectorscopeContext *s = outlink->src->priv;
170     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(outlink->format);
171     int depth = desc->comp[0].depth_minus1 + 1;
172
173     s->is_yuv = !(desc->flags & AV_PIX_FMT_FLAG_RGB);
174     outlink->h = outlink->w = 1 << depth;
175     outlink->sample_aspect_ratio = (AVRational){1,1};
176     return 0;
177 }
178
179 static void vectorscope(VectorscopeContext *s, AVFrame *in, AVFrame *out, int pd)
180 {
181     const uint8_t * const *src = (const uint8_t * const *)in->data;
182     const int slinesizex = in->linesize[s->x];
183     const int slinesizey = in->linesize[s->y];
184     const int dlinesize = out->linesize[0];
185     int i, j, px = s->x, py = s->y;
186     const int h = s->planeheight[py];
187     const int w = s->planewidth[px];
188     const uint8_t *spx = src[px];
189     const uint8_t *spy = src[py];
190     uint8_t **dst = out->data;
191     uint8_t *dpx = dst[px];
192     uint8_t *dpy = dst[py];
193     uint8_t *dpd = dst[pd];
194
195     switch (s->mode) {
196     case COLOR:
197     case GRAY:
198         if (s->is_yuv) {
199             for (i = 0; i < h; i++) {
200                 const int iwx = i * slinesizex;
201                 const int iwy = i * slinesizey;
202                 for (j = 0; j < w; j++) {
203                     const int x = spx[iwx + j];
204                     const int y = spy[iwy + j];
205                     const int pos = y * dlinesize + x;
206
207                     dpd[pos] = FFMIN(dpd[pos] + 1, 255);
208                     if (dst[3])
209                         dst[3][pos] = 255;
210                 }
211             }
212         } else {
213             for (i = 0; i < h; i++) {
214                 const int iwx = i * slinesizex;
215                 const int iwy = i * slinesizey;
216                 for (j = 0; j < w; j++) {
217                     const int x = spx[iwx + j];
218                     const int y = spy[iwy + j];
219                     const int pos = y * dlinesize + x;
220
221                     dst[0][pos] = FFMIN(dst[0][pos] + 1, 255);
222                     dst[1][pos] = FFMIN(dst[1][pos] + 1, 255);
223                     dst[2][pos] = FFMIN(dst[2][pos] + 1, 255);
224                     if (dst[3])
225                         dst[3][pos] = 255;
226                 }
227             }
228         }
229         if (s->mode == COLOR) {
230             for (i = 0; i < out->height; i++) {
231                 for (j = 0; j < out->width; j++) {
232                     if (!dpd[i * out->linesize[pd] + j]) {
233                         dpx[i * out->linesize[px] + j] = j;
234                         dpy[i * out->linesize[py] + j] = i;
235                     }
236                 }
237             }
238         }
239         break;
240     case COLOR2:
241         if (s->is_yuv) {
242             for (i = 0; i < h; i++) {
243                 const int iw1 = i * slinesizex;
244                 const int iw2 = i * slinesizey;
245                 for (j = 0; j < w; j++) {
246                     const int x = spx[iw1 + j];
247                     const int y = spy[iw2 + j];
248                     const int pos = y * dlinesize + x;
249
250                     if (!dpd[pos])
251                         dpd[pos] = FFABS(128 - x) + FFABS(128 - y);
252                     dpx[pos] = x;
253                     dpy[pos] = y;
254                     if (dst[3])
255                         dst[3][pos] = 255;
256                 }
257             }
258         } else {
259             for (i = 0; i < h; i++) {
260                 const int iw1 = i * slinesizex;
261                 const int iw2 = i * slinesizey;
262                 for (j = 0; j < w; j++) {
263                     const int x = spx[iw1 + j];
264                     const int y = spy[iw2 + j];
265                     const int pos = y * dlinesize + x;
266
267                     if (!dpd[pos])
268                         dpd[pos] = FFMIN(x + y, 255);
269                     dpx[pos] = x;
270                     dpy[pos] = y;
271                     if (dst[3])
272                         dst[3][pos] = 255;
273                 }
274             }
275         }
276         break;
277     case COLOR3:
278         for (i = 0; i < h; i++) {
279             const int iw1 = i * slinesizex;
280             const int iw2 = i * slinesizey;
281             for (j = 0; j < w; j++) {
282                 const int x = spx[iw1 + j];
283                 const int y = spy[iw2 + j];
284                 const int pos = y * dlinesize + x;
285
286                 dpd[pos] = FFMIN(255, dpd[pos] + 1);
287                 dpx[pos] = x;
288                 dpy[pos] = y;
289                 if (dst[3])
290                     dst[3][pos] = 255;
291             }
292         }
293         break;
294     default:
295         av_assert0(0);
296     }
297 }
298
299 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
300 {
301     AVFilterContext *ctx  = inlink->dst;
302     VectorscopeContext *s = ctx->priv;
303     AVFilterLink *outlink = ctx->outputs[0];
304     AVFrame *out;
305     uint8_t **dst;
306     int i, k;
307
308     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
309     if (!out) {
310         av_frame_free(&in);
311         return AVERROR(ENOMEM);
312     }
313     out->pts = in->pts;
314     dst = out->data;
315
316     for (k = 0; k < 4 && dst[k]; k++)
317         for (i = 0; i < outlink->h ; i++)
318             memset(dst[k] + i * out->linesize[k],
319                    s->mode == COLOR && k == s->pd ? 0 : s->bg_color[k], outlink->w);
320
321     vectorscope(s, in, out, s->pd);
322
323     av_frame_free(&in);
324     return ff_filter_frame(outlink, out);
325 }
326
327 static const AVFilterPad inputs[] = {
328     {
329         .name         = "default",
330         .type         = AVMEDIA_TYPE_VIDEO,
331         .filter_frame = filter_frame,
332         .config_props = config_input,
333     },
334     { NULL }
335 };
336
337 static const AVFilterPad outputs[] = {
338     {
339         .name         = "default",
340         .type         = AVMEDIA_TYPE_VIDEO,
341         .config_props = config_output,
342     },
343     { NULL }
344 };
345
346 AVFilter ff_vf_vectorscope = {
347     .name          = "vectorscope",
348     .description   = NULL_IF_CONFIG_SMALL("Video vectorscope."),
349     .priv_size     = sizeof(VectorscopeContext),
350     .priv_class    = &vectorscope_class,
351     .query_formats = query_formats,
352     .inputs        = inputs,
353     .outputs       = outputs,
354 };