From patchwork Fri May 11 14:03:07 2018 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Paul B Mahol X-Patchwork-Id: 8922 Delivered-To: ffmpegpatchwork@gmail.com Received: by 2002:a02:155:0:0:0:0:0 with SMTP id c82-v6csp727978jad; Fri, 11 May 2018 07:10:20 -0700 (PDT) X-Google-Smtp-Source: AB8JxZrRsg8a0b7lbBhcLl4FuKP26Ffrd+4z8A6FwFyQYzvk6UkXKMTXVGjABeQgG0m8hpzm0M7J X-Received: by 2002:adf:a6f8:: with SMTP id t111-v6mr4651687wrc.161.1526047820305; Fri, 11 May 2018 07:10:20 -0700 (PDT) ARC-Seal: i=1; a=rsa-sha256; t=1526047820; cv=none; d=google.com; s=arc-20160816; b=cAMEW9gAeZUvDQRmdo21wbhcxkHam3OQCZzSEmzdR9tjzfFlP4OqXwF0pQqUmM+x48 YZpssa7mG+nhM44P2FxPeDCSzp5rZ3ynKKLBaZFfqs82ZjUDq8H7a94j0NkQzSvDFBIX RJskFFCHjBowAshyGwUd7BR6KWz0cK55SIvi3tMcchxRe0n3m6nJuiV3FaTsBlmnj3qM rnJ/HXPaoWPvMHs/5ppX0NdFNyHmZw+iqac6qgvz3tyBp6AA7RINWXiDuUCA+j32gZhe u/zQCf1dtdvgpNUwfgFVjifkNkdDKFVsyAim4RpSr2rzkCL9a8HSR79PuG00I/CcxjYg uljw== ARC-Message-Signature: i=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=arc-20160816; h=sender:errors-to:content-transfer-encoding:mime-version:reply-to :list-subscribe:list-help:list-post:list-archive:list-unsubscribe :list-id:precedence:subject:message-id:date:to:from:dkim-signature :delivered-to:arc-authentication-results; bh=5fGRW18GW2PWg19dYop24qhZ5UB28UYRVjj5LiVVAPU=; b=Ps4L18a+MH59tTcmcYZrcS4T9/HNuwfdIHdd+fO22H5Q7FfnuQyACfCoDa9WJwgeDf uunAtBi6t80Syb/LCzpp+Z5KYhC7XzsLuRL+kbRG4neTfbXg08nszpRiOmJ+2Mz/qz7u onV/pREBqpJ6THVUXgNPczleIDfBHMlZWHmI/OWcJV0t3bVSFNBFG/j0cXOwx2B6FE0F lkhXi4C3WvgPqvpK84PEjvk18ZF2a5busXoVOJxaV9tZTqhDUzxvMNM+vV+vBHxKv5vT iRmGWKw+5HZyeIBBEMOkDssv1YCfIkAkzLo60FiC7xTMlw2GaFovW/4YcYEbzcSA1J3Y TLbg== ARC-Authentication-Results: i=1; mx.google.com; dkim=neutral (body hash did not verify) header.i=@gmail.com header.s=20161025 header.b=VSypB93E; spf=pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) smtp.mailfrom=ffmpeg-devel-bounces@ffmpeg.org; dmarc=fail (p=NONE sp=QUARANTINE dis=NONE) header.from=gmail.com Return-Path: Received: from ffbox0-bg.mplayerhq.hu (ffbox0-bg.ffmpeg.org. [79.124.17.100]) by mx.google.com with ESMTP id f3-v6si1008659wmh.178.2018.05.11.07.10.05; Fri, 11 May 2018 07:10:20 -0700 (PDT) Received-SPF: pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) client-ip=79.124.17.100; Authentication-Results: mx.google.com; dkim=neutral (body hash did not verify) header.i=@gmail.com header.s=20161025 header.b=VSypB93E; spf=pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) smtp.mailfrom=ffmpeg-devel-bounces@ffmpeg.org; dmarc=fail (p=NONE sp=QUARANTINE dis=NONE) header.from=gmail.com Received: from [127.0.1.1] (localhost [127.0.0.1]) by ffbox0-bg.mplayerhq.hu (Postfix) with ESMTP id 4B2F468A8C9; Fri, 11 May 2018 17:09:27 +0300 (EEST) X-Original-To: ffmpeg-devel@ffmpeg.org Delivered-To: ffmpeg-devel@ffmpeg.org Received: from mail-wr0-f194.google.com (mail-wr0-f194.google.com [209.85.128.194]) by ffbox0-bg.mplayerhq.hu (Postfix) with ESMTPS id F385F68A609 for ; Fri, 11 May 2018 17:09:20 +0300 (EEST) Received: by mail-wr0-f194.google.com with SMTP id y15-v6so5479359wrg.11 for ; Fri, 11 May 2018 07:09:57 -0700 (PDT) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=gmail.com; s=20161025; h=from:to:subject:date:message-id; bh=DIxNnkuuBFMbjCbCWlVwMe7NvxxB3G7Cl6vQys3hu5I=; b=VSypB93EdDhAedSfP3OkafltjKqoPcz+eyArd2GMpuMSQKIX3Zy/5jzxyYG/yGePGB irsenvH8JiDrKY6bvzGd+Y3ExKgJ/0sNhgEx7vGMPCsBaGmYdaDkMzWMCQCb7Glzbb+9 qN1DpG4RTzfj1gjGpd+vN+yJOJ/w+pc0FX86XWhAyX9gxPQ+2kWXdoI/2XavSm7u+4D5 MKNuCL5wvNbpSFlJqIbOGL3u/pchkwxSkLLzr6iOdQ8ENEiCFUsuPKhnZXu7kjY9rTdH 6JmgFOt3NQpzKujVBWt820uByY8S5DHH1SXa2WDQ6RuvTLUQocXWbR1sTLQZnH4E3+Mi oILQ== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20161025; h=x-gm-message-state:from:to:subject:date:message-id; bh=DIxNnkuuBFMbjCbCWlVwMe7NvxxB3G7Cl6vQys3hu5I=; b=AqBIszolc09wS1ikoszdjdWfOH8A+2B9ou0B3qs5iGqCPeF5YwH3Xj1KN63l2fkh0b U/ObUgu+aBskngaGeRcBf/3+86GCwPvlVoaB0XkaWSis2nbo3kUd7nBNr6uGF78mMgr8 E++UVnRdUIcjHa8Ndh93N2BrmGoCX+Z3gAsfk/DippVis3LG/QsqpplGs9Dq8+aTEZaa V+lOO/vuWuEh0y/nQVr7Te/oZhqfu7DnY0spDsCTH/XEN7X8TfUMsJWXJGIDzoRfW4mi UY0r1+cBjnQC0ouTnhoWrwkZzc9Njza+TXQqdO9XdKfDa6zM3maLmVvDOrprhWCyRZDL mCrg== X-Gm-Message-State: ALKqPweC6M89SYaN6d5XC4GOg+KxtUqSPMsvwfU7rqJsd+bURA6/99YQ hLm1CeuIPDMX92DyJCevH0CiHQ== X-Received: by 2002:adf:e4c8:: with SMTP id v8-v6mr4255859wrm.133.1526047404791; Fri, 11 May 2018 07:03:24 -0700 (PDT) Received: from localhost.localdomain ([94.250.174.60]) by smtp.gmail.com with ESMTPSA id k126-v6sm1591810wmd.45.2018.05.11.07.03.23 for (version=TLS1_2 cipher=ECDHE-RSA-AES128-GCM-SHA256 bits=128/128); Fri, 11 May 2018 07:03:24 -0700 (PDT) From: Paul B Mahol To: ffmpeg-devel@ffmpeg.org Date: Fri, 11 May 2018 16:03:07 +0200 Message-Id: <20180511140307.21095-1-onemda@gmail.com> X-Mailer: git-send-email 2.11.0 Subject: [FFmpeg-devel] [PATCH] avfilter: add bm3d filter X-BeenThere: ffmpeg-devel@ffmpeg.org X-Mailman-Version: 2.1.20 Precedence: list List-Id: FFmpeg development discussions and patches List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , Reply-To: FFmpeg development discussions and patches MIME-Version: 1.0 Errors-To: ffmpeg-devel-bounces@ffmpeg.org Sender: "ffmpeg-devel" Signed-off-by: Paul B Mahol --- libavfilter/Makefile | 1 + libavfilter/allfilters.c | 1 + libavfilter/vf_bm3d.c | 1002 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 1004 insertions(+) create mode 100644 libavfilter/vf_bm3d.c diff --git a/libavfilter/Makefile b/libavfilter/Makefile index 92f526a275..b2d6756e79 100644 --- a/libavfilter/Makefile +++ b/libavfilter/Makefile @@ -150,6 +150,7 @@ OBJS-$(CONFIG_BITPLANENOISE_FILTER) += vf_bitplanenoise.o OBJS-$(CONFIG_BLACKDETECT_FILTER) += vf_blackdetect.o OBJS-$(CONFIG_BLACKFRAME_FILTER) += vf_blackframe.o OBJS-$(CONFIG_BLEND_FILTER) += vf_blend.o framesync.o +OBJS-$(CONFIG_BM3D_FILTER) += vf_bm3d.o OBJS-$(CONFIG_BOXBLUR_FILTER) += vf_boxblur.o OBJS-$(CONFIG_BWDIF_FILTER) += vf_bwdif.o OBJS-$(CONFIG_CHROMAKEY_FILTER) += vf_chromakey.o diff --git a/libavfilter/allfilters.c b/libavfilter/allfilters.c index f3a76c4652..f28f6e47ee 100644 --- a/libavfilter/allfilters.c +++ b/libavfilter/allfilters.c @@ -143,6 +143,7 @@ extern AVFilter ff_vf_bitplanenoise; extern AVFilter ff_vf_blackdetect; extern AVFilter ff_vf_blackframe; extern AVFilter ff_vf_blend; +extern AVFilter ff_vf_bm3d; extern AVFilter ff_vf_boxblur; extern AVFilter ff_vf_bwdif; extern AVFilter ff_vf_chromakey; diff --git a/libavfilter/vf_bm3d.c b/libavfilter/vf_bm3d.c new file mode 100644 index 0000000000..61f9db17e3 --- /dev/null +++ b/libavfilter/vf_bm3d.c @@ -0,0 +1,1002 @@ +/* + * Copyright (c) 2015-2016 mawen1250 + * Copyright (c) 2018 Paul B Mahol + * + * This file is part of FFmpeg. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +/** + * @todo + * - non-power of 2 DCT + * - slice threading + * - opponent color space + * - temporal support + * - documentation + */ + +#include + +#include "libavutil/avassert.h" +#include "libavutil/imgutils.h" +#include "libavutil/opt.h" +#include "libavutil/pixdesc.h" +#include "libavcodec/avfft.h" +#include "avfilter.h" +#include "filters.h" +#include "formats.h" +#include "framesync.h" +#include "internal.h" +#include "video.h" + +enum FilterModes { + BASIC, + FINAL, + NB_MODES, +}; + +typedef struct PosCode { + int x, y; +} PosCode; + +typedef struct PosPairCode { + double score; + int x, y; +} PosPairCode; + +typedef struct BM3DContext { + const AVClass *class; + + float sigma; + int block_size; + int block_step; + int group_size; + int bm_range; + int bm_step; + float th_mse; + float hard_threshold; + int mode; + int ref; + int planes; + + int depth; + int max; + int nb_planes; + int planewidth[4]; + int planeheight[4]; + int group_bits; + int pgroup_size; + + DCTContext *gdctf, *gdcti; + DCTContext *dctf, *dcti; + FFTSample *bufferh; + FFTSample *bufferv; + FFTSample *bufferz; + FFTSample *buffer; + FFTSample *rbufferh; + FFTSample *rbufferv; + FFTSample *rbufferz; + FFTSample *rbuffer; + float *num, *den; + PosPairCode match_codes[256]; + int nb_match_codes; + + FFFrameSync fs; + + void (*get_block_row)(const uint8_t *srcp, int src_linesize, + int y, int x, int block_size, float *dst); + double (*do_block_ssd)(struct BM3DContext *s, PosCode *pos, + const uint8_t *src, int src_stride, + int r_y, int r_x); + void (*do_output)(struct BM3DContext *s, uint8_t *dst, int dst_linesize, int plane); + void (*block_filtering)(struct BM3DContext *s, const uint8_t *src, int src_linesize, + const uint8_t *ref, int ref_linesize, + int y, int x, int plane); +} BM3DContext; + +#define OFFSET(x) offsetof(BM3DContext, x) +#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM +static const AVOption bm3d_options[] = { + { "sigma", "set denoising strength", + OFFSET(sigma), AV_OPT_TYPE_FLOAT, {.dbl=1}, 0, 999.9, FLAGS }, + { "block", "set log2(size) of local patch", + OFFSET(block_size), AV_OPT_TYPE_INT, {.i64=4}, 4, 6, FLAGS }, + { "bstep", "set sliding step for processing blocks", + OFFSET(block_step), AV_OPT_TYPE_INT, {.i64=8}, 1, 64, FLAGS }, + { "group", "set maximal number of similar blocks", + OFFSET(group_size), AV_OPT_TYPE_INT, {.i64=1}, 1, 256, FLAGS }, + { "range", "set block matching range", + OFFSET(bm_range), AV_OPT_TYPE_INT, {.i64=9}, 1, INT32_MAX, FLAGS }, + { "mstep", "set step for block matching", + OFFSET(bm_step), AV_OPT_TYPE_INT, {.i64=1}, 1, 64, FLAGS }, + { "thmse", "set threshold of mean square error for block matching", + OFFSET(th_mse), AV_OPT_TYPE_FLOAT, {.dbl=0}, 0, INT32_MAX, FLAGS }, + { "hdthr", "set hard threshold for 3D transfer domain", + OFFSET(hard_threshold), AV_OPT_TYPE_FLOAT, {.dbl=2.7}, 0, INT32_MAX, FLAGS }, + { "estim", "set filtering estimation mode", + OFFSET(mode), AV_OPT_TYPE_INT, {.i64=BASIC}, 0, NB_MODES-1, FLAGS, "mode" }, + { "basic", "basic estimate", + 0, AV_OPT_TYPE_CONST, {.i64=BASIC}, 0, 0, FLAGS, "mode" }, + { "final", "final estimate", + 0, AV_OPT_TYPE_CONST, {.i64=FINAL}, 0, 0, FLAGS, "mode" }, + { "ref", "have reference stream", + OFFSET(ref), AV_OPT_TYPE_INT, {.i64=0}, 0, 1, FLAGS }, + { "planes", "set planes to filter", + OFFSET(planes), AV_OPT_TYPE_INT, {.i64=7}, 0, 15, FLAGS }, + { NULL } +}; + +AVFILTER_DEFINE_CLASS(bm3d); + +static int query_formats(AVFilterContext *ctx) +{ + static const enum AVPixelFormat pix_fmts[] = { + AV_PIX_FMT_GRAY8, + AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, + AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY16, + AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, + AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, + AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV444P, + AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ422P, + AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ444P, + AV_PIX_FMT_YUVJ411P, + 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_YUV440P10, + AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV420P12, + AV_PIX_FMT_YUV440P12, + AV_PIX_FMT_YUV444P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV420P14, + AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16, + AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10, + AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16, + AV_PIX_FMT_NONE + }; + + AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts); + if (!fmts_list) + return AVERROR(ENOMEM); + return ff_set_common_formats(ctx, fmts_list); +} + +static int do_search_boundary(int pos, int plane_boundary, int search_range, int search_step) +{ + int search_boundary; + + search_range = search_range / search_step * search_step; + + if (pos == plane_boundary) { + search_boundary = plane_boundary; + } else if (pos > plane_boundary) { + search_boundary = pos - search_range; + + while (search_boundary < plane_boundary) { + search_boundary += search_step; + } + } else { + search_boundary = pos + search_range; + + while (search_boundary > plane_boundary) { + search_boundary -= search_step; + } + } + + return search_boundary; +} + +static int search_boundary(int plane_boundary, int search_range, int search_step, int vertical, int y, int x) +{ + return do_search_boundary(vertical ? y : x, plane_boundary, search_range, search_step); +} + +static int cmp_scores(const void *a, const void *b) +{ + const struct PosPairCode *pair1 = a; + const struct PosPairCode *pair2 = b; + return FFDIFFSIGN(pair1->score, pair2->score); +} + +static double do_block_ssd(BM3DContext *s, PosCode *pos, const uint8_t *src, int src_stride, int r_y, int r_x) +{ + const uint8_t *srcp = src + pos->y * src_stride + pos->x; + const uint8_t *refp = src + r_y * src_stride + r_x; + const int block_size = s->block_size; + double dist = 0.; + int x, y; + + for (y = 0; y < block_size; y++) { + for (x = 0; x < block_size; x++) { + double temp = refp[x] - srcp[x]; + dist += temp * temp; + } + + srcp += src_stride; + refp += src_stride; + } + + return dist; +} + +static double do_block_ssd16(BM3DContext *s, PosCode *pos, const uint8_t *src, int src_stride, int r_y, int r_x) +{ + const uint16_t *srcp = (uint16_t *)src + pos->y * src_stride / 2 + pos->x; + const uint16_t *refp = (uint16_t *)src + r_y * src_stride / 2 + r_x; + const int block_size = s->block_size; + double dist = 0.; + int x, y; + + for (y = 0; y < block_size; y++) { + for (x = 0; x < block_size; x++) { + double temp = refp[x] - srcp[x]; + dist += temp * temp; + } + + srcp += src_stride / 2; + refp += src_stride / 2; + } + + return dist; +} + +static void do_block_matching_multi(BM3DContext *s, const uint8_t *src, int src_stride, int src_range, + const PosCode *search_pos, int search_size, float thMSE, + int r_y, int r_x, int plane) +{ + double MSE2SSE = s->group_size * s->block_size * s->block_size * src_range * src_range / (s->max * s->max); + double distMul = 1. / MSE2SSE; + double thSSE = thMSE * MSE2SSE; + int i, index = s->nb_match_codes; + + for (i = 0; i < search_size; i++) { + PosCode pos = search_pos[i]; + double dist; + + dist = s->do_block_ssd(s, &pos, src, src_stride, r_y, r_x); + + // Only match similar blocks but not identical blocks + if (dist <= thSSE && dist != 0) { + const double score = dist * distMul; + + if (index >= s->group_size && score >= s->match_codes[index - 1].score) { + continue; + } + + if (index >= s->group_size) + index = s->group_size - 1; + + s->match_codes[index].score = score; + s->match_codes[index].y = pos.y; + s->match_codes[index].x = pos.x; + index++; + qsort(s->match_codes, index, sizeof(PosPairCode), cmp_scores); + } + } + + s->nb_match_codes = index; +} + +static void block_matching_multi(BM3DContext *s, const uint8_t *ref, int ref_linesize, int y, int x, + int exclude_cur_pos, int plane) +{ + const int width = s->planewidth[plane]; + const int height = s->planeheight[plane]; + const int block_size = s->block_size; + const int step = s->bm_step; + const int range = s->bm_range / step * step; + int l = search_boundary(0, range, step, 0, y, x); + int r = search_boundary(width - block_size, range, step, 0, y, x); + int t = search_boundary(0, range, step, 1, y, x); + int b = search_boundary(height - block_size, range, step, 1, y, x); + int j, i, index = 0; + + PosCode search_pos[((r - l) / step + 1) * ((b - t) / step + 1)]; + + for (j = t; j <= b; j += step) { + for (i = l; i <= r; i += step) { + PosCode pos; + + if (exclude_cur_pos > 0 && j == y && i == x) { + continue; + } + + pos.y = j; + pos.x = i; + search_pos[index++] = pos; + } + } + + if (exclude_cur_pos == 1) { + s->match_codes[0].score = 0; + s->match_codes[0].y = y; + s->match_codes[0].x = x; + s->nb_match_codes = 1; + } + + do_block_matching_multi(s, ref, ref_linesize, s->bm_range, + search_pos, index, s->th_mse, y, x, plane); +} + +static void block_matching(BM3DContext *s, const uint8_t *ref, int ref_linesize, + int j, int i, int plane) +{ + if (s->group_size == 1 || s->th_mse <= 0) { + s->match_codes[0].score = 1; + s->match_codes[0].x = i; + s->match_codes[0].y = j; + s->nb_match_codes = 1; + return; + } + + s->nb_match_codes = 0; + block_matching_multi(s, ref, ref_linesize, j, i, 1, plane); +} + +static void get_block_row(const uint8_t *srcp, int src_linesize, + int y, int x, int block_size, float *dst) +{ + const uint8_t *src = srcp + y * src_linesize + x; + int j; + + for (j = 0; j < block_size; j++) { + dst[j] = src[j]; + } +} + +static void get_block_row16(const uint8_t *srcp, int src_linesize, + int y, int x, int block_size, float *dst) +{ + const uint16_t *src = (uint16_t *)srcp + y * src_linesize / 2 + x; + int j; + + for (j = 0; j < block_size; j++) { + dst[j] = src[j]; + } +} + +static void basic_block_filtering(BM3DContext *s, const uint8_t *src, int src_linesize, + const uint8_t *ref, int ref_linesize, + int y, int x, int plane) +{ + const int buffer_linesize = s->block_size * s->block_size; + const int nb_match_codes = s->nb_match_codes; + const int block_size = s->block_size; + const int width = s->planewidth[plane]; + const int pgroup_size = s->pgroup_size; + const int group_size = s->group_size; + float *buffer = s->buffer; + float *bufferh = s->bufferh; + float *bufferv = s->bufferv; + float *bufferz = s->bufferz; + float threshold[4]; + float den_weight, num_weight; + int retained = 0; + int i, j, k; + + for (k = 0; k < nb_match_codes; k++) { + const int y = s->match_codes[k].y; + const int x = s->match_codes[k].x; + + for (i = 0; i < block_size; i++) { + s->get_block_row(src, src_linesize, y + i, x, block_size, bufferh + block_size * i); + av_dct_calc(s->dctf, bufferh + block_size * i); + } + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + bufferv[i * block_size + j] = bufferh[j * block_size + i]; + } + av_dct_calc(s->dctf, bufferv + i * block_size); + } + + for (i = 0; i < block_size; i++) { + memcpy(buffer + k * buffer_linesize + i * block_size, + bufferv + i * block_size, block_size * 4); + } + } + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + for (k = 0; k < nb_match_codes; k++) + bufferz[k] = buffer[buffer_linesize * k + i * block_size + j]; + if (group_size > 1) + av_dct_calc(s->gdctf, bufferz); + bufferz += pgroup_size; + } + } + + threshold[0] = s->hard_threshold * s->sigma * 2.f; + threshold[1] = threshold[0] * sqrtf(2.f); + threshold[2] = threshold[0] * 2.f; + threshold[3] = threshold[0] * sqrtf(8.f); + bufferz = s->bufferz; + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + for (k = 0; k < nb_match_codes; k++) { + const float thresh = threshold[(j == 0) + (i == 0) + (k == 0)]; + + if (bufferz[k] > thresh || bufferz[k] < -thresh) { + retained++; + } else { + bufferz[k] = 0; + } + } + bufferz += pgroup_size; + } + } + + bufferz = s->bufferz; + buffer = s->buffer; + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + if (group_size > 1) + av_dct_calc(s->gdcti, bufferz); + for (k = 0; k < nb_match_codes; k++) { + buffer[buffer_linesize * k + i * block_size + j] = bufferz[k]; + } + bufferz += pgroup_size; + } + } + + den_weight = retained < 1 ? 1.f : 1.f / retained; + num_weight = den_weight; + + buffer = s->buffer; + for (k = 0; k < nb_match_codes; k++) { + float *num = s->num + y * width + x; + float *den = s->den + y * width + x; + + for (i = 0; i < block_size; i++) { + memcpy(bufferv + i * block_size, + buffer + k * buffer_linesize + i * block_size, + block_size * 4); + } + + for (i = 0; i < block_size; i++) { + av_dct_calc(s->dcti, bufferv + block_size * i); + for (j = 0; j < block_size; j++) { + bufferh[j * block_size + i] = bufferv[i * block_size + j]; + } + } + + for (i = 0; i < block_size; i++) { + av_dct_calc(s->dcti, bufferh + block_size * i); + for (j = 0; j < block_size; j++) { + num[j] += bufferh[i * block_size + j] * num_weight; + den[j] += den_weight; + } + num += width; + den += width; + } + } +} + +static void final_block_filtering(BM3DContext *s, const uint8_t *src, int src_linesize, + const uint8_t *ref, int ref_linesize, + int y, int x, int plane) +{ + const int buffer_linesize = s->block_size * s->block_size; + const int nb_match_codes = s->nb_match_codes; + const int block_size = s->block_size; + const int width = s->planewidth[plane]; + const int pgroup_size = s->pgroup_size; + const int group_size = s->group_size; + const float sigma_sqr = s->sigma * s->sigma; + float *buffer = s->buffer; + float *bufferh = s->bufferh; + float *bufferv = s->bufferv; + float *bufferz = s->bufferz; + float *rbuffer = s->rbuffer; + float *rbufferh = s->rbufferh; + float *rbufferv = s->rbufferv; + float *rbufferz = s->rbufferz; + float den_weight, num_weight; + float l2_wiener = 0; + int i, j, k; + + for (k = 0; k < nb_match_codes; k++) { + const int y = s->match_codes[k].y; + const int x = s->match_codes[k].x; + + for (i = 0; i < block_size; i++) { + s->get_block_row(src, src_linesize, y + i, x, block_size, bufferh + block_size * i); + s->get_block_row(ref, ref_linesize, y + i, x, block_size, rbufferh + block_size * i); + av_dct_calc(s->dctf, bufferh + block_size * i); + av_dct_calc(s->dctf, rbufferh + block_size * i); + } + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + bufferv[i * block_size + j] = bufferh[j * block_size + i]; + rbufferv[i * block_size + j] = rbufferh[j * block_size + i]; + } + av_dct_calc(s->dctf, bufferv + i * block_size); + av_dct_calc(s->dctf, rbufferv + i * block_size); + } + + for (i = 0; i < block_size; i++) { + memcpy(buffer + k * buffer_linesize + i * block_size, + bufferv + i * block_size, block_size * 4); + memcpy(rbuffer + k * buffer_linesize + i * block_size, + rbufferv + i * block_size, block_size * 4); + } + } + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + for (k = 0; k < nb_match_codes; k++) { + bufferz[k] = buffer[buffer_linesize * k + i * block_size + j]; + rbufferz[k] = rbuffer[buffer_linesize * k + i * block_size + j]; + } + if (group_size > 1) { + av_dct_calc(s->gdctf, bufferz); + av_dct_calc(s->gdctf, rbufferz); + } + bufferz += pgroup_size; + rbufferz += pgroup_size; + } + } + + bufferz = s->bufferz; + rbufferz = s->rbufferz; + + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + for (k = 0; k < nb_match_codes; k++) { + const float ref_sqr = rbufferz[k] * rbufferz[k]; + float wiener_coef = ref_sqr / (ref_sqr + sigma_sqr); + + if (isnan(wiener_coef)) + wiener_coef = 1; + bufferz[k] *= wiener_coef; + l2_wiener += wiener_coef * wiener_coef; + } + bufferz += pgroup_size; + rbufferz += pgroup_size; + } + } + + bufferz = s->bufferz; + buffer = s->buffer; + for (i = 0; i < block_size; i++) { + for (j = 0; j < block_size; j++) { + if (group_size > 1) + av_dct_calc(s->gdcti, bufferz); + for (k = 0; k < nb_match_codes; k++) { + buffer[buffer_linesize * k + i * block_size + j] = bufferz[k]; + } + bufferz += pgroup_size; + } + } + + l2_wiener = FFMAX(l2_wiener, 1e-15f); + den_weight = 1.f / l2_wiener; + num_weight = den_weight; + + for (k = 0; k < nb_match_codes; k++) { + float *num = s->num + y * width + x; + float *den = s->den + y * width + x; + + for (i = 0; i < block_size; i++) { + memcpy(bufferv + i * block_size, + buffer + k * buffer_linesize + i * block_size, + block_size * 4); + } + + for (i = 0; i < block_size; i++) { + av_dct_calc(s->dcti, bufferv + block_size * i); + for (j = 0; j < block_size; j++) { + bufferh[j * block_size + i] = bufferv[i * block_size + j]; + } + } + + for (i = 0; i < block_size; i++) { + av_dct_calc(s->dcti, bufferh + block_size * i); + for (j = 0; j < block_size; j++) { + num[j] += bufferh[i * block_size + j] * num_weight; + den[j] += den_weight; + } + num += width; + den += width; + } + } +} + +static void do_output(BM3DContext *s, uint8_t *dst, int dst_linesize, int plane) +{ + const int height = s->planeheight[plane]; + const int width = s->planewidth[plane]; + int i, j; + + for (i = 0; i < height; i++) { + uint8_t *dstp = dst + i * dst_linesize; + float *num = s->num + i * width; + float *den = s->den + i * width; + + for (j = 0; j < width; j++) { + dstp[j] = av_clip_uint8(num[j] / den[j]); + } + } +} + +static void do_output16(BM3DContext *s, uint8_t *dst, int dst_linesize, int plane) +{ + const int height = s->planeheight[plane]; + const int width = s->planewidth[plane]; + const int depth = s->depth; + int i, j; + + for (i = 0; i < height; i++) { + uint16_t *dstp = (uint16_t *)dst + i * dst_linesize / 2; + float *num = s->num + i * width; + float *den = s->den + i * width; + + for (j = 0; j < width; j++) { + dstp[j] = av_clip_uintp2(num[j] / den[j], depth); + } + } +} + +static void kernel(BM3DContext *s, const uint8_t *src, int src_linesize, + const uint8_t *ref, int ref_linesize, + uint8_t *dst, int dst_linesize, int plane) +{ + const int width = s->planewidth[plane]; + const int height = s->planeheight[plane]; + const int block_pos_bottom = height - s->block_size; + const int block_pos_right = width - s->block_size; + int i, j; + + memset(s->num, 0, width * height * sizeof(*s->num)); + memset(s->den, 0, width * height * sizeof(*s->den)); + + for (j = 0; ; j += s->block_step) { + if (j >= block_pos_bottom + s->block_step) { + break; + } else if (j > block_pos_bottom) { + j = block_pos_bottom; + } + for (i = 0; ; i += s->block_step) { + if (i >= block_pos_right + s->block_step) { + break; + } else if (i > block_pos_right) { + i = block_pos_right; + } + + block_matching(s, ref, ref_linesize, j, i, plane); + + s->block_filtering(s, src, src_linesize, + ref, ref_linesize, j, i, plane); + } + } + + s->do_output(s, dst, dst_linesize, plane); +} + +static int filter_frame(AVFilterContext *ctx, AVFrame **out, AVFrame *in, AVFrame *ref) +{ + BM3DContext *s = ctx->priv; + AVFilterLink *outlink = ctx->outputs[0]; + int p; + + *out = ff_get_video_buffer(outlink, outlink->w, outlink->h); + if (!*out) + return AVERROR(ENOMEM); + av_frame_copy_props(*out, in); + + for (p = 0; p < s->nb_planes; p++) { + if (!((1 << p) & s->planes)) { + av_image_copy_plane((*out)->data[p], (*out)->linesize[p], + in->data[p], in->linesize[p], + s->planewidth[p], s->planeheight[p]); + continue; + } + kernel(s, in->data[p], in->linesize[p], + ref->data[p], ref->linesize[p], + (*out)->data[p], (*out)->linesize[p], p); + } + + return 0; +} + +static int config_input(AVFilterLink *inlink) +{ + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format); + AVFilterContext *ctx = inlink->dst; + BM3DContext *s = ctx->priv; + int group_bits; + + s->nb_planes = av_pix_fmt_count_planes(inlink->format); + s->depth = desc->comp[0].depth; + s->max = (1 << s->depth) - 1; + s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h); + s->planeheight[0] = s->planeheight[3] = inlink->h; + s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w); + s->planewidth[0] = s->planewidth[3] = inlink->w; + + s->num = av_calloc(s->planewidth[0] * s->planeheight[0], sizeof(*s->num)); + s->den = av_calloc(s->planewidth[0] * s->planeheight[0], sizeof(*s->den)); + if (!s->num || !s->den) + return AVERROR(ENOMEM); + + for (group_bits = 4; 1 << group_bits < s->group_size; group_bits++); + s->group_bits = group_bits; + s->pgroup_size = 1 << group_bits; + + if (s->group_bits > 1) { + s->gdctf = av_dct_init(s->group_bits, DCT_II); + s->gdcti = av_dct_init(s->group_bits, DCT_III); + if (!s->gdctf || !s->gdcti) + return AVERROR(ENOMEM); + } + + s->buffer = av_calloc(s->block_size * s->block_size * s->pgroup_size, sizeof(*s->buffer)); + s->bufferz = av_calloc(s->block_size * s->block_size * s->pgroup_size, sizeof(*s->bufferz)); + s->bufferh = av_calloc(s->block_size * s->block_size, sizeof(*s->bufferh)); + s->bufferv = av_calloc(s->block_size * s->block_size, sizeof(*s->bufferv)); + if (!s->bufferh || !s->bufferv || !s->buffer || !s->bufferz) + return AVERROR(ENOMEM); + + if (s->mode == FINAL) { + s->rbuffer = av_calloc(s->block_size * s->block_size * s->pgroup_size, sizeof(*s->rbuffer)); + s->rbufferz = av_calloc(s->block_size * s->block_size * s->pgroup_size, sizeof(*s->rbufferz)); + s->rbufferh = av_calloc(s->block_size * s->block_size, sizeof(*s->rbufferh)); + s->rbufferv = av_calloc(s->block_size * s->block_size, sizeof(*s->rbufferv)); + if (!s->rbufferh || !s->rbufferv || !s->rbuffer || !s->rbufferz) + return AVERROR(ENOMEM); + } + + s->dctf = av_dct_init(av_log2(s->block_size), DCT_II); + s->dcti = av_dct_init(av_log2(s->block_size), DCT_III); + if (!s->dctf || !s->dcti) + return AVERROR(ENOMEM); + + s->do_output = do_output; + s->do_block_ssd = do_block_ssd; + s->get_block_row = get_block_row; + + if (s->depth > 8) { + s->do_output = do_output16; + s->do_block_ssd = do_block_ssd16; + s->get_block_row = get_block_row16; + } + + return 0; +} + +static int activate(AVFilterContext *ctx) +{ + BM3DContext *s = ctx->priv; + + if (!s->ref) { + AVFrame *frame = NULL; + AVFrame *out = NULL; + int ret, status; + int64_t pts; + + if ((ret = ff_inlink_consume_frame(ctx->inputs[0], &frame)) > 0) { + ret = filter_frame(ctx, &out, frame, frame); + av_frame_free(&frame); + if (ret < 0) + return ret; + ret = ff_filter_frame(ctx->outputs[0], out); + } + if (ret < 0) { + return ret; + } else if (ff_inlink_acknowledge_status(ctx->inputs[0], &status, &pts)) { + ff_outlink_set_status(ctx->outputs[0], status, pts); + return 0; + } else { + if (ff_outlink_frame_wanted(ctx->outputs[0])) + ff_inlink_request_frame(ctx->inputs[0]); + return 0; + } + } else { + return ff_framesync_activate(&s->fs); + } +} + +static int process_frame(FFFrameSync *fs) +{ + AVFilterContext *ctx = fs->parent; + BM3DContext *s = fs->opaque; + AVFilterLink *outlink = ctx->outputs[0]; + AVFrame *out = NULL, *src, *ref; + int ret; + + if ((ret = ff_framesync_get_frame(&s->fs, 0, &src, 0)) < 0 || + (ret = ff_framesync_get_frame(&s->fs, 1, &ref, 0)) < 0) + return ret; + + if ((ret = filter_frame(ctx, &out, src, ref)) < 0) + return ret; + + out->pts = av_rescale_q(src->pts, s->fs.time_base, outlink->time_base); + + return ff_filter_frame(outlink, out); +} + +static av_cold int init(AVFilterContext *ctx) +{ + BM3DContext *s = ctx->priv; + AVFilterPad pad = { 0 }; + int ret; + + if (s->mode == BASIC) { + if (s->th_mse == 0) + s->th_mse = 400.f + s->sigma * 80.f; + s->block_filtering = basic_block_filtering; + } else if (s->mode == FINAL) { + if (!s->ref) { + av_log(ctx, AV_LOG_WARNING, "Reference stream is mandatory in final estimation mode.\n"); + s->ref = 1; + } + if (s->th_mse == 0) + s->th_mse = 200.f + s->sigma * 10.f; + + s->block_filtering = final_block_filtering; + } else { + return AVERROR_BUG; + } + + s->block_size = 1 << s->block_size; + + if (s->block_step > s->block_size) { + av_log(ctx, AV_LOG_WARNING, "bstep: %d can't be bigger than block size. Changing to %d.\n", + s->block_step, s->block_size); + s->block_step = s->block_size; + } + if (s->bm_step > s->bm_range) { + av_log(ctx, AV_LOG_WARNING, "mstep: %d can't be bigger than block matching range. Changing to %d.\n", + s->bm_step, s->bm_range); + s->bm_step = s->bm_range; + } + + pad.type = AVMEDIA_TYPE_VIDEO; + pad.name = av_strdup("source"); + pad.config_props = config_input; + if (!pad.name) + return AVERROR(ENOMEM); + + if ((ret = ff_insert_inpad(ctx, 0, &pad)) < 0) { + av_freep(&pad.name); + return ret; + } + + if (s->ref) { + pad.type = AVMEDIA_TYPE_VIDEO; + pad.name = av_strdup("reference"); + pad.config_props = NULL; + if (!pad.name) + return AVERROR(ENOMEM); + + if ((ret = ff_insert_inpad(ctx, 1, &pad)) < 0) { + av_freep(&pad.name); + return ret; + } + } + + return 0; +} + +static int config_output(AVFilterLink *outlink) +{ + AVFilterContext *ctx = outlink->src; + BM3DContext *s = ctx->priv; + AVFilterLink *src = ctx->inputs[0]; + AVFilterLink *ref; + FFFrameSyncIn *in; + int ret; + + if (s->ref) { + ref = ctx->inputs[1]; + + if (src->format != ref->format) { + av_log(ctx, AV_LOG_ERROR, "inputs must be of same pixel format\n"); + return AVERROR(EINVAL); + } + if (src->w != ref->w || + src->h != ref->h) { + av_log(ctx, AV_LOG_ERROR, "First input link %s parameters " + "(size %dx%d) do not match the corresponding " + "second input link %s parameters (%dx%d) ", + ctx->input_pads[0].name, src->w, src->h, + ctx->input_pads[1].name, ref->w, ref->h); + return AVERROR(EINVAL); + } + } + + outlink->w = src->w; + outlink->h = src->h; + outlink->time_base = src->time_base; + outlink->sample_aspect_ratio = src->sample_aspect_ratio; + outlink->frame_rate = src->frame_rate; + + if (!s->ref) + return 0; + + if ((ret = ff_framesync_init(&s->fs, ctx, 2)) < 0) + return ret; + + in = s->fs.in; + in[0].time_base = src->time_base; + in[1].time_base = ref->time_base; + in[0].sync = 1; + in[0].before = EXT_STOP; + in[0].after = EXT_STOP; + in[1].sync = 1; + in[1].before = EXT_STOP; + in[1].after = EXT_STOP; + s->fs.opaque = s; + s->fs.on_event = process_frame; + + return ff_framesync_configure(&s->fs); +} + +static av_cold void uninit(AVFilterContext *ctx) +{ + BM3DContext *s = ctx->priv; + int i; + + for (i = 0; i < ctx->nb_inputs; i++) + av_freep(&ctx->input_pads[i].name); + + if (s->ref) + ff_framesync_uninit(&s->fs); + + av_freep(&s->num); + av_freep(&s->den); + av_freep(&s->buffer); + av_freep(&s->bufferh); + av_freep(&s->bufferv); + av_freep(&s->bufferz); + av_freep(&s->rbuffer); + av_freep(&s->rbufferh); + av_freep(&s->rbufferv); + av_freep(&s->rbufferz); + + av_dct_end(s->gdctf); + av_dct_end(s->gdcti); + av_dct_end(s->dctf); + av_dct_end(s->dcti); +} + +static const AVFilterPad bm3d_outputs[] = { + { + .name = "default", + .type = AVMEDIA_TYPE_VIDEO, + .config_props = config_output, + }, + { NULL } +}; + +AVFilter ff_vf_bm3d = { + .name = "bm3d", + .description = NULL_IF_CONFIG_SMALL("BM3D denoiser."), + .priv_size = sizeof(BM3DContext), + .init = init, + .uninit = uninit, + .activate = activate, + .query_formats = query_formats, + .inputs = NULL, + .outputs = bm3d_outputs, + .priv_class = &bm3d_class, + .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL | + AVFILTER_FLAG_DYNAMIC_INPUTS, +};