avfilter: add rgbashift filter

1 files changed, 223 insertions, 9 deletions

diff --git a/libavfilter/vf_chromashift.c b/libavfilter/vf_chromashift.c
index fe1c79e5a7..068c3c1b68 100644
--- a/libavfilter/vf_chromashift.c
+++ b/libavfilter/vf_chromashift.c
@@ -35,8 +35,13 @@ typedef struct ChromaShiftContext {
     const AVClass *class;
     int cbh, cbv;
     int crh, crv;
+    int rh, rv;
+    int gh, gv;
+    int bh, bv;
+    int ah, av;
     int edge;
 
+    int nb_planes;
     int depth;
     int height[4];
     int width[4];
@@ -44,12 +49,13 @@ typedef struct ChromaShiftContext {
 
     AVFrame *in;
 
+    int is_rgbashift;
     int (*filter_slice)(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
 } ChromaShiftContext;
 
 static int query_formats(AVFilterContext *ctx)
 {
-    static const enum AVPixelFormat pix_fmts[] = {
+    static const enum AVPixelFormat yuv_pix_fmts[] = {
         AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUVA420P,
         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ422P,AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ411P,
         AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P,
@@ -62,6 +68,19 @@ static int query_formats(AVFilterContext *ctx)
         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
         AV_PIX_FMT_NONE
     };
+    static const enum AVPixelFormat rgb_pix_fmts[] = {
+        AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRP9,
+        AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRP12,
+        AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
+        AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
+        AV_PIX_FMT_NONE
+    };
+    const enum AVPixelFormat *pix_fmts;
+
+    if (!strcmp(ctx->filter->name, "rgbashift"))
+        pix_fmts = rgb_pix_fmts;
+    else
+        pix_fmts = yuv_pix_fmts;
 
     AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
     if (!fmts_list)
@@ -166,6 +185,162 @@ static int wrap_slice ## depth(AVFilterContext *ctx, void *arg, int jobnr, int n
 DEFINE_WRAP(8, uint8_t, 1)
 DEFINE_WRAP(16, uint16_t, 2)
 
+#define DEFINE_RGBASMEAR(depth, type, div)                                                    \
+static int rgbasmear_slice ## depth(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)  \
+{                                                                                         \
+    ChromaShiftContext *s = ctx->priv;                                                    \
+    AVFrame *in = s->in;                                                                  \
+    AVFrame *out = arg;                                                                   \
+    const int srlinesize = in->linesize[2] / div;                                         \
+    const int sglinesize = in->linesize[0] / div;                                         \
+    const int sblinesize = in->linesize[1] / div;                                         \
+    const int salinesize = in->linesize[3] / div;                                         \
+    const int rlinesize = out->linesize[2] / div;                                         \
+    const int glinesize = out->linesize[0] / div;                                         \
+    const int blinesize = out->linesize[1] / div;                                         \
+    const int alinesize = out->linesize[3] / div;                                         \
+    const int rh = s->rh;                                                                 \
+    const int rv = s->rv;                                                                 \
+    const int gh = s->gh;                                                                 \
+    const int gv = s->gv;                                                                 \
+    const int bh = s->bh;                                                                 \
+    const int bv = s->bv;                                                                 \
+    const int ah = s->ah;                                                                 \
+    const int av = s->av;                                                                 \
+    const int h = s->height[1];                                                           \
+    const int w = s->width[1];                                                            \
+    const int slice_start = (h * jobnr) / nb_jobs;                                        \
+    const int slice_end = (h * (jobnr+1)) / nb_jobs;                                      \
+    const type *sr = (const type *)in->data[2];                                           \
+    const type *sg = (const type *)in->data[0];                                           \
+    const type *sb = (const type *)in->data[1];                                           \
+    const type *sa = (const type *)in->data[3];                                           \
+    type *dr = (type *)out->data[2] + slice_start * rlinesize;                            \
+    type *dg = (type *)out->data[0] + slice_start * glinesize;                            \
+    type *db = (type *)out->data[1] + slice_start * blinesize;                            \
+    type *da = (type *)out->data[3] + slice_start * alinesize;                            \
+                                                                                          \
+    for (int y = slice_start; y < slice_end; y++) {                                       \
+        const int ry = av_clip(y - rv, 0, h-1) * srlinesize;                              \
+        const int gy = av_clip(y - gv, 0, h-1) * sglinesize;                              \
+        const int by = av_clip(y - bv, 0, h-1) * sblinesize;                              \
+        int ay;                                                                           \
+                                                                                          \
+        for (int x = 0; x < w; x++) {                                                     \
+            dr[x] = sr[av_clip(x - rh, 0, w - 1) + ry];                                   \
+            dg[x] = sg[av_clip(x - gh, 0, w - 1) + gy];                                   \
+            db[x] = sb[av_clip(x - bh, 0, w - 1) + by];                                   \
+        }                                                                                 \
+                                                                                          \
+        dr += rlinesize;                                                                  \
+        dg += glinesize;                                                                  \
+        db += blinesize;                                                                  \
+                                                                                          \
+        if (s->nb_planes < 4)                                                             \
+            continue;                                                                     \
+        ay = av_clip(y - av, 0, h-1) * salinesize;                                        \
+        for (int x = 0; x < w; x++) {                                                     \
+            da[x] = sa[av_clip(x - ah, 0, w - 1) + ay];                                   \
+        }                                                                                 \
+                                                                                          \
+        da += alinesize;                                                                  \
+    }                                                                                     \
+                                                                                          \
+    return 0;                                                                             \
+}
+
+DEFINE_RGBASMEAR(8, uint8_t, 1)
+DEFINE_RGBASMEAR(16, uint16_t, 2)
+
+#define DEFINE_RGBAWRAP(depth, type, div)                                                     \
+static int rgbawrap_slice ## depth(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)   \
+{                                                                                         \
+    ChromaShiftContext *s = ctx->priv;                                                    \
+    AVFrame *in = s->in;                                                                  \
+    AVFrame *out = arg;                                                                   \
+    const int srlinesize = in->linesize[2] / div;                                         \
+    const int sglinesize = in->linesize[0] / div;                                         \
+    const int sblinesize = in->linesize[1] / div;                                         \
+    const int salinesize = in->linesize[3] / div;                                         \
+    const int rlinesize = out->linesize[2] / div;                                         \
+    const int glinesize = out->linesize[0] / div;                                         \
+    const int blinesize = out->linesize[1] / div;                                         \
+    const int alinesize = out->linesize[3] / div;                                         \
+    const int rh = s->rh;                                                                 \
+    const int rv = s->rv;                                                                 \
+    const int gh = s->gh;                                                                 \
+    const int gv = s->gv;                                                                 \
+    const int bh = s->bh;                                                                 \
+    const int bv = s->bv;                                                                 \
+    const int ah = s->ah;                                                                 \
+    const int av = s->av;                                                                 \
+    const int h = s->height[1];                                                           \
+    const int w = s->width[1];                                                            \
+    const int slice_start = (h * jobnr) / nb_jobs;                                        \
+    const int slice_end = (h * (jobnr+1)) / nb_jobs;                                      \
+    const type *sr = (const type *)in->data[2];                                           \
+    const type *sg = (const type *)in->data[0];                                           \
+    const type *sb = (const type *)in->data[1];                                           \
+    const type *sa = (const type *)in->data[3];                                           \
+    type *dr = (type *)out->data[2] + slice_start * rlinesize;                            \
+    type *dg = (type *)out->data[0] + slice_start * glinesize;                            \
+    type *db = (type *)out->data[1] + slice_start * blinesize;                            \
+    type *da = (type *)out->data[3] + slice_start * alinesize;                            \
+                                                                                          \
+    for (int y = slice_start; y < slice_end; y++) {                                       \
+        int ry = (y - rv) % h;                                                            \
+        int gy = (y - gv) % h;                                                            \
+        int by = (y - bv) % h;                                                            \
+                                                                                          \
+        if (ry < 0)                                                                       \
+            ry += h;                                                                      \
+        if (gy < 0)                                                                       \
+            gy += h;                                                                      \
+        if (by < 0)                                                                       \
+            by += h;                                                                      \
+                                                                                          \
+        for (int x = 0; x < w; x++) {                                                     \
+            int rx = (x - rh) % w;                                                        \
+            int gx = (x - gh) % w;                                                        \
+            int bx = (x - bh) % w;                                                        \
+                                                                                          \
+            if (rx < 0)                                                                   \
+                rx += w;                                                                  \
+            if (gx < 0)                                                                   \
+                gx += w;                                                                  \
+            if (bx < 0)                                                                   \
+                bx += w;                                                                  \
+            dr[x] = sr[rx + ry * srlinesize];                                             \
+            dg[x] = sg[gx + gy * sglinesize];                                             \
+            db[x] = sb[bx + by * sblinesize];                                             \
+        }                                                                                 \
+                                                                                          \
+        dr += rlinesize;                                                                  \
+        dg += glinesize;                                                                  \
+        db += blinesize;                                                                  \
+                                                                                          \
+        if (s->nb_planes < 4)                                                             \
+            continue;                                                                     \
+        for (int x = 0; x < w; x++) {                                                     \
+            int ax = (x - ah) % w;                                                        \
+            int ay = (x - av) % h;                                                        \
+                                                                                          \
+            if (ax < 0)                                                                   \
+                ax += w;                                                                  \
+            if (ay < 0)                                                                   \
+                ay += h;                                                                  \
+            da[x] = sa[ax + ay * salinesize];                                             \
+        }                                                                                 \
+                                                                                          \
+        da += alinesize;                                                                  \
+    }                                                                                     \
+                                                                                          \
+    return 0;                                                                             \
+}
+
+DEFINE_RGBAWRAP(8, uint8_t, 1)
+DEFINE_RGBAWRAP(16, uint16_t, 2)
+
 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 {
     AVFilterContext *ctx = inlink->dst;
@@ -181,10 +356,12 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
     av_frame_copy_props(out, in);
 
     s->in = in;
-    av_image_copy_plane(out->data[0] + out->linesize[0],
-                        out->linesize[0],
-                        in->data[0], in->linesize[0],
-                        s->linesize[0], s->height[0]);
+    if (!s->is_rgbashift) {
+        av_image_copy_plane(out->data[0] + out->linesize[0],
+                            out->linesize[0],
+                            in->data[0], in->linesize[0],
+                            s->linesize[0], s->height[0]);
+    }
     ctx->internal->execute(ctx, s->filter_slice, out, NULL,
                            FFMIN3(s->height[1],
                                   s->height[2],
@@ -200,11 +377,20 @@ static int config_input(AVFilterLink *inlink)
     ChromaShiftContext *s = ctx->priv;
     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
 
+    s->is_rgbashift = !strcmp(ctx->filter->name, "rgbashift");
     s->depth = desc->comp[0].depth;
-    if (s->edge)
-        s->filter_slice = s->depth > 8 ? wrap_slice16 : wrap_slice8;
-    else
-        s->filter_slice = s->depth > 8 ? smear_slice16 : smear_slice8;
+    s->nb_planes = desc->nb_components;
+    if (s->is_rgbashift) {
+        if (s->edge)
+            s->filter_slice = s->depth > 8 ? rgbawrap_slice16 : rgbawrap_slice8;
+        else
+            s->filter_slice = s->depth > 8 ? rgbasmear_slice16 : rgbasmear_slice8;
+    } else {
+        if (s->edge)
+            s->filter_slice = s->depth > 8 ? wrap_slice16 : wrap_slice8;
+        else
+            s->filter_slice = s->depth > 8 ? smear_slice16 : smear_slice8;
+    }
     s->height[1] = s->height[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
     s->height[0] = s->height[3] = inlink->h;
     s->width[1] = s->width[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
@@ -257,3 +443,31 @@ AVFilter ff_vf_chromashift = {
     .inputs        = inputs,
     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
 };
+
+static const AVOption rgbashift_options[] = {
+    { "rh", "shift red horizontally",   OFFSET(rh),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "rv", "shift red vertically",     OFFSET(rv),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "gh", "shift green horizontally", OFFSET(gh),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "gv", "shift green vertically",   OFFSET(gv),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "bh", "shift blue horizontally",  OFFSET(bh),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "bv", "shift blue vertically",    OFFSET(bv),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "ah", "shift alpha horizontally", OFFSET(ah),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "av", "shift alpha vertically",   OFFSET(av),   AV_OPT_TYPE_INT,   {.i64=0}, -255, 255, .flags = VF },
+    { "edge", "set edge operation",     OFFSET(edge), AV_OPT_TYPE_INT,   {.i64=0},    0,   1, .flags = VF, "edge" },
+    { "smear",                          0,         0, AV_OPT_TYPE_CONST, {.i64=0},    0,   0, .flags = VF, "edge" },
+    { "wrap",                           0,         0, AV_OPT_TYPE_CONST, {.i64=1},    0,   0, .flags = VF, "edge" },
+    { NULL },
+};
+
+AVFILTER_DEFINE_CLASS(rgbashift);
+
+AVFilter ff_vf_rgbashift = {
+    .name          = "rgbashift",
+    .description   = NULL_IF_CONFIG_SMALL("Shift RGBA."),
+    .priv_size     = sizeof(ChromaShiftContext),
+    .priv_class    = &rgbashift_class,
+    .query_formats = query_formats,
+    .outputs       = outputs,
+    .inputs        = inputs,
+    .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
+};