503 lines
21 KiB
C
503 lines
21 KiB
C
/*
|
|
* Copyright (c) 2021 Xuewei Meng
|
|
*
|
|
* This file is part of FFmpeg.
|
|
*
|
|
* FFmpeg is free software; you can redistribute it and/or
|
|
* modify it under the terms of the GNU Lesser General Public
|
|
* License as published by the Free Software Foundation; either
|
|
* version 2.1 of the License, or (at your option) any later version.
|
|
*
|
|
* FFmpeg is distributed in the hope that it will be useful,
|
|
* but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
|
|
* Lesser General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Lesser General Public
|
|
* License along with FFmpeg; if not, write to the Free Software
|
|
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
|
|
*/
|
|
|
|
#include "libavutil/imgutils.h"
|
|
#include "libavutil/mem.h"
|
|
#include "libavutil/opt.h"
|
|
#include "libavutil/pixdesc.h"
|
|
#include "avfilter.h"
|
|
#include "filters.h"
|
|
#include "framesync.h"
|
|
#include "internal.h"
|
|
#include "video.h"
|
|
|
|
enum FilterModes {
|
|
BASIC,
|
|
FAST,
|
|
NB_MODES,
|
|
};
|
|
|
|
enum GuidanceModes {
|
|
OFF,
|
|
ON,
|
|
NB_GUIDANCE_MODES,
|
|
};
|
|
|
|
typedef struct GuidedContext {
|
|
const AVClass *class;
|
|
FFFrameSync fs;
|
|
|
|
int radius;
|
|
float eps;
|
|
int mode;
|
|
int sub;
|
|
int guidance;
|
|
int planes;
|
|
|
|
int width;
|
|
int height;
|
|
|
|
int nb_planes;
|
|
int depth;
|
|
int planewidth[4];
|
|
int planeheight[4];
|
|
|
|
float *I;
|
|
float *II;
|
|
float *P;
|
|
float *IP;
|
|
float *meanI;
|
|
float *meanII;
|
|
float *meanP;
|
|
float *meanIP;
|
|
|
|
float *A;
|
|
float *B;
|
|
float *meanA;
|
|
float *meanB;
|
|
|
|
int (*box_slice)(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
|
|
} GuidedContext;
|
|
|
|
#define OFFSET(x) offsetof(GuidedContext, x)
|
|
#define TFLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
|
|
#define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
|
|
|
|
static const AVOption guided_options[] = {
|
|
{ "radius", "set the box radius", OFFSET(radius), AV_OPT_TYPE_INT, {.i64 = 3 }, 1, 20, TFLAGS },
|
|
{ "eps", "set the regularization parameter (with square)", OFFSET(eps), AV_OPT_TYPE_FLOAT, {.dbl = 0.01 }, 0.0, 1, TFLAGS },
|
|
{ "mode", "set filtering mode (0: basic mode; 1: fast mode)", OFFSET(mode), AV_OPT_TYPE_INT, {.i64 = BASIC}, BASIC, NB_MODES - 1, TFLAGS, .unit = "mode" },
|
|
{ "basic", "basic guided filter", 0, AV_OPT_TYPE_CONST, {.i64 = BASIC}, 0, 0, TFLAGS, .unit = "mode" },
|
|
{ "fast", "fast guided filter", 0, AV_OPT_TYPE_CONST, {.i64 = FAST }, 0, 0, TFLAGS, .unit = "mode" },
|
|
{ "sub", "subsampling ratio for fast mode", OFFSET(sub), AV_OPT_TYPE_INT, {.i64 = 4 }, 2, 64, TFLAGS },
|
|
{ "guidance", "set guidance mode (0: off mode; 1: on mode)", OFFSET(guidance), AV_OPT_TYPE_INT, {.i64 = OFF }, OFF, NB_GUIDANCE_MODES - 1, FLAGS, .unit = "guidance" },
|
|
{ "off", "only one input is enabled", 0, AV_OPT_TYPE_CONST, {.i64 = OFF }, 0, 0, FLAGS, .unit = "guidance" },
|
|
{ "on", "two inputs are required", 0, AV_OPT_TYPE_CONST, {.i64 = ON }, 0, 0, FLAGS, .unit = "guidance" },
|
|
{ "planes", "set planes to filter", OFFSET(planes), AV_OPT_TYPE_INT, {.i64 = 1 }, 0, 0xF, TFLAGS },
|
|
{ NULL }
|
|
};
|
|
|
|
AVFILTER_DEFINE_CLASS(guided);
|
|
|
|
typedef struct ThreadData {
|
|
int width;
|
|
int height;
|
|
float *src;
|
|
float *dst;
|
|
int srcStride;
|
|
int dstStride;
|
|
} ThreadData;
|
|
|
|
static int box_slice(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
|
|
{
|
|
GuidedContext *s = ctx->priv;
|
|
ThreadData *t = arg;
|
|
|
|
const int width = t->width;
|
|
const int height = t->height;
|
|
const int src_stride = t->srcStride;
|
|
const int dst_stride = t->dstStride;
|
|
const int slice_start = (height * jobnr) / nb_jobs;
|
|
const int slice_end = (height * (jobnr + 1)) / nb_jobs;
|
|
const int radius = s->radius;
|
|
const float *src = t->src;
|
|
float *dst = t->dst;
|
|
|
|
int w;
|
|
int numPix;
|
|
w = (radius << 1) + 1;
|
|
numPix = w * w;
|
|
for (int i = slice_start;i < slice_end;i++) {
|
|
for (int j = 0;j < width;j++) {
|
|
float temp = 0.0;
|
|
for (int row = -radius;row <= radius;row++) {
|
|
for (int col = -radius;col <= radius;col++) {
|
|
int x = i + row;
|
|
int y = j + col;
|
|
x = (x < 0) ? 0 : (x >= height ? height - 1 : x);
|
|
y = (y < 0) ? 0 : (y >= width ? width - 1 : y);
|
|
temp += src[x * src_stride + y];
|
|
}
|
|
}
|
|
dst[i * dst_stride + j] = temp / numPix;
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
static const enum AVPixelFormat pix_fmts[] = {
|
|
AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV440P,
|
|
AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P,
|
|
AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P,
|
|
AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P,
|
|
AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P,
|
|
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_YUV420P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV440P12,
|
|
AV_PIX_FMT_YUV420P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV444P14,
|
|
AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16,
|
|
AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9,
|
|
AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10,
|
|
AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
|
|
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_GBRAP, AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
|
|
AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY14, AV_PIX_FMT_GRAY16,
|
|
AV_PIX_FMT_NONE
|
|
};
|
|
|
|
static int config_input(AVFilterLink *inlink)
|
|
{
|
|
AVFilterContext *ctx = inlink->dst;
|
|
GuidedContext *s = ctx->priv;
|
|
const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
|
|
|
|
if (s->mode == BASIC) {
|
|
s->sub = 1;
|
|
} else if (s->mode == FAST) {
|
|
if (s->radius >= s->sub)
|
|
s->radius = s->radius / s->sub;
|
|
else {
|
|
s->radius = 1;
|
|
}
|
|
}
|
|
|
|
s->depth = desc->comp[0].depth;
|
|
s->width = ctx->inputs[0]->w;
|
|
s->height = ctx->inputs[0]->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->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
|
|
s->planeheight[0] = s->planeheight[3] = inlink->h;
|
|
|
|
s->nb_planes = av_pix_fmt_count_planes(inlink->format);
|
|
s->box_slice = box_slice;
|
|
return 0;
|
|
}
|
|
|
|
#define GUIDED(type, name) \
|
|
static int guided_##name(AVFilterContext *ctx, GuidedContext *s, \
|
|
const uint8_t *ssrc, const uint8_t *ssrcRef, \
|
|
uint8_t *ddst, int radius, float eps, int width, int height, \
|
|
int src_stride, int src_ref_stride, int dst_stride, \
|
|
float maxval) \
|
|
{ \
|
|
int ret = 0; \
|
|
type *dst = (type *)ddst; \
|
|
const type *src = (const type *)ssrc; \
|
|
const type *srcRef = (const type *)ssrcRef; \
|
|
\
|
|
int sub = s->sub; \
|
|
int h = (height % sub) == 0 ? height / sub : height / sub + 1; \
|
|
int w = (width % sub) == 0 ? width / sub : width / sub + 1; \
|
|
\
|
|
ThreadData t; \
|
|
const int nb_threads = ff_filter_get_nb_threads(ctx); \
|
|
float *I = s->I; \
|
|
float *II = s->II; \
|
|
float *P = s->P; \
|
|
float *IP = s->IP; \
|
|
float *meanI = s->meanI; \
|
|
float *meanII = s->meanII; \
|
|
float *meanP = s->meanP; \
|
|
float *meanIP = s->meanIP; \
|
|
float *A = s->A; \
|
|
float *B = s->B; \
|
|
float *meanA = s->meanA; \
|
|
float *meanB = s->meanB; \
|
|
\
|
|
for (int i = 0;i < h;i++) { \
|
|
for (int j = 0;j < w;j++) { \
|
|
int x = i * w + j; \
|
|
I[x] = src[(i * src_stride + j) * sub] / maxval; \
|
|
II[x] = I[x] * I[x]; \
|
|
P[x] = srcRef[(i * src_ref_stride + j) * sub] / maxval; \
|
|
IP[x] = I[x] * P[x]; \
|
|
} \
|
|
} \
|
|
\
|
|
t.width = w; \
|
|
t.height = h; \
|
|
t.srcStride = w; \
|
|
t.dstStride = w; \
|
|
t.src = I; \
|
|
t.dst = meanI; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
t.src = II; \
|
|
t.dst = meanII; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
t.src = P; \
|
|
t.dst = meanP; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
t.src = IP; \
|
|
t.dst = meanIP; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
\
|
|
for (int i = 0;i < h;i++) { \
|
|
for (int j = 0;j < w;j++) { \
|
|
int x = i * w + j; \
|
|
float varI = meanII[x] - (meanI[x] * meanI[x]); \
|
|
float covIP = meanIP[x] - (meanI[x] * meanP[x]); \
|
|
A[x] = covIP / (varI + eps); \
|
|
B[x] = meanP[x] - A[x] * meanI[x]; \
|
|
} \
|
|
} \
|
|
\
|
|
t.src = A; \
|
|
t.dst = meanA; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
t.src = B; \
|
|
t.dst = meanB; \
|
|
ff_filter_execute(ctx, s->box_slice, &t, NULL, FFMIN(h, nb_threads)); \
|
|
\
|
|
for (int i = 0;i < height;i++) { \
|
|
for (int j = 0;j < width;j++) { \
|
|
int x = i / sub * w + j / sub; \
|
|
dst[i * dst_stride + j] = meanA[x] * src[i * src_stride + j] + \
|
|
meanB[x] * maxval; \
|
|
} \
|
|
} \
|
|
\
|
|
return ret; \
|
|
}
|
|
|
|
GUIDED(uint8_t, byte)
|
|
GUIDED(uint16_t, word)
|
|
|
|
static int filter_frame(AVFilterContext *ctx, AVFrame **out, AVFrame *in, AVFrame *ref)
|
|
{
|
|
GuidedContext *s = ctx->priv;
|
|
AVFilterLink *outlink = ctx->outputs[0];
|
|
*out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
|
|
if (!*out)
|
|
return AVERROR(ENOMEM);
|
|
av_frame_copy_props(*out, in);
|
|
|
|
for (int plane = 0; plane < s->nb_planes; plane++) {
|
|
if (!(s->planes & (1 << plane))) {
|
|
av_image_copy_plane((*out)->data[plane], (*out)->linesize[plane],
|
|
in->data[plane], in->linesize[plane],
|
|
s->planewidth[plane] * ((s->depth + 7) / 8), s->planeheight[plane]);
|
|
continue;
|
|
}
|
|
if (s->depth <= 8)
|
|
guided_byte(ctx, s, in->data[plane], ref->data[plane], (*out)->data[plane], s->radius, s->eps,
|
|
s->planewidth[plane], s->planeheight[plane],
|
|
in->linesize[plane], ref->linesize[plane], (*out)->linesize[plane], (1 << s->depth) - 1.f);
|
|
else
|
|
guided_word(ctx, s, in->data[plane], ref->data[plane], (*out)->data[plane], s->radius, s->eps,
|
|
s->planewidth[plane], s->planeheight[plane],
|
|
in->linesize[plane] / 2, ref->linesize[plane] / 2, (*out)->linesize[plane] / 2, (1 << s->depth) - 1.f);
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int process_frame(FFFrameSync *fs)
|
|
{
|
|
AVFilterContext *ctx = fs->parent;
|
|
AVFilterLink *outlink = ctx->outputs[0];
|
|
AVFrame *out_frame = NULL, *main_frame = NULL, *ref_frame = NULL;
|
|
int ret;
|
|
ret = ff_framesync_dualinput_get(fs, &main_frame, &ref_frame);
|
|
if (ret < 0)
|
|
return ret;
|
|
|
|
if (ctx->is_disabled)
|
|
return ff_filter_frame(outlink, main_frame);
|
|
|
|
ret = filter_frame(ctx, &out_frame, main_frame, ref_frame);
|
|
if (ret < 0)
|
|
return ret;
|
|
av_frame_free(&main_frame);
|
|
|
|
return ff_filter_frame(outlink, out_frame);
|
|
}
|
|
|
|
static int config_output(AVFilterLink *outlink)
|
|
{
|
|
AVFilterContext *ctx = outlink->src;
|
|
GuidedContext *s = ctx->priv;
|
|
AVFilterLink *mainlink = ctx->inputs[0];
|
|
FFFrameSyncIn *in;
|
|
int w, h, ret;
|
|
|
|
if (s->guidance == ON) {
|
|
if (ctx->inputs[0]->w != ctx->inputs[1]->w ||
|
|
ctx->inputs[0]->h != ctx->inputs[1]->h) {
|
|
av_log(ctx, AV_LOG_ERROR, "Width and height of input videos must be same.\n");
|
|
return AVERROR(EINVAL);
|
|
}
|
|
}
|
|
|
|
outlink->w = w = mainlink->w;
|
|
outlink->h = h = mainlink->h;
|
|
outlink->time_base = mainlink->time_base;
|
|
outlink->sample_aspect_ratio = mainlink->sample_aspect_ratio;
|
|
outlink->frame_rate = mainlink->frame_rate;
|
|
|
|
s->I = av_calloc(w * h, sizeof(*s->I));
|
|
s->II = av_calloc(w * h, sizeof(*s->II));
|
|
s->P = av_calloc(w * h, sizeof(*s->P));
|
|
s->IP = av_calloc(w * h, sizeof(*s->IP));
|
|
s->meanI = av_calloc(w * h, sizeof(*s->meanI));
|
|
s->meanII = av_calloc(w * h, sizeof(*s->meanII));
|
|
s->meanP = av_calloc(w * h, sizeof(*s->meanP));
|
|
s->meanIP = av_calloc(w * h, sizeof(*s->meanIP));
|
|
|
|
s->A = av_calloc(w * h, sizeof(*s->A));
|
|
s->B = av_calloc(w * h, sizeof(*s->B));
|
|
s->meanA = av_calloc(w * h, sizeof(*s->meanA));
|
|
s->meanB = av_calloc(w * h, sizeof(*s->meanA));
|
|
|
|
if (!s->I || !s->II || !s->P || !s->IP || !s->meanI || !s->meanII || !s->meanP ||
|
|
!s->meanIP || !s->A || !s->B || !s->meanA || !s->meanB)
|
|
return AVERROR(ENOMEM);
|
|
|
|
if (s->guidance == OFF)
|
|
return 0;
|
|
|
|
if ((ret = ff_framesync_init(&s->fs, ctx, 2)) < 0)
|
|
return ret;
|
|
|
|
outlink->time_base = s->fs.time_base;
|
|
|
|
in = s->fs.in;
|
|
in[0].time_base = mainlink->time_base;
|
|
in[1].time_base = ctx->inputs[1]->time_base;
|
|
in[0].sync = 2;
|
|
in[0].before = EXT_INFINITY;
|
|
in[0].after = EXT_INFINITY;
|
|
in[1].sync = 1;
|
|
in[1].before = EXT_INFINITY;
|
|
in[1].after = EXT_INFINITY;
|
|
s->fs.opaque = s;
|
|
s->fs.on_event = process_frame;
|
|
|
|
return ff_framesync_configure(&s->fs);
|
|
}
|
|
|
|
static int activate(AVFilterContext *ctx)
|
|
{
|
|
GuidedContext *s = ctx->priv;
|
|
AVFilterLink *outlink = ctx->outputs[0];
|
|
AVFilterLink *inlink = ctx->inputs[0];
|
|
AVFrame *frame = NULL;
|
|
AVFrame *out = NULL;
|
|
int ret, status;
|
|
int64_t pts;
|
|
if (s->guidance)
|
|
return ff_framesync_activate(&s->fs);
|
|
|
|
FF_FILTER_FORWARD_STATUS_BACK(outlink, inlink);
|
|
|
|
if ((ret = ff_inlink_consume_frame(inlink, &frame)) > 0) {
|
|
if (ctx->is_disabled)
|
|
return ff_filter_frame(outlink, frame);
|
|
|
|
ret = filter_frame(ctx, &out, frame, frame);
|
|
av_frame_free(&frame);
|
|
if (ret < 0)
|
|
return ret;
|
|
ret = ff_filter_frame(outlink, out);
|
|
}
|
|
if (ret < 0)
|
|
return ret;
|
|
if (ff_inlink_acknowledge_status(inlink, &status, &pts)) {
|
|
ff_outlink_set_status(outlink, status, pts);
|
|
return 0;
|
|
}
|
|
if (ff_outlink_frame_wanted(outlink))
|
|
ff_inlink_request_frame(inlink);
|
|
return 0;
|
|
}
|
|
|
|
static av_cold int init(AVFilterContext *ctx)
|
|
{
|
|
GuidedContext *s = ctx->priv;
|
|
AVFilterPad pad = { 0 };
|
|
int ret;
|
|
|
|
pad.type = AVMEDIA_TYPE_VIDEO;
|
|
pad.name = "source";
|
|
pad.config_props = config_input;
|
|
|
|
if ((ret = ff_append_inpad(ctx, &pad)) < 0)
|
|
return ret;
|
|
|
|
if (s->guidance == ON) {
|
|
pad.type = AVMEDIA_TYPE_VIDEO;
|
|
pad.name = "guidance";
|
|
pad.config_props = NULL;
|
|
|
|
if ((ret = ff_append_inpad(ctx, &pad)) < 0)
|
|
return ret;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static av_cold void uninit(AVFilterContext *ctx)
|
|
{
|
|
GuidedContext *s = ctx->priv;
|
|
if (s->guidance == ON)
|
|
ff_framesync_uninit(&s->fs);
|
|
|
|
av_freep(&s->I);
|
|
av_freep(&s->II);
|
|
av_freep(&s->P);
|
|
av_freep(&s->IP);
|
|
av_freep(&s->meanI);
|
|
av_freep(&s->meanII);
|
|
av_freep(&s->meanP);
|
|
av_freep(&s->meanIP);
|
|
av_freep(&s->A);
|
|
av_freep(&s->B);
|
|
av_freep(&s->meanA);
|
|
av_freep(&s->meanB);
|
|
|
|
return;
|
|
}
|
|
|
|
static const AVFilterPad guided_outputs[] = {
|
|
{
|
|
.name = "default",
|
|
.type = AVMEDIA_TYPE_VIDEO,
|
|
.config_props = config_output,
|
|
},
|
|
};
|
|
|
|
const AVFilter ff_vf_guided = {
|
|
.name = "guided",
|
|
.description = NULL_IF_CONFIG_SMALL("Apply Guided filter."),
|
|
.init = init,
|
|
.uninit = uninit,
|
|
.priv_size = sizeof(GuidedContext),
|
|
.priv_class = &guided_class,
|
|
.activate = activate,
|
|
.inputs = NULL,
|
|
FILTER_OUTPUTS(guided_outputs),
|
|
FILTER_PIXFMTS_ARRAY(pix_fmts),
|
|
.flags = AVFILTER_FLAG_DYNAMIC_INPUTS | AVFILTER_FLAG_SLICE_THREADS |
|
|
AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL,
|
|
.process_command = ff_filter_process_command,
|
|
};
|