avfilter/vf_edgedetect: add more formats support to canny mode

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

1 files changed, 28 insertions, 20 deletions

diff --git a/libavfilter/vf_edgedetect.c b/libavfilter/vf_edgedetect.c
index 476c434349..a0ddcbbf5c 100644
--- a/libavfilter/vf_edgedetect.c
+++ b/libavfilter/vf_edgedetect.c
@@ -52,6 +52,7 @@ struct plane_info {
     uint8_t  *tmpbuf;
     uint16_t *gradients;
     char     *directions;
+    int      width, height;
 };
 
 typedef struct EdgeDetectContext {
@@ -98,7 +99,7 @@ static int query_formats(AVFilterContext *ctx)
 {
     const EdgeDetectContext *edgedetect = ctx->priv;
     static const enum AVPixelFormat wires_pix_fmts[] = {AV_PIX_FMT_GRAY8, AV_PIX_FMT_NONE};
-    static const enum AVPixelFormat canny_pix_fmts[] = {AV_PIX_FMT_YUV444P, AV_PIX_FMT_GBRP, AV_PIX_FMT_GRAY8, AV_PIX_FMT_NONE};
+    static const enum AVPixelFormat canny_pix_fmts[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_GBRP, AV_PIX_FMT_GRAY8, AV_PIX_FMT_NONE};
     static const enum AVPixelFormat colormix_pix_fmts[] = {AV_PIX_FMT_GBRP, AV_PIX_FMT_GRAY8, AV_PIX_FMT_NONE};
     AVFilterFormats *fmts_list;
     const enum AVPixelFormat *pix_fmts = NULL;
@@ -123,14 +124,19 @@ static int config_props(AVFilterLink *inlink)
     int p;
     AVFilterContext *ctx = inlink->dst;
     EdgeDetectContext *edgedetect = ctx->priv;
+    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
 
     edgedetect->nb_planes = inlink->format == AV_PIX_FMT_GRAY8 ? 1 : 3;
     for (p = 0; p < edgedetect->nb_planes; p++) {
         struct plane_info *plane = &edgedetect->planes[p];
-
-        plane->tmpbuf     = av_malloc(inlink->w * inlink->h);
-        plane->gradients  = av_calloc(inlink->w * inlink->h, sizeof(*plane->gradients));
-        plane->directions = av_malloc(inlink->w * inlink->h);
+        int vsub = p ? desc->log2_chroma_h : 0;
+        int hsub = p ? desc->log2_chroma_w : 0;
+
+        plane->width      = AV_CEIL_RSHIFT(inlink->w, hsub);
+        plane->height     = AV_CEIL_RSHIFT(inlink->h, vsub);
+        plane->tmpbuf     = av_malloc(plane->width * plane->height);
+        plane->gradients  = av_calloc(plane->width * plane->height, sizeof(*plane->gradients));
+        plane->directions = av_malloc(plane->width * plane->height);
         if (!plane->tmpbuf || !plane->gradients || !plane->directions)
             return AVERROR(ENOMEM);
     }
@@ -338,42 +344,44 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
         uint8_t  *tmpbuf     = plane->tmpbuf;
         uint16_t *gradients  = plane->gradients;
         int8_t   *directions = plane->directions;
+        const int width      = plane->width;
+        const int height     = plane->height;
 
         if (!((1 << p) & edgedetect->filter_planes)) {
             if (!direct)
                 av_image_copy_plane(out->data[p], out->linesize[p],
                                     in->data[p], in->linesize[p],
-                                    inlink->w, inlink->h);
+                                    width, height);
             continue;
         }
 
         /* gaussian filter to reduce noise  */
-        gaussian_blur(ctx, inlink->w, inlink->h,
-                      tmpbuf,      inlink->w,
+        gaussian_blur(ctx, width, height,
+                      tmpbuf,      width,
                       in->data[p], in->linesize[p]);
 
         /* compute the 16-bits gradients and directions for the next step */
-        sobel(inlink->w, inlink->h,
-              gradients, inlink->w,
-              directions,inlink->w,
-              tmpbuf,    inlink->w);
+        sobel(width, height,
+              gradients, width,
+              directions,width,
+              tmpbuf,    width);
 
         /* non_maximum_suppression() will actually keep & clip what's necessary and
          * ignore the rest, so we need a clean output buffer */
-        memset(tmpbuf, 0, inlink->w * inlink->h);
-        non_maximum_suppression(inlink->w, inlink->h,
-                                tmpbuf,    inlink->w,
-                                directions,inlink->w,
-                                gradients, inlink->w);
+        memset(tmpbuf, 0, width * height);
+        non_maximum_suppression(width, height,
+                                tmpbuf,    width,
+                                directions,width,
+                                gradients, width);
 
         /* keep high values, or low values surrounded by high values */
         double_threshold(edgedetect->low_u8, edgedetect->high_u8,
-                         inlink->w, inlink->h,
+                         width, height,
                          out->data[p], out->linesize[p],
-                         tmpbuf,       inlink->w);
+                         tmpbuf,       width);
 
         if (edgedetect->mode == MODE_COLORMIX) {
-            color_mix(inlink->w, inlink->h,
+            color_mix(width, height,
                       out->data[p], out->linesize[p],
                       in->data[p], in->linesize[p]);
         }