| @@ -0,0 +1,440 @@ |
| 1 | +/* |
| 2 | + * Copyright (c) 2020 Paul B Mahol |
| 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 | +#include "libavutil/imgutils.h" |
| 22 | +#include "libavutil/pixdesc.h" |
| 23 | +#include "libavutil/opt.h" |
| 24 | +#include "avfilter.h" |
| 25 | +#include "formats.h" |
| 26 | +#include "internal.h" |
| 27 | +#include "video.h" |
| 28 | + |
| 29 | +typedef struct TMidEqualizerContext { |
| 30 | + const AVClass *class; |
| 31 | + |
| 32 | + int planes; |
| 33 | + int radius; |
| 34 | + float sigma; |
| 35 | + |
| 36 | + int plane_width[4], plane_height[4]; |
| 37 | + int nb_frames; |
| 38 | + int depth; |
| 39 | + int f_frames; |
| 40 | + int l_frames; |
| 41 | + int del_frame; |
| 42 | + int cur_frame; |
| 43 | + int nb_planes; |
| 44 | + int histogram_size; |
| 45 | + float kernel[127]; |
| 46 | + float *histogram[4][256]; |
| 47 | + float *change[4]; |
| 48 | + |
| 49 | + AVFrame **frames; |
| 50 | + |
| 51 | + void (*compute_histogram)(const uint8_t *ssrc, ptrdiff_t linesize, |
| 52 | + int w, int h, float *histogram, size_t hsize); |
| 53 | + void (*apply_contrast_change)(const uint8_t *src, ptrdiff_t src_linesize, |
| 54 | + uint8_t *dst, ptrdiff_t dst_linesize, |
| 55 | + int w, int h, float *change, float *orig); |
| 56 | +} TMidEqualizerContext; |
| 57 | + |
| 58 | +#define OFFSET(x) offsetof(TMidEqualizerContext, x) |
| 59 | +#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM |
| 60 | + |
| 61 | +static const AVOption tmidequalizer_options[] = { |
| 62 | + { "radius", "set radius", OFFSET(radius), AV_OPT_TYPE_INT, {.i64=5}, 1, 127, FLAGS }, |
| 63 | + { "sigma", "set sigma", OFFSET(sigma), AV_OPT_TYPE_FLOAT, {.dbl=0.5}, 0, 1, FLAGS }, |
| 64 | + { "planes", "set planes", OFFSET(planes), AV_OPT_TYPE_INT, {.i64=0xF}, 0, 0xF, FLAGS }, |
| 65 | + { NULL } |
| 66 | +}; |
| 67 | + |
| 68 | +AVFILTER_DEFINE_CLASS(tmidequalizer); |
| 69 | + |
| 70 | +static int query_formats(AVFilterContext *ctx) |
| 71 | +{ |
| 72 | + static const enum AVPixelFormat pix_fmts[] = { |
| 73 | + AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV440P, |
| 74 | + AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P, |
| 75 | + AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P, |
| 76 | + AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P, |
| 77 | + AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P, |
| 78 | + AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP, |
| 79 | + AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY14, |
| 80 | + AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9, |
| 81 | + AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10, |
| 82 | + AV_PIX_FMT_YUV420P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV444P12, |
| 83 | + AV_PIX_FMT_YUV420P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV444P14, |
| 84 | + AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, |
| 85 | + AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9, |
| 86 | + AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10, |
| 87 | + AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12, |
| 88 | + AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, |
| 89 | + AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16, |
| 90 | + AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16, |
| 91 | + AV_PIX_FMT_GBRP16, AV_PIX_FMT_GBRAP16, |
| 92 | + AV_PIX_FMT_GRAY16, |
| 93 | + AV_PIX_FMT_NONE |
| 94 | + }; |
| 95 | + |
| 96 | + return ff_set_common_formats(ctx, ff_make_format_list(pix_fmts)); |
| 97 | +} |
| 98 | + |
| 99 | +static void compute_contrast_function(const float *const histograms[256], |
| 100 | + const float *const kernel, |
| 101 | + int nb_frames, int radius, int hsize, |
| 102 | + float *f, int idx) |
| 103 | +{ |
| 104 | + const float *const h1 = histograms[idx]; |
| 105 | + int p2[256] = { 0 }; |
| 106 | + |
| 107 | + for (int p1 = 0; p1 < hsize; p1++) { |
| 108 | + float weight = 1.f; |
| 109 | + float sum = p1 * weight; |
| 110 | + |
| 111 | + for (int j = 0; j < radius; j++) { |
| 112 | + const int nidx = ((idx - radius + j) % nb_frames); |
| 113 | + const float *const h2 = histograms[nidx < 0 ? nidx + nb_frames: nidx]; |
| 114 | + int k = j; |
| 115 | + |
| 116 | + for (; p2[k] < hsize && h2[p2[k]] < h1[p1]; p2[k]++); |
| 117 | + if (p2[k] == hsize) |
| 118 | + p2[k]--; |
| 119 | + |
| 120 | + weight += kernel[j]; |
| 121 | + sum += kernel[j] * p2[k]; |
| 122 | + } |
| 123 | + |
| 124 | + for (int j = radius + 1; j < nb_frames; j++) { |
| 125 | + const int nidx = (idx - radius + j) % nb_frames; |
| 126 | + const float *const h2 = histograms[nidx < 0 ? nidx + nb_frames: nidx]; |
| 127 | + int k = j; |
| 128 | + |
| 129 | + for (; p2[k] < hsize && h2[p2[k]] < h1[p1]; p2[k]++); |
| 130 | + if (p2[k] == hsize) |
| 131 | + p2[k]--; |
| 132 | + |
| 133 | + weight += kernel[j - radius - 1]; |
| 134 | + sum += kernel[j - radius - 1] * p2[k]; |
| 135 | + } |
| 136 | + |
| 137 | + f[p1] = sum / weight; |
| 138 | + } |
| 139 | +} |
| 140 | + |
| 141 | +static void apply_contrast_change8(const uint8_t *src, ptrdiff_t src_linesize, |
| 142 | + uint8_t *dst, ptrdiff_t dst_linesize, |
| 143 | + int w, int h, float *change, float *orig) |
| 144 | +{ |
| 145 | + for (int y = 0; y < h; y++) { |
| 146 | + for (int x = 0; x < w; x++) |
| 147 | + dst[x] = lrintf(change[src[x]]); |
| 148 | + |
| 149 | + dst += dst_linesize; |
| 150 | + src += src_linesize; |
| 151 | + } |
| 152 | +} |
| 153 | + |
| 154 | +static void apply_contrast_change16(const uint8_t *ssrc, ptrdiff_t src_linesize, |
| 155 | + uint8_t *ddst, ptrdiff_t dst_linesize, |
| 156 | + int w, int h, float *change, float *orig) |
| 157 | +{ |
| 158 | + const uint16_t *src = (const uint16_t *)ssrc; |
| 159 | + uint16_t *dst = (uint16_t *)ddst; |
| 160 | + |
| 161 | + for (int y = 0; y < h; y++) { |
| 162 | + for (int x = 0; x < w; x++) |
| 163 | + dst[x] = lrintf(change[src[x]]); |
| 164 | + |
| 165 | + dst += dst_linesize / 2; |
| 166 | + src += src_linesize / 2; |
| 167 | + } |
| 168 | +} |
| 169 | + |
| 170 | +static int filter_frame(AVFilterLink *inlink, AVFrame *in) |
| 171 | +{ |
| 172 | + AVFilterContext *ctx = inlink->dst; |
| 173 | + TMidEqualizerContext *s = ctx->priv; |
| 174 | + AVFilterLink *outlink = ctx->outputs[0]; |
| 175 | + AVFrame *out; |
| 176 | + int eof = 0; |
| 177 | + |
| 178 | + if (!in) { |
| 179 | + int idx = s->f_frames < s->nb_frames ? s->radius : s->del_frame ? s->del_frame - 1 : s->nb_frames - 1; |
| 180 | + |
| 181 | + if (s->f_frames < s->nb_frames) { |
| 182 | + s->l_frames = s->nb_frames - s->f_frames; |
| 183 | + } else { |
| 184 | + s->l_frames++; |
| 185 | + } |
| 186 | + in = av_frame_clone(s->frames[idx]); |
| 187 | + if (!in) |
| 188 | + return AVERROR(ENOMEM); |
| 189 | + eof = 1; |
| 190 | + } |
| 191 | + |
| 192 | + if (s->f_frames < s->nb_frames) { |
| 193 | + s->frames[s->f_frames] = in; |
| 194 | + |
| 195 | + for (int p = 0; p < s->nb_planes; p++) { |
| 196 | + s->compute_histogram(in->data[p], in->linesize[p], |
| 197 | + s->plane_width[p], s->plane_height[p], |
| 198 | + s->histogram[p][s->f_frames], |
| 199 | + s->histogram_size); |
| 200 | + } |
| 201 | + |
| 202 | + s->f_frames++; |
| 203 | + |
| 204 | + while (s->f_frames <= s->radius) { |
| 205 | + s->frames[s->f_frames] = av_frame_clone(in); |
| 206 | + if (!s->frames[s->f_frames]) |
| 207 | + return AVERROR(ENOMEM); |
| 208 | + for (int p = 0; p < s->nb_planes; p++) { |
| 209 | + memcpy(s->histogram[p][s->f_frames], |
| 210 | + s->histogram[p][s->f_frames - 1], |
| 211 | + s->histogram_size * sizeof(float)); |
| 212 | + } |
| 213 | + s->f_frames++; |
| 214 | + } |
| 215 | + |
| 216 | + if (!eof && s->f_frames < s->nb_frames) { |
| 217 | + return 0; |
| 218 | + } else { |
| 219 | + while (s->f_frames < s->nb_frames) { |
| 220 | + s->frames[s->f_frames] = av_frame_clone(in); |
| 221 | + if (!s->frames[s->f_frames]) |
| 222 | + return AVERROR(ENOMEM); |
| 223 | + for (int p = 0; p < s->nb_planes; p++) { |
| 224 | + memcpy(s->histogram[p][s->f_frames], |
| 225 | + s->histogram[p][s->f_frames - 1], |
| 226 | + s->histogram_size * sizeof(float)); |
| 227 | + } |
| 228 | + s->f_frames++; |
| 229 | + } |
| 230 | + } |
| 231 | + s->cur_frame = s->radius; |
| 232 | + s->del_frame = 0; |
| 233 | + } else { |
| 234 | + av_frame_free(&s->frames[s->del_frame]); |
| 235 | + s->frames[s->del_frame] = in; |
| 236 | + |
| 237 | + for (int p = 0; p < s->nb_planes; p++) { |
| 238 | + s->compute_histogram(in->data[p], in->linesize[p], |
| 239 | + s->plane_width[p], s->plane_height[p], |
| 240 | + s->histogram[p][s->del_frame], |
| 241 | + s->histogram_size); |
| 242 | + } |
| 243 | + |
| 244 | + s->del_frame++; |
| 245 | + if (s->del_frame >= s->nb_frames) |
| 246 | + s->del_frame = 0; |
| 247 | + } |
| 248 | + |
| 249 | + if (ctx->is_disabled) { |
| 250 | + const int idx = s->cur_frame; |
| 251 | + |
| 252 | + out = av_frame_clone(s->frames[idx]); |
| 253 | + if (!out) |
| 254 | + return AVERROR(ENOMEM); |
| 255 | + } else { |
| 256 | + const int idx = s->cur_frame; |
| 257 | + |
| 258 | + in = s->frames[idx]; |
| 259 | + out = ff_get_video_buffer(outlink, outlink->w, outlink->h); |
| 260 | + if (!out) |
| 261 | + return AVERROR(ENOMEM); |
| 262 | + av_frame_copy_props(out, in); |
| 263 | + |
| 264 | + for (int p = 0; p < s->nb_planes; p++) { |
| 265 | + if (!((1 << p) & s->planes)) { |
| 266 | + av_image_copy_plane(out->data[p], out->linesize[p], in->data[p], in->linesize[p], |
| 267 | + s->plane_width[p] * (1 + (s->depth > 8)), s->plane_height[p]); |
| 268 | + continue; |
| 269 | + } |
| 270 | + |
| 271 | + compute_contrast_function((const float *const *)s->histogram[p], s->kernel, |
| 272 | + s->nb_frames, s->radius, s->histogram_size, s->change[p], idx); |
| 273 | + |
| 274 | + s->apply_contrast_change(in->data[p], in->linesize[p], |
| 275 | + out->data[p], out->linesize[p], |
| 276 | + s->plane_width[p], s->plane_height[p], |
| 277 | + s->change[p], s->histogram[p][idx]); |
| 278 | + } |
| 279 | + } |
| 280 | + |
| 281 | + s->cur_frame++; |
| 282 | + if (s->cur_frame >= s->nb_frames) |
| 283 | + s->cur_frame = 0; |
| 284 | + |
| 285 | + return ff_filter_frame(outlink, out); |
| 286 | +} |
| 287 | + |
| 288 | +static void compute_histogram8(const uint8_t *src, ptrdiff_t linesize, |
| 289 | + int w, int h, float *histogram, size_t hsize) |
| 290 | +{ |
| 291 | + memset(histogram, 0, hsize * sizeof(*histogram)); |
| 292 | + |
| 293 | + for (int y = 0; y < h; y++) { |
| 294 | + for (int x = 0; x < w; x++) |
| 295 | + histogram[src[x]] += 1; |
| 296 | + src += linesize; |
| 297 | + } |
| 298 | + |
| 299 | + for (int x = 0; x < hsize; x++) |
| 300 | + histogram[x] /= hsize; |
| 301 | + |
| 302 | + for (int x = 1; x < hsize; x++) |
| 303 | + histogram[x] += histogram[x-1]; |
| 304 | +} |
| 305 | + |
| 306 | +static void compute_histogram16(const uint8_t *ssrc, ptrdiff_t linesize, |
| 307 | + int w, int h, float *histogram, size_t hsize) |
| 308 | +{ |
| 309 | + const uint16_t *src = (const uint16_t *)ssrc; |
| 310 | + |
| 311 | + memset(histogram, 0, hsize * sizeof(*histogram)); |
| 312 | + |
| 313 | + for (int y = 0; y < h; y++) { |
| 314 | + for (int x = 0; x < w; x++) |
| 315 | + histogram[src[x]] += 1; |
| 316 | + src += linesize / 2; |
| 317 | + } |
| 318 | + |
| 319 | + for (int x = 0; x < hsize; x++) |
| 320 | + histogram[x] /= hsize; |
| 321 | + |
| 322 | + for (int x = 1; x < hsize; x++) |
| 323 | + histogram[x] += histogram[x-1]; |
| 324 | +} |
| 325 | + |
| 326 | +static int config_input(AVFilterLink *inlink) |
| 327 | +{ |
| 328 | + AVFilterContext *ctx = inlink->dst; |
| 329 | + TMidEqualizerContext *s = ctx->priv; |
| 330 | + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format); |
| 331 | + float sigma = s->radius * s->sigma; |
| 332 | + int vsub, hsub; |
| 333 | + |
| 334 | + s->depth = desc->comp[0].depth; |
| 335 | + s->nb_frames = s->radius * 2 + 1; |
| 336 | + s->nb_planes = av_pix_fmt_count_planes(inlink->format); |
| 337 | + |
| 338 | + hsub = desc->log2_chroma_w; |
| 339 | + vsub = desc->log2_chroma_h; |
| 340 | + |
| 341 | + s->plane_height[0] = s->plane_height[3] = inlink->h; |
| 342 | + s->plane_width[0] = s->plane_width[3] = inlink->w; |
| 343 | + s->plane_height[1] = s->plane_height[2] = AV_CEIL_RSHIFT(inlink->h, vsub); |
| 344 | + s->plane_width[1] = s->plane_width[2] = AV_CEIL_RSHIFT(inlink->w, hsub); |
| 345 | + |
| 346 | + s->histogram_size = 1 << s->depth; |
| 347 | + |
| 348 | + for (int n = 0; n < s->radius; n++) |
| 349 | + s->kernel[n] = expf(-0.5 * (n + 1) * (n + 1) / (sigma * sigma)); |
| 350 | + |
| 351 | + for (int p = 0; p < s->nb_planes; p++) { |
| 352 | + for (int n = 0; n < s->nb_frames; n++) { |
| 353 | + s->histogram[p][n] = av_calloc(s->histogram_size, sizeof(float)); |
| 354 | + if (!s->histogram[p][n]) |
| 355 | + return AVERROR(ENOMEM); |
| 356 | + } |
| 357 | + |
| 358 | + s->change[p] = av_calloc(s->histogram_size, sizeof(float)); |
| 359 | + if (!s->change[p]) |
| 360 | + return AVERROR(ENOMEM); |
| 361 | + } |
| 362 | + |
| 363 | + if (!s->frames) |
| 364 | + s->frames = av_calloc(s->nb_frames, sizeof(*s->frames)); |
| 365 | + if (!s->frames) |
| 366 | + return AVERROR(ENOMEM); |
| 367 | + |
| 368 | + s->compute_histogram = s->depth <= 8 ? compute_histogram8 : compute_histogram16; |
| 369 | + s->apply_contrast_change = s->depth <= 8 ? apply_contrast_change8 : apply_contrast_change16; |
| 370 | + |
| 371 | + return 0; |
| 372 | +} |
| 373 | + |
| 374 | +static int request_frame(AVFilterLink *outlink) |
| 375 | +{ |
| 376 | + AVFilterContext *ctx = outlink->src; |
| 377 | + TMidEqualizerContext *s = ctx->priv; |
| 378 | + int ret; |
| 379 | + |
| 380 | + ret = ff_request_frame(ctx->inputs[0]); |
| 381 | + if (ret == AVERROR_EOF && s->l_frames < s->radius) { |
| 382 | + ret = filter_frame(ctx->inputs[0], NULL); |
| 383 | + } |
| 384 | + |
| 385 | + return ret; |
| 386 | +} |
| 387 | + |
| 388 | +static void free_histograms(AVFilterContext *ctx, int x, int nb_frames) |
| 389 | +{ |
| 390 | + TMidEqualizerContext *s = ctx->priv; |
| 391 | + |
| 392 | + for (int n = 0; n < nb_frames; n++) |
| 393 | + av_freep(&s->histogram[x][n]); |
| 394 | + av_freep(&s->change[x]); |
| 395 | +} |
| 396 | + |
| 397 | +static av_cold void uninit(AVFilterContext *ctx) |
| 398 | +{ |
| 399 | + TMidEqualizerContext *s = ctx->priv; |
| 400 | + |
| 401 | + free_histograms(ctx, 0, s->nb_frames); |
| 402 | + free_histograms(ctx, 1, s->nb_frames); |
| 403 | + free_histograms(ctx, 2, s->nb_frames); |
| 404 | + free_histograms(ctx, 3, s->nb_frames); |
| 405 | + |
| 406 | + for (int i = 0; i < s->nb_frames && s->frames; i++) |
| 407 | + av_frame_free(&s->frames[i]); |
| 408 | + av_freep(&s->frames); |
| 409 | +} |
| 410 | + |
| 411 | +static const AVFilterPad tmidequalizer_inputs[] = { |
| 412 | + { |
| 413 | + .name = "default", |
| 414 | + .type = AVMEDIA_TYPE_VIDEO, |
| 415 | + .config_props = config_input, |
| 416 | + .filter_frame = filter_frame, |
| 417 | + }, |
| 418 | + { NULL } |
| 419 | +}; |
| 420 | + |
| 421 | +static const AVFilterPad tmidequalizer_outputs[] = { |
| 422 | + { |
| 423 | + .name = "default", |
| 424 | + .type = AVMEDIA_TYPE_VIDEO, |
| 425 | + .request_frame = request_frame, |
| 426 | + }, |
| 427 | + { NULL } |
| 428 | +}; |
| 429 | + |
| 430 | +AVFilter ff_vf_tmidequalizer = { |
| 431 | + .name = "tmidequalizer", |
| 432 | + .description = NULL_IF_CONFIG_SMALL("Apply Temporal Midway Equalization."), |
| 433 | + .priv_size = sizeof(TMidEqualizerContext), |
| 434 | + .uninit = uninit, |
| 435 | + .query_formats = query_formats, |
| 436 | + .inputs = tmidequalizer_inputs, |
| 437 | + .outputs = tmidequalizer_outputs, |
| 438 | + .priv_class = &tmidequalizer_class, |
| 439 | + .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL, |
| 440 | +}; |