git.sesse.net Git - ffmpeg/blob - libavfilter/vf_yaepblur.c

   1 /*
   2  * Copyright (C) 2019 Leo Zhang <leozhang@qiyi.com>
   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 /**
  22  * @file
  23  * yaep(yet another edge preserving) blur filter
  24  *
  25  * This implementation is based on an algorithm described in
  26  * "J. S. Lee, Digital image enhancement and noise filtering by use of local statistics, IEEE Trans. Pattern
  27  * Anal. Mach. Intell. PAMI-2, 1980."
  28  */
  29
  30 #include "libavutil/opt.h"
  31 #include "libavutil/imgutils.h"
  32 #include "avfilter.h"
  33 #include "internal.h"
  34
  35 typedef struct YAEPContext {
  36     const AVClass *class;
  37
  38     int planes;
  39     int radius;
  40     int sigma;
  41
  42     int nb_planes;
  43     int planewidth[4];
  44     int planeheight[4];
  45     int depth;
  46
  47     uint64_t *sat;        ///< summed area table
  48     uint64_t *square_sat; ///< square summed area table
  49     int sat_linesize;
  50
  51     int (*pre_calculate_row)(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
  52     int (*filter_slice     )(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
  53 } YAEPContext;
  54
  55 static av_cold void uninit(AVFilterContext *ctx)
  56 {
  57     YAEPContext *s = ctx->priv;
  58     av_freep(&s->sat);
  59     av_freep(&s->square_sat);
  60 }
  61
  62 static int query_formats(AVFilterContext *ctx)
  63 {
  64     static const enum AVPixelFormat pix_fmts[] = {
  65         AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV440P,
  66         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P,
  67         AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P,
  68         AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P,
  69         AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P,
  70         AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9,
  71         AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
  72         AV_PIX_FMT_YUV420P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV440P12,
  73         AV_PIX_FMT_YUV420P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV444P14,
  74         AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16,
  75         AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9,
  76         AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10,
  77         AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12,
  78         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
  79         AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
  80         AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
  81         AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
  82         AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY14, AV_PIX_FMT_GRAY16,
  83         AV_PIX_FMT_NONE
  84     };
  85
  86     return ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
  87 }
  88
  89 typedef struct ThreadData {
  90     int width;
  91     int height;
  92     int src_linesize;
  93     int dst_linesize;
  94     uint8_t *src;
  95     uint8_t *dst;
  96 } ThreadData;
  97
  98 #define PRE_CALCULATE_ROW(type, name)                                    \
  99 static int pre_calculate_row_##name(AVFilterContext *ctx, void *arg,     \
 100                                    int jobnr, int nb_jobs)               \
 101 {                                                                        \
 102     ThreadData *td = arg;                                                \
 103     YAEPContext *s = ctx->priv;                                          \
 104                                                                          \
 105     const int width        = td->width;                                  \
 106     const int height       = td->height;                                 \
 107     const int linesize     = td->src_linesize / sizeof(type);            \
 108     const int sat_linesize = s->sat_linesize;                            \
 109                                                                          \
 110     const int starty = height * jobnr     / nb_jobs;                     \
 111     const int endy   = height * (jobnr+1) / nb_jobs;                     \
 112                                                                          \
 113     uint64_t *sat        = s->sat + (starty + 1) * sat_linesize;         \
 114     uint64_t *square_sat = s->square_sat + (starty + 1) * sat_linesize;  \
 115     const type *src      = (const type *)td->src + starty * linesize;    \
 116                                                                          \
 117     int x, y;                                                            \
 118                                                                          \
 119     for (y = starty; y < endy; y++) {                                    \
 120         for (x = 0; x < width; x++) {                                    \
 121             sat[x+1]        = sat[x] + src[x];                           \
 122             square_sat[x+1] = square_sat[x] + (uint64_t)src[x] * src[x]; \
 123         }                                                                \
 124         sat               += sat_linesize;                               \
 125         square_sat        += sat_linesize;                               \
 126         src               += linesize;                                   \
 127     }                                                                    \
 128                                                                          \
 129     return 0;                                                            \
 130 }
 131
 132 PRE_CALCULATE_ROW(uint8_t,  byte)
 133 PRE_CALCULATE_ROW(uint16_t, word)
 134
 135 static int pre_calculate_col(AVFilterContext *ctx, void *arg,
 136                              int jobnr, int nb_jobs)
 137 {
 138     ThreadData *td = arg;
 139     YAEPContext *s = ctx->priv;
 140
 141     const int width        = td->width;
 142     const int height       = td->height;
 143     const int sat_linesize = s->sat_linesize;
 144
 145     const int startx = width * jobnr       / nb_jobs;
 146     const int endx   = width * (jobnr + 1) / nb_jobs;
 147
 148     uint64_t *sat, *square_sat;
 149     int x, y;
 150
 151     for (x = startx; x < endx; x++) {
 152         sat = s->sat + x + 1;
 153         square_sat = s->square_sat + x + 1;
 154         for (y = 0; y < height; y++) {
 155             *(sat+sat_linesize)        += *sat;
 156             *(square_sat+sat_linesize) += *square_sat;
 157             sat         += sat_linesize;
 158             square_sat  += sat_linesize;
 159         }
 160     }
 161
 162     return 0;
 163 }
 164
 165 #define FILTER_SLICE(type, name)                                                                          \
 166 static int filter_slice_##name(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)                   \
 167 {                                                                                                         \
 168     ThreadData *td = arg;                                                                                 \
 169     YAEPContext *s = ctx->priv;                                                                           \
 170                                                                                                           \
 171     const int width = td->width;                                                                          \
 172     const int height = td->height;                                                                        \
 173     const int src_linesize = td->src_linesize / sizeof(type);                                             \
 174     const int dst_linesize = td->dst_linesize / sizeof(type);                                             \
 175     const int sat_linesize = s->sat_linesize;                                                             \
 176     const int sigma = s->sigma;                                                                           \
 177     const int radius = s->radius;                                                                         \
 178                                                                                                           \
 179     uint64_t *sat = s->sat;                                                                               \
 180     uint64_t *square_sat = s->square_sat;                                                                 \
 181     const type *src = (const type *)td->src;                                                              \
 182     type *dst = (type *)td->dst;                                                                          \
 183                                                                                                           \
 184     const int starty = height * jobnr       / nb_jobs;                                                    \
 185     const int endy   = height * (jobnr + 1) / nb_jobs;                                                    \
 186                                                                                                           \
 187     int x, y;                                                                                             \
 188     int lower_x, higher_x;                                                                                \
 189     int lower_y, higher_y;                                                                                \
 190     int dist_y, count;                                                                                    \
 191     uint64_t sum, square_sum, mean, var;                                                                  \
 192                                                                                                           \
 193     for (y = starty; y < endy; y++) {                                                                     \
 194         lower_y  = y - radius     < 0      ? 0      : y - radius;                                         \
 195         higher_y = y + radius + 1 > height ? height : y + radius + 1;                                     \
 196         dist_y = higher_y - lower_y;                                                                      \
 197         for (x = 0; x < width; x++) {                                                                     \
 198             lower_x  = x - radius     < 0     ? 0     : x - radius;                                       \
 199             higher_x = x + radius + 1 > width ? width : x + radius + 1;                                   \
 200             count = dist_y * (higher_x - lower_x);                                                        \
 201             sum = sat[higher_y * sat_linesize + higher_x]                                                 \
 202                 - sat[higher_y * sat_linesize + lower_x]                                                  \
 203                 - sat[lower_y  * sat_linesize + higher_x]                                                 \
 204                 + sat[lower_y  * sat_linesize + lower_x];                                                 \
 205             square_sum = square_sat[higher_y * sat_linesize + higher_x]                                   \
 206                        - square_sat[higher_y * sat_linesize + lower_x]                                    \
 207                        - square_sat[lower_y  * sat_linesize + higher_x]                                   \
 208                        + square_sat[lower_y  * sat_linesize + lower_x];                                   \
 209             mean = sum / count;                                                                           \
 210             var = (square_sum - sum * sum / count) / count;                                               \
 211             dst[y * dst_linesize + x] = (sigma * mean + var * src[y * src_linesize + x]) / (sigma + var); \
 212         }                                                                                                 \
 213     }                                                                                                     \
 214     return 0;                                                                                             \
 215 }
 216
 217 FILTER_SLICE(uint8_t,  byte)
 218 FILTER_SLICE(uint16_t, word)
 219
 220 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 221 {
 222     AVFilterContext *ctx = inlink->dst;
 223     YAEPContext *s = ctx->priv;
 224     AVFilterLink *outlink = ctx->outputs[0];
 225     AVFrame *out;
 226     int plane;
 227     const int nb_threads = ff_filter_get_nb_threads(ctx);
 228     ThreadData td;
 229
 230     if (av_frame_is_writable(in)) {
 231         out = in;
 232     } else {
 233         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 234         if (!out) {
 235             av_frame_free(&in);
 236             return AVERROR(ENOMEM);
 237         }
 238         av_frame_copy_props(out, in);
 239     }
 240
 241     for (plane = 0; plane < s->nb_planes; plane++) {
 242         if (!s->radius || !(s->planes & (1<<plane))) {
 243             if (out != in) {
 244                 av_image_copy_plane(out->data[plane], out->linesize[plane],
 245                                     in->data[plane], in->linesize[plane],
 246                                     s->planewidth[plane] * ((s->depth + 7) / 8),
 247                                     s->planeheight[plane]);
 248             }
 249             continue;
 250         }
 251
 252         td.width        = s->planewidth[plane];
 253         td.height       = s->planeheight[plane];
 254         td.src          = in->data[plane];
 255         td.src_linesize = in->linesize[plane];
 256         ctx->internal->execute(ctx, s->pre_calculate_row, &td, NULL, FFMIN(td.height, nb_threads));
 257         ctx->internal->execute(ctx, pre_calculate_col, &td, NULL, FFMIN(td.width,  nb_threads));
 258
 259         td.dst          = out->data[plane];
 260         td.dst_linesize = out->linesize[plane];
 261         ctx->internal->execute(ctx, s->filter_slice, &td, NULL, FFMIN(td.height, nb_threads));
 262     }
 263
 264     if (out != in)
 265         av_frame_free(&in);
 266
 267     return ff_filter_frame(outlink, out);
 268 }
 269
 270 static int config_input(AVFilterLink *inlink)
 271 {
 272     YAEPContext *s = inlink->dst->priv;
 273     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
 274
 275     s->depth = desc->comp[0].depth;
 276     s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
 277     s->planewidth[0] = s->planewidth[3] = inlink->w;
 278     s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
 279     s->planeheight[0] = s->planeheight[3] = inlink->h;
 280     s->nb_planes = av_pix_fmt_count_planes(inlink->format);
 281
 282     s->radius = FFMIN(s->radius, AV_CEIL_RSHIFT(FFMIN(inlink->w, inlink->h), 1));
 283
 284     if (s->depth <= 8) {
 285         s->pre_calculate_row = pre_calculate_row_byte;
 286         s->filter_slice      = filter_slice_byte;
 287     } else {
 288         s->pre_calculate_row = pre_calculate_row_word;
 289         s->filter_slice      = filter_slice_word;
 290     }
 291
 292     // padding one row on the top, and padding one col on the left, that is why + 1 below
 293     s->sat_linesize = inlink->w + 1;
 294     s->sat = av_mallocz_array(inlink->h + 1, s->sat_linesize*sizeof(*s->sat));
 295     if (!s->sat)
 296         return AVERROR(ENOMEM);
 297
 298     s->square_sat = av_mallocz_array(inlink->h + 1, s->sat_linesize*sizeof(*s->square_sat));
 299     if (!s->square_sat)
 300         return AVERROR(ENOMEM);
 301
 302     return 0;
 303 }
 304
 305 static const AVFilterPad yaep_inputs[] = {
 306     {
 307         .name = "default",
 308         .type = AVMEDIA_TYPE_VIDEO,
 309         .config_props = config_input,
 310         .filter_frame = filter_frame,
 311     },
 312     { NULL }
 313 };
 314
 315 static const AVFilterPad yaep_outputs[] = {
 316     {
 317         .name = "default",
 318         .type = AVMEDIA_TYPE_VIDEO,
 319     },
 320     { NULL }
 321 };
 322
 323 #define OFFSET(x) offsetof(YAEPContext, x)
 324 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
 325
 326 static const AVOption yaepblur_options[] = {
 327     { "radius", "set window radius",    OFFSET(radius), AV_OPT_TYPE_INT, {.i64=3},   0, INT_MAX, .flags=FLAGS },
 328     { "r"     , "set window radius",    OFFSET(radius), AV_OPT_TYPE_INT, {.i64=3},   0, INT_MAX, .flags=FLAGS },
 329     { "planes", "set planes to filter", OFFSET(planes), AV_OPT_TYPE_INT, {.i64=1},   0,     0xF, .flags=FLAGS },
 330     { "p",      "set planes to filter", OFFSET(planes), AV_OPT_TYPE_INT, {.i64=1},   0,     0xF, .flags=FLAGS },
 331     { "sigma",  "set blur strength",    OFFSET(sigma),  AV_OPT_TYPE_INT, {.i64=128}, 1, INT_MAX, .flags=FLAGS },
 332     { "s",      "set blur strength",    OFFSET(sigma),  AV_OPT_TYPE_INT, {.i64=128}, 1, INT_MAX, .flags=FLAGS },
 333     { NULL }
 334 };
 335
 336 AVFILTER_DEFINE_CLASS(yaepblur);
 337
 338 const AVFilter ff_vf_yaepblur = {
 339     .name            = "yaepblur",
 340     .description     = NULL_IF_CONFIG_SMALL("Yet another edge preserving blur filter."),
 341     .priv_size       = sizeof(YAEPContext),
 342     .priv_class      = &yaepblur_class,
 343     .uninit          = uninit,
 344     .query_formats   = query_formats,
 345     .inputs          = yaep_inputs,
 346     .outputs         = yaep_outputs,
 347     .flags           = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
 348     .process_command = ff_filter_process_command,
 349 };