avfilter/vf_vectorscope: 9 & 10 bit depth support

Signed-off-by: Paul B Mahol <onemda@gmail.com>

libavfilter/vf_vectorscope.c

@@ -19,6 +19,7 @@
  */
 
 #include "libavutil/avassert.h"
+#include "libavutil/intreadwrite.h"
 #include "libavutil/opt.h"
 #include "libavutil/parseutils.h"
 #include "libavutil/pixdesc.h"
@@ -47,8 +48,13 @@ typedef struct VectorscopeContext {
     int hsub, vsub;
     int x, y, pd;
     int is_yuv;
+    int size;
+    int mult;
     int envelope;
-    uint8_t peak[256][256];
+    uint8_t peak[1024][1024];
+
+    void (*vectorscope)(struct VectorscopeContext *s,
+                        AVFrame *in, AVFrame *out, int pd);
 } VectorscopeContext;
 
 #define OFFSET(x) offsetof(VectorscopeContext, x)
@@ -77,29 +83,54 @@ static const AVOption vectorscope_options[] = {
 
 AVFILTER_DEFINE_CLASS(vectorscope);
 
-static const enum AVPixelFormat out_yuv_pix_fmts[] = {
+static const enum AVPixelFormat out_yuv8_pix_fmts[] = {
     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P,
     AV_PIX_FMT_NONE
 };
 
-static const enum AVPixelFormat out_rgb_pix_fmts[] = {
+static const enum AVPixelFormat out_yuv9_pix_fmts[] = {
+    AV_PIX_FMT_YUV444P9,
+    AV_PIX_FMT_NONE
+};
+
+static const enum AVPixelFormat out_yuv10_pix_fmts[] = {
+    AV_PIX_FMT_YUV444P10,
+    AV_PIX_FMT_NONE
+};
+
+static const enum AVPixelFormat out_rgb8_pix_fmts[] = {
     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
     AV_PIX_FMT_NONE
 };
 
+static const enum AVPixelFormat out_rgb9_pix_fmts[] = {
+    AV_PIX_FMT_GBRP9,
+    AV_PIX_FMT_NONE
+};
+
+static const enum AVPixelFormat out_rgb10_pix_fmts[] = {
+    AV_PIX_FMT_GBRP10,
+    AV_PIX_FMT_NONE
+};
+
 static const enum AVPixelFormat in1_pix_fmts[] = {
     AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
+    AV_PIX_FMT_YUV444P9, AV_PIX_FMT_YUV444P10,
     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
+    AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
     AV_PIX_FMT_NONE
 };
 
 static const enum AVPixelFormat in2_pix_fmts[] = {
     AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUVJ420P,
     AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVJ422P,
+    AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
     AV_PIX_FMT_YUV411P,  AV_PIX_FMT_YUVJ411P,
     AV_PIX_FMT_YUV440P,  AV_PIX_FMT_YUV410P,
-    AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP,
+    AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
+    AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9,
+    AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
     AV_PIX_FMT_NONE
 };
 
@@ -109,7 +140,7 @@ static int query_formats(AVFilterContext *ctx)
     const enum AVPixelFormat *out_pix_fmts;
     const AVPixFmtDescriptor *desc;
     AVFilterFormats *avff;
-    int rgb, i;
+    int depth, rgb, i;
 
     if (!ctx->inputs[0]->in_formats ||
         !ctx->inputs[0]->in_formats->nb_formats) {
@@ -129,16 +160,26 @@ static int query_formats(AVFilterContext *ctx)
     avff = ctx->inputs[0]->in_formats;
     desc = av_pix_fmt_desc_get(avff->formats[0]);
     rgb = desc->flags & AV_PIX_FMT_FLAG_RGB;
+    depth = desc->comp[0].depth;
     for (i = 1; i < avff->nb_formats; i++) {
         desc = av_pix_fmt_desc_get(avff->formats[i]);
-        if (rgb != (desc->flags & AV_PIX_FMT_FLAG_RGB))
+        if (rgb != (desc->flags & AV_PIX_FMT_FLAG_RGB) ||
+            depth != desc->comp[0].depth)
             return AVERROR(EAGAIN);
     }
 
-    if (rgb)
-        out_pix_fmts = out_rgb_pix_fmts;
+    if (rgb && depth == 8)
+        out_pix_fmts = out_rgb8_pix_fmts;
+    else if (rgb && depth == 9)
+        out_pix_fmts = out_rgb9_pix_fmts;
+    else if (rgb && depth == 10)
+        out_pix_fmts = out_rgb10_pix_fmts;
+    else if (depth == 9)
+        out_pix_fmts = out_yuv9_pix_fmts;
+    else if (depth == 10)
+        out_pix_fmts = out_yuv10_pix_fmts;
     else
-        out_pix_fmts = out_yuv_pix_fmts;
+        out_pix_fmts = out_yuv8_pix_fmts;
     ff_formats_ref(ff_make_format_list(out_pix_fmts), &ctx->outputs[0]->in_formats);
 
     return 0;
@@ -147,53 +188,66 @@ static int query_formats(AVFilterContext *ctx)
 static const uint8_t black_yuva_color[4] = { 0, 127, 127, 0 };
 static const uint8_t black_gbrp_color[4] = { 0, 0, 0, 0 };
 
-static int config_input(AVFilterLink *inlink)
+static int config_output(AVFilterLink *outlink)
 {
-    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
-    VectorscopeContext *s = inlink->dst->priv;
-
-    s->is_yuv = !(desc->flags & AV_PIX_FMT_FLAG_RGB);
+    VectorscopeContext *s = outlink->src->priv;
 
-    if (s->mode == GRAY && s->is_yuv)
-        s->pd = 0;
-    else {
-        if ((s->x == 1 && s->y == 2) || (s->x == 2 && s->y == 1))
-            s->pd = 0;
-        else if ((s->x == 0 && s->y == 2) || (s->x == 2 && s->y == 0))
-            s->pd = 1;
-        else if ((s->x == 0 && s->y == 1) || (s->x == 1 && s->y == 0))
-            s->pd = 2;
-    }
+    s->intensity = s->fintensity * (s->size - 1);
+    outlink->h = outlink->w = s->size;
+    outlink->sample_aspect_ratio = (AVRational){1,1};
+    return 0;
+}
 
-    switch (inlink->format) {
-    case AV_PIX_FMT_GBRAP:
-    case AV_PIX_FMT_GBRP:
-        s->bg_color = black_gbrp_color;
-        break;
-    default:
-        s->bg_color = black_yuva_color;
-    }
+static void envelope_instant16(VectorscopeContext *s, AVFrame *out)
+{
+    const int dlinesize = out->linesize[0] / 2;
+    uint16_t *dpd = s->mode == COLOR || !s->is_yuv ? (uint16_t *)out->data[s->pd] : (uint16_t *)out->data[0];
+    const int max = s->size - 1;
+    int i, j;
 
-    s->hsub = desc->log2_chroma_w;
-    s->vsub = desc->log2_chroma_h;
-    s->planeheight[1] = s->planeheight[2] = FF_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
-    s->planeheight[0] = s->planeheight[3] = inlink->h;
-    s->planewidth[1]  = s->planewidth[2]  = FF_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
-    s->planewidth[0]  = s->planewidth[3]  = inlink->w;
+    for (i = 0; i < out->height; i++) {
+        for (j = 0; j < out->width; j++) {
+            const int pos = i * dlinesize + j;
+            const int poa = (i - 1) * dlinesize + j;
+            const int pob = (i + 1) * dlinesize + j;
 
-    return 0;
+            if (dpd[pos] && (((!j || !dpd[pos - 1]) || ((j == (out->width - 1)) || !dpd[pos + 1]))
+                         || ((!i || !dpd[poa]) || ((i == (out->height - 1)) || !dpd[pob])))) {
+                dpd[pos] = max;
+            }
+        }
+    }
 }
 
-static int config_output(AVFilterLink *outlink)
+static void envelope_peak16(VectorscopeContext *s, AVFrame *out)
 {
-    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(outlink->format);
-    const int depth = desc->comp[0].depth;
-    VectorscopeContext *s = outlink->src->priv;
+    const int dlinesize = out->linesize[0] / 2;
+    uint16_t *dpd = s->mode == COLOR || !s->is_yuv ? (uint16_t *)out->data[s->pd] : (uint16_t *)out->data[0];
+    const int max = s->size - 1;
+    int i, j;
 
-    s->intensity = s->fintensity * ((1 << depth) - 1);
-    outlink->h = outlink->w = 1 << depth;
-    outlink->sample_aspect_ratio = (AVRational){1,1};
-    return 0;
+    for (i = 0; i < out->height; i++) {
+        for (j = 0; j < out->width; j++) {
+            const int pos = i * dlinesize + j;
+
+            if (dpd[pos])
+                s->peak[i][j] = 1;
+        }
+    }
+
+    if (s->envelope == 3)
+        envelope_instant16(s, out);
+
+    for (i = 0; i < out->height; i++) {
+        for (j = 0; j < out->width; j++) {
+            const int pos = i * dlinesize + j;
+
+            if (s->peak[i][j] && (((!j || !s->peak[i][j-1]) || ((j == (out->width - 1)) || !s->peak[i][j + 1]))
+                              || ((!i || !s->peak[i-1][j]) || ((i == (out->height - 1)) || !s->peak[i + 1][j])))) {
+                dpd[pos] = max;
+            }
+        }
+    }
 }
 
 static void envelope_instant(VectorscopeContext *s, AVFrame *out)
@@ -227,7 +281,7 @@ static void envelope_peak(VectorscopeContext *s, AVFrame *out)
             const int pos = i * dlinesize + j;
 
             if (dpd[pos])
-                s->peak[i][j] = 255;
+                s->peak[i][j] = 1;
         }
     }
 
@@ -246,6 +300,17 @@ static void envelope_peak(VectorscopeContext *s, AVFrame *out)
     }
 }
 
+static void envelope16(VectorscopeContext *s, AVFrame *out)
+{
+    if (!s->envelope) {
+        return;
+    } else if (s->envelope == 1) {
+        envelope_instant16(s, out);
+    } else {
+        envelope_peak16(s, out);
+    }
+}
+
 static void envelope(VectorscopeContext *s, AVFrame *out)
 {
     if (!s->envelope) {
@@ -257,7 +322,166 @@ static void envelope(VectorscopeContext *s, AVFrame *out)
     }
 }
 
-static void vectorscope(VectorscopeContext *s, AVFrame *in, AVFrame *out, int pd)
+static void vectorscope16(VectorscopeContext *s, AVFrame *in, AVFrame *out, int pd)
+{
+    const uint16_t * const *src = (const uint16_t * const *)in->data;
+    const int slinesizex = in->linesize[s->x] / 2;
+    const int slinesizey = in->linesize[s->y] / 2;
+    const int slinesized = in->linesize[pd] / 2;
+    const int dlinesize = out->linesize[0] / 2;
+    const int intensity = s->intensity;
+    const int px = s->x, py = s->y;
+    const int h = s->planeheight[py];
+    const int w = s->planewidth[px];
+    const uint16_t *spx = src[px];
+    const uint16_t *spy = src[py];
+    const uint16_t *spd = src[pd];
+    const int hsub = s->hsub;
+    const int vsub = s->vsub;
+    uint16_t **dst = (uint16_t **)out->data;
+    uint16_t *dpx = dst[px];
+    uint16_t *dpy = dst[py];
+    uint16_t *dpd = dst[pd];
+    const int max = s->size - 1;
+    const int mid = s->size / 2;
+    int i, j, k;
+
+    for (k = 0; k < 4 && dst[k]; k++) {
+        const int mult = s->mult;
+
+        for (i = 0; i < out->height ; i++)
+            for (j = 0; j < out->width; j++)
+                AV_WN16(out->data[k] + i * out->linesize[k] + j * 2,
+                        s->mode == COLOR && k == s->pd ? 0 : s->bg_color[k] * mult);
+    }
+
+    switch (s->mode) {
+    case COLOR:
+    case GRAY:
+        if (s->is_yuv) {
+            for (i = 0; i < h; i++) {
+                const int iwx = i * slinesizex;
+                const int iwy = i * slinesizey;
+                for (j = 0; j < w; j++) {
+                    const int x = FFMIN(spx[iwx + j], max);
+                    const int y = FFMIN(spy[iwy + j], max);
+                    const int pos = y * dlinesize + x;
+
+                    dpd[pos] = FFMIN(dpd[pos] + intensity, max);
+                    if (dst[3])
+                        dst[3][pos] = max;
+                }
+            }
+        } else {
+            for (i = 0; i < h; i++) {
+                const int iwx = i * slinesizex;
+                const int iwy = i * slinesizey;
+                for (j = 0; j < w; j++) {
+                    const int x = FFMIN(spx[iwx + j], max);
+                    const int y = FFMIN(spy[iwy + j], max);
+                    const int pos = y * dlinesize + x;
+
+                    dst[0][pos] = FFMIN(dst[0][pos] + intensity, max);
+                    dst[1][pos] = FFMIN(dst[1][pos] + intensity, max);
+                    dst[2][pos] = FFMIN(dst[2][pos] + intensity, max);
+                    if (dst[3])
+                        dst[3][pos] = max;
+                }
+            }
+        }
+        break;
+    case COLOR2:
+        if (s->is_yuv) {
+            for (i = 0; i < h; i++) {
+                const int iw1 = i * slinesizex;
+                const int iw2 = i * slinesizey;
+                for (j = 0; j < w; j++) {
+                    const int x = FFMIN(spx[iw1 + j], max);
+                    const int y = FFMIN(spy[iw2 + j], max);
+                    const int pos = y * dlinesize + x;
+
+                    if (!dpd[pos])
+                        dpd[pos] = FFABS(mid - x) + FFABS(mid - y);
+                    dpx[pos] = x;
+                    dpy[pos] = y;
+                    if (dst[3])
+                        dst[3][pos] = max;
+                }
+            }
+        } else {
+            for (i = 0; i < h; i++) {
+                const int iw1 = i * slinesizex;
+                const int iw2 = i * slinesizey;
+                for (j = 0; j < w; j++) {
+                    const int x = FFMIN(spx[iw1 + j], max);
+                    const int y = FFMIN(spy[iw2 + j], max);
+                    const int pos = y * dlinesize + x;
+
+                    if (!dpd[pos])
+                        dpd[pos] = FFMIN(x + y, max);
+                    dpx[pos] = x;
+                    dpy[pos] = y;
+                    if (dst[3])
+                        dst[3][pos] = max;
+                }
+            }
+        }
+        break;
+    case COLOR3:
+        for (i = 0; i < h; i++) {
+            const int iw1 = i * slinesizex;
+            const int iw2 = i * slinesizey;
+            for (j = 0; j < w; j++) {
+                const int x = FFMIN(spx[iw1 + j], max);
+                const int y = FFMIN(spy[iw2 + j], max);
+                const int pos = y * dlinesize + x;
+
+                dpd[pos] = FFMIN(max, dpd[pos] + intensity);
+                dpx[pos] = x;
+                dpy[pos] = y;
+                if (dst[3])
+                    dst[3][pos] = max;
+            }
+        }
+        break;
+    case COLOR4:
+        for (i = 0; i < in->height; i++) {
+            const int iwx = (i >> vsub) * slinesizex;
+            const int iwy = (i >> vsub) * slinesizey;
+            const int iwd = i * slinesized;
+            for (j = 0; j < in->width; j++) {
+                const int x = FFMIN(spx[iwx + (j >> hsub)], max);
+                const int y = FFMIN(spy[iwy + (j >> hsub)], max);
+                const int pos = y * dlinesize + x;
+
+                dpd[pos] = FFMAX(spd[iwd + j], dpd[pos]);
+                dpx[pos] = x;
+                dpy[pos] = y;
+                if (dst[3])
+                    dst[3][pos] = max;
+            }
+        }
+        break;
+    default:
+        av_assert0(0);
+    }
+
+    envelope16(s, out);
+
+    if (s->mode == COLOR) {
+        for (i = 0; i < out->height; i++) {
+            for (j = 0; j < out->width; j++) {
+                if (!dpd[i * dlinesize + j]) {
+                    dpx[i * dlinesize + j] = j;
+                    dpy[i * dlinesize + j] = i;
+                    dpd[i * dlinesize + j] = mid;
+                }
+            }
+        }
+    }
+}
+
+static void vectorscope8(VectorscopeContext *s, AVFrame *in, AVFrame *out, int pd)
 {
     const uint8_t * const *src = (const uint8_t * const *)in->data;
     const int slinesizex = in->linesize[s->x];
@@ -277,7 +501,12 @@ static void vectorscope(VectorscopeContext *s, AVFrame *in, AVFrame *out, int pd
     uint8_t *dpx = dst[px];
     uint8_t *dpy = dst[py];
     uint8_t *dpd = dst[pd];
-    int i, j;
+    int i, j, k;
+
+    for (k = 0; k < 4 && dst[k]; k++)
+        for (i = 0; i < out->height ; i++)
+            memset(dst[k] + i * out->linesize[k],
+                   s->mode == COLOR && k == s->pd ? 0 : s->bg_color[k], out->width);
 
     switch (s->mode) {
     case COLOR:
@@ -411,8 +640,6 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
     VectorscopeContext *s = ctx->priv;
     AVFilterLink *outlink = ctx->outputs[0];
     AVFrame *out;
-    uint8_t **dst;
-    int i, k;
 
     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
     if (!out) {
@@ -420,19 +647,59 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
         return AVERROR(ENOMEM);
     }
     out->pts = in->pts;
-    dst = out->data;
-
-    for (k = 0; k < 4 && dst[k]; k++)
-        for (i = 0; i < outlink->h ; i++)
-            memset(dst[k] + i * out->linesize[k],
-                   s->mode == COLOR && k == s->pd ? 0 : s->bg_color[k], outlink->w);
 
-    vectorscope(s, in, out, s->pd);
+    s->vectorscope(s, in, out, s->pd);
 
     av_frame_free(&in);
     return ff_filter_frame(outlink, out);
 }
 
+static int config_input(AVFilterLink *inlink)
+{
+    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
+    VectorscopeContext *s = inlink->dst->priv;
+
+    s->is_yuv = !(desc->flags & AV_PIX_FMT_FLAG_RGB);
+    s->size = 1 << desc->comp[0].depth;
+    s->mult = s->size / 256;
+
+    if (s->mode == GRAY && s->is_yuv)
+        s->pd = 0;
+    else {
+        if ((s->x == 1 && s->y == 2) || (s->x == 2 && s->y == 1))
+            s->pd = 0;
+        else if ((s->x == 0 && s->y == 2) || (s->x == 2 && s->y == 0))
+            s->pd = 1;
+        else if ((s->x == 0 && s->y == 1) || (s->x == 1 && s->y == 0))
+            s->pd = 2;
+    }
+
+    if (s->size == 256)
+        s->vectorscope = vectorscope8;
+    else
+        s->vectorscope = vectorscope16;
+
+    switch (inlink->format) {
+    case AV_PIX_FMT_GBRP10:
+    case AV_PIX_FMT_GBRP9:
+    case AV_PIX_FMT_GBRAP:
+    case AV_PIX_FMT_GBRP:
+        s->bg_color = black_gbrp_color;
+        break;
+    default:
+        s->bg_color = black_yuva_color;
+    }
+
+    s->hsub = desc->log2_chroma_w;
+    s->vsub = desc->log2_chroma_h;
+    s->planeheight[1] = s->planeheight[2] = FF_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
+    s->planeheight[0] = s->planeheight[3] = inlink->h;
+    s->planewidth[1]  = s->planewidth[2]  = FF_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
+    s->planewidth[0]  = s->planewidth[3]  = inlink->w;
+
+    return 0;
+}
+
 static const AVFilterPad inputs[] = {
     {
         .name         = "default",