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

   1 /*
   2  * Copyright (C) 2010-2011 Kevin Stone
   3  * Copyright (C) 2016 Paul B Mahol
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or modify
   8  * it under the terms of the GNU General Public License as published by
   9  * the Free Software Foundation; either version 2 of the License, or
  10  * (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15  * GNU General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU General Public License along
  18  * with FFmpeg; if not, write to the Free Software Foundation, Inc.,
  19  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  20  */
  21
  22 #include <float.h>
  23
  24 #include "libavutil/avassert.h"
  25 #include "libavutil/common.h"
  26 #include "libavutil/float_dsp.h"
  27 #include "libavutil/imgutils.h"
  28 #include "libavutil/mem_internal.h"
  29 #include "libavutil/opt.h"
  30 #include "libavutil/pixdesc.h"
  31 #include "avfilter.h"
  32 #include "formats.h"
  33 #include "internal.h"
  34 #include "video.h"
  35
  36 static const size_t NNEDI_WEIGHTS_SIZE = 13574928;
  37 static const uint8_t NNEDI_XDIM[] = { 8, 16, 32, 48, 8, 16, 32 };
  38 static const uint8_t NNEDI_YDIM[] = { 6, 6, 6, 6, 4, 4, 4 };
  39 static const uint16_t NNEDI_NNS[] = { 16, 32, 64, 128, 256 };
  40
  41 typedef struct PrescreenerCoefficients {
  42     DECLARE_ALIGNED(32, float, kernel_l0)[4][16 * 4];
  43     DECLARE_ALIGNED(32, float, bias_l0)[4];
  44
  45     DECLARE_ALIGNED(32, float, kernel_l1)[4][4];
  46     DECLARE_ALIGNED(32, float, bias_l1)[4];
  47
  48     DECLARE_ALIGNED(32, float, kernel_l2)[4][8];
  49     DECLARE_ALIGNED(32, float, bias_l2)[4];
  50 } PrescreenerCoefficients;
  51
  52 typedef struct PredictorCoefficients {
  53     int xdim, ydim, nns, nsize;
  54     float *data;
  55     float *softmax_q1;
  56     float *elliott_q1;
  57     float *softmax_bias_q1;
  58     float *elliott_bias_q1;
  59     float *softmax_q2;
  60     float *elliott_q2;
  61     float *softmax_bias_q2;
  62     float *elliott_bias_q2;
  63 } PredictorCoefficients;
  64
  65 typedef struct NNEDIContext {
  66     const AVClass *class;
  67
  68     char *weights_file;
  69
  70     AVFrame *prev;
  71     int eof;
  72     int64_t pts;
  73
  74     AVFloatDSPContext *fdsp;
  75     int depth;
  76     int nb_planes;
  77     int nb_threads;
  78     int linesize[4];
  79     int planewidth[4];
  80     int planeheight[4];
  81     int field_n;
  82
  83     PrescreenerCoefficients prescreener[4];
  84     PredictorCoefficients coeffs[2][5][7];
  85
  86     float half;
  87     float in_scale;
  88     float out_scale;
  89
  90     // Parameters
  91     int deint;
  92     int field;
  93     int process_plane;
  94     int nsize;
  95     int nnsparam;
  96     int qual;
  97     int etype;
  98     int pscrn;
  99
 100     int input_size;
 101     uint8_t **prescreen_buf;
 102     float **input_buf;
 103     float **output_buf;
 104
 105     void (*read)(const uint8_t *src, float *dst,
 106                  int src_stride, int dst_stride,
 107                  int width, int height, float scale);
 108     void (*write)(const float *src, uint8_t *dst,
 109                   int src_stride, int dst_stride,
 110                   int width, int height, int depth, float scale);
 111     void (*prescreen[2])(AVFilterContext *ctx,
 112                          const void *src, ptrdiff_t src_stride,
 113                          uint8_t *prescreen, int N,
 114                          const PrescreenerCoefficients *const coeffs);
 115 } NNEDIContext;
 116
 117 #define OFFSET(x) offsetof(NNEDIContext, x)
 118 #define RFLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
 119 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
 120
 121 static const AVOption nnedi_options[] = {
 122     {"weights",  "set weights file", OFFSET(weights_file),  AV_OPT_TYPE_STRING, {.str="nnedi3_weights.bin"}, 0, 0, FLAGS },
 123     {"deint",         "set which frames to deinterlace", OFFSET(deint),         AV_OPT_TYPE_INT, {.i64=0}, 0, 1, RFLAGS, "deint" },
 124         {"all",        "deinterlace all frames",                       0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "deint" },
 125         {"interlaced", "only deinterlace frames marked as interlaced", 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "deint" },
 126     {"field",  "set mode of operation", OFFSET(field),         AV_OPT_TYPE_INT, {.i64=-1}, -2, 3, RFLAGS, "field" },
 127         {"af", "use frame flags, both fields",  0, AV_OPT_TYPE_CONST, {.i64=-2}, 0, 0, RFLAGS, "field" },
 128         {"a",  "use frame flags, single field", 0, AV_OPT_TYPE_CONST, {.i64=-1}, 0, 0, RFLAGS, "field" },
 129         {"t",  "use top field only",            0, AV_OPT_TYPE_CONST, {.i64=0},  0, 0, RFLAGS, "field" },
 130         {"b",  "use bottom field only",         0, AV_OPT_TYPE_CONST, {.i64=1},  0, 0, RFLAGS, "field" },
 131         {"tf", "use both fields, top first",    0, AV_OPT_TYPE_CONST, {.i64=2},  0, 0, RFLAGS, "field" },
 132         {"bf", "use both fields, bottom first", 0, AV_OPT_TYPE_CONST, {.i64=3},  0, 0, RFLAGS, "field" },
 133     {"planes", "set which planes to process", OFFSET(process_plane), AV_OPT_TYPE_INT, {.i64=7}, 0, 15, RFLAGS },
 134     {"nsize",  "set size of local neighborhood around each pixel, used by the predictor neural network", OFFSET(nsize), AV_OPT_TYPE_INT, {.i64=6}, 0, 6, RFLAGS, "nsize" },
 135         {"s8x6",     NULL, 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "nsize" },
 136         {"s16x6",    NULL, 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "nsize" },
 137         {"s32x6",    NULL, 0, AV_OPT_TYPE_CONST, {.i64=2}, 0, 0, RFLAGS, "nsize" },
 138         {"s48x6",    NULL, 0, AV_OPT_TYPE_CONST, {.i64=3}, 0, 0, RFLAGS, "nsize" },
 139         {"s8x4",     NULL, 0, AV_OPT_TYPE_CONST, {.i64=4}, 0, 0, RFLAGS, "nsize" },
 140         {"s16x4",    NULL, 0, AV_OPT_TYPE_CONST, {.i64=5}, 0, 0, RFLAGS, "nsize" },
 141         {"s32x4",    NULL, 0, AV_OPT_TYPE_CONST, {.i64=6}, 0, 0, RFLAGS, "nsize" },
 142     {"nns",    "set number of neurons in predictor neural network", OFFSET(nnsparam), AV_OPT_TYPE_INT, {.i64=1}, 0, 4, RFLAGS, "nns" },
 143         {"n16",       NULL, 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "nns" },
 144         {"n32",       NULL, 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "nns" },
 145         {"n64",       NULL, 0, AV_OPT_TYPE_CONST, {.i64=2}, 0, 0, RFLAGS, "nns" },
 146         {"n128",      NULL, 0, AV_OPT_TYPE_CONST, {.i64=3}, 0, 0, RFLAGS, "nns" },
 147         {"n256",      NULL, 0, AV_OPT_TYPE_CONST, {.i64=4}, 0, 0, RFLAGS, "nns" },
 148     {"qual",  "set quality", OFFSET(qual), AV_OPT_TYPE_INT, {.i64=1}, 1, 2, RFLAGS, "qual" },
 149         {"fast", NULL, 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "qual" },
 150         {"slow", NULL, 0, AV_OPT_TYPE_CONST, {.i64=2}, 0, 0, RFLAGS, "qual" },
 151     {"etype", "set which set of weights to use in the predictor", OFFSET(etype), AV_OPT_TYPE_INT, {.i64=0}, 0, 1, RFLAGS, "etype" },
 152         {"a",  "weights trained to minimize absolute error", 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "etype" },
 153         {"abs","weights trained to minimize absolute error", 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "etype" },
 154         {"s",  "weights trained to minimize squared error",  0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "etype" },
 155         {"mse","weights trained to minimize squared error",  0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "etype" },
 156     {"pscrn", "set prescreening", OFFSET(pscrn), AV_OPT_TYPE_INT, {.i64=2}, 0, 4, RFLAGS, "pscrn" },
 157         {"none",      NULL, 0, AV_OPT_TYPE_CONST, {.i64=0}, 0, 0, RFLAGS, "pscrn" },
 158         {"original",  NULL, 0, AV_OPT_TYPE_CONST, {.i64=1}, 0, 0, RFLAGS, "pscrn" },
 159         {"new",       NULL, 0, AV_OPT_TYPE_CONST, {.i64=2}, 0, 0, RFLAGS, "pscrn" },
 160         {"new2",      NULL, 0, AV_OPT_TYPE_CONST, {.i64=3}, 0, 0, RFLAGS, "pscrn" },
 161         {"new3",      NULL, 0, AV_OPT_TYPE_CONST, {.i64=4}, 0, 0, RFLAGS, "pscrn" },
 162     { NULL }
 163 };
 164
 165 AVFILTER_DEFINE_CLASS(nnedi);
 166
 167 static int config_output(AVFilterLink *outlink)
 168 {
 169     AVFilterContext *ctx = outlink->src;
 170
 171     outlink->time_base.num = ctx->inputs[0]->time_base.num;
 172     outlink->time_base.den = ctx->inputs[0]->time_base.den * 2;
 173     outlink->w             = ctx->inputs[0]->w;
 174     outlink->h             = ctx->inputs[0]->h;
 175
 176     outlink->frame_rate = av_mul_q(ctx->inputs[0]->frame_rate,
 177                                    (AVRational){2, 1});
 178
 179     return 0;
 180 }
 181
 182 static int query_formats(AVFilterContext *ctx)
 183 {
 184     static const enum AVPixelFormat pix_fmts[] = {
 185         AV_PIX_FMT_GRAY8,
 186         AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY14, AV_PIX_FMT_GRAY16,
 187         AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P,
 188         AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P,
 189         AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV444P,
 190         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P,
 191         AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P,
 192         AV_PIX_FMT_YUVJ411P,
 193         AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUVA444P,
 194         AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP,
 195         AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9,
 196         AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
 197         AV_PIX_FMT_YUV440P10,
 198         AV_PIX_FMT_YUV420P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV444P12,
 199         AV_PIX_FMT_YUV440P12,
 200         AV_PIX_FMT_YUV420P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV444P14,
 201         AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16,
 202         AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10, AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
 203         AV_PIX_FMT_YUVA444P9, AV_PIX_FMT_YUVA444P10, AV_PIX_FMT_YUVA444P12, AV_PIX_FMT_YUVA444P16,
 204         AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA422P16,
 205         AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA420P16,
 206         AV_PIX_FMT_GBRAP10,   AV_PIX_FMT_GBRAP12,    AV_PIX_FMT_GBRAP16,
 207         AV_PIX_FMT_NONE
 208     };
 209
 210     AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
 211     if (!fmts_list)
 212         return AVERROR(ENOMEM);
 213     return ff_set_common_formats(ctx, fmts_list);
 214 }
 215
 216 static float dot_dsp(const NNEDIContext *const s, const float *kernel, const float *input,
 217                      int n, float scale, float bias)
 218 {
 219     float sum, y;
 220
 221     sum = s->fdsp->scalarproduct_float(kernel, input, n);
 222
 223     y = sum * scale + bias + 1e-20f;
 224
 225     return y;
 226 }
 227
 228 static float elliott(float x)
 229 {
 230     return x / (1.0f + fabsf(x));
 231 }
 232
 233 static void transform_elliott(float *input, int size)
 234 {
 235     for (int i = 0; i < size; i++)
 236         input[i] = elliott(input[i]);
 237 }
 238
 239 static void process_old(AVFilterContext *ctx,
 240                         const void *src, ptrdiff_t src_stride,
 241                         uint8_t *prescreen, int N,
 242                         const PrescreenerCoefficients *const m_data)
 243 {
 244     NNEDIContext *s = ctx->priv;
 245     const float *src_p = src;
 246
 247     // Adjust source pointer to point to top-left of filter window.
 248     const float *window = src_p - 2 * src_stride - 5;
 249
 250     for (int j = 0; j < N; j++) {
 251         LOCAL_ALIGNED_32(float, input, [48]);
 252         float state[12];
 253
 254         for (int i = 0; i < 4; i++)
 255             memcpy(input + i * 12, window + i * src_stride + j, 12 * sizeof(float));
 256
 257         // Layer 0.
 258         for (int n = 0; n < 4; n++)
 259             state[n] = dot_dsp(s, m_data->kernel_l0[n], input, 48, 1.0f, m_data->bias_l0[n]);
 260         transform_elliott(state + 1, 3);
 261
 262         // Layer 1.
 263         for (int n = 0; n < 4; n++)
 264             state[n + 4] = dot_dsp(s, m_data->kernel_l1[n], state, 4, 1.0f, m_data->bias_l1[n]);
 265         transform_elliott(state + 4, 3);
 266
 267         // Layer 2.
 268         for (int n = 0; n < 4; n++)
 269             state[n + 8] = dot_dsp(s, m_data->kernel_l2[n], state, 8, 1.0f, m_data->bias_l2[n]);
 270
 271         prescreen[j] = FFMAX(state[10], state[11]) <= FFMAX(state[8], state[9]) ? 255 : 0;
 272     }
 273 }
 274
 275 static void process_new(AVFilterContext *ctx,
 276                         const void *src, ptrdiff_t src_stride,
 277                         uint8_t *prescreen, int N,
 278                         const PrescreenerCoefficients *const m_data)
 279 {
 280     NNEDIContext *s = ctx->priv;
 281     const float *src_p = src;
 282
 283     // Adjust source pointer to point to top-left of filter window.
 284     const float *window = src_p - 2 * src_stride - 6;
 285
 286     for (int j = 0; j < N; j += 4) {
 287         LOCAL_ALIGNED_32(float, input, [64]);
 288         float state[8];
 289
 290         for (int i = 0; i < 4; i++)
 291             memcpy(input + i * 16, window + i * src_stride + j, 16 * sizeof(float));
 292
 293         for (int n = 0; n < 4; n++)
 294             state[n] = dot_dsp(s, m_data->kernel_l0[n], input, 64, 1.0f, m_data->bias_l0[n]);
 295         transform_elliott(state, 4);
 296
 297         for (int n = 0; n < 4; n++)
 298             state[n + 4] = dot_dsp(s, m_data->kernel_l1[n], state, 4, 1.0f, m_data->bias_l1[n]);
 299
 300         for (int n = 0; n < 4; n++)
 301             prescreen[j + n] = state[n + 4] > 0.f;
 302     }
 303 }
 304
 305 static int filter_offset(int nn, const PredictorCoefficients *const model)
 306 {
 307     return nn * model->nsize;
 308 }
 309
 310 static const float *softmax_q1_filter(int nn,
 311                                       const PredictorCoefficients *const model)
 312 {
 313     return model->softmax_q1 + filter_offset(nn, model);
 314 }
 315
 316 static const float *elliott_q1_filter(int nn,
 317                                       const PredictorCoefficients *const model)
 318 {
 319     return model->elliott_q1 + filter_offset(nn, model);
 320 }
 321
 322 static const float *softmax_q2_filter(int nn,
 323                                       const PredictorCoefficients *const model)
 324 {
 325     return model->softmax_q2 + filter_offset(nn, model);
 326 }
 327
 328 static const float *elliott_q2_filter(int nn,
 329                                       const PredictorCoefficients *const model)
 330 {
 331     return model->elliott_q2 + filter_offset(nn, model);
 332 }
 333
 334 static void gather_input(const float *src, ptrdiff_t src_stride,
 335                          float *buf, float mstd[4],
 336                          const PredictorCoefficients *const model)
 337 {
 338     const float scale = 1.f / model->nsize;
 339     float sum = 0.f;
 340     float sum_sq = 0.f;
 341     float tmp;
 342
 343     for (int i = 0; i < model->ydim; i++) {
 344         memcpy(buf, src, model->xdim * sizeof(float));
 345
 346         for (int j = 0; j < model->xdim; j++) {
 347             const float val = src[j];
 348
 349             sum += val;
 350             sum_sq += val * val;
 351         }
 352
 353         src += src_stride;
 354         buf += model->xdim;
 355     }
 356
 357     mstd[0] = sum * scale;
 358     mstd[3] = 0.f;
 359
 360     tmp = sum_sq * scale - mstd[0] * mstd[0];
 361     if (tmp < FLT_EPSILON) {
 362         mstd[1] = 0.0f;
 363         mstd[2] = 0.0f;
 364     } else {
 365         mstd[1] = sqrtf(tmp);
 366         mstd[2] = 1.0f / mstd[1];
 367     }
 368 }
 369
 370 static float softmax_exp(float x)
 371 {
 372     return expf(av_clipf(x, -80.f, 80.f));
 373 }
 374
 375 static void transform_softmax_exp(float *input, int size)
 376 {
 377     for (int i = 0; i < size; i++)
 378         input[i] = softmax_exp(input[i]);
 379 }
 380
 381 static void wae5(const float *softmax, const float *el,
 382                  int n, float mstd[4])
 383 {
 384     float vsum = 0.0f, wsum = 0.0f;
 385
 386     for (int i = 0; i < n; i++) {
 387         vsum += softmax[i] * elliott(el[i]);
 388         wsum += softmax[i];
 389     }
 390
 391     if (wsum > 1e-10f)
 392         mstd[3] += (5.0f * vsum) / wsum * mstd[1] + mstd[0];
 393     else
 394         mstd[3] += mstd[0];
 395 }
 396
 397 static void predictor(AVFilterContext *ctx,
 398                       const void *src, ptrdiff_t src_stride, void *dst,
 399                       const uint8_t *prescreen, int N,
 400                       const PredictorCoefficients *const model, int use_q2)
 401 {
 402     const NNEDIContext *const s = ctx->priv;
 403     const float *src_p = src;
 404     float *dst_p = dst;
 405
 406     // Adjust source pointer to point to top-left of filter window.
 407     const float *window = src_p - (model->ydim / 2) * src_stride - (model->xdim / 2 - 1);
 408     const int filter_size = model->nsize;
 409     const int nns = model->nns;
 410
 411     for (int i = 0; i < N; i++) {
 412         LOCAL_ALIGNED_32(float, input, [48 * 6]);
 413         float activation[256 * 2];
 414         float mstd[4];
 415         float scale;
 416
 417         if (prescreen[i])
 418             continue;
 419
 420         gather_input(window + i, src_stride, input, mstd, model);
 421         scale = mstd[2];
 422
 423         for (int nn = 0; nn < nns; nn++)
 424             activation[nn] = dot_dsp(s, softmax_q1_filter(nn, model), input, filter_size, scale, model->softmax_bias_q1[nn]);
 425
 426         for (int nn = 0; nn < nns; nn++)
 427             activation[nns + nn] = dot_dsp(s, elliott_q1_filter(nn, model), input, filter_size, scale, model->elliott_bias_q1[nn]);
 428
 429         transform_softmax_exp(activation, nns);
 430         wae5(activation, activation + nns, nns, mstd);
 431
 432         if (use_q2) {
 433             for (int nn = 0; nn < nns; nn++)
 434                 activation[nn] = dot_dsp(s, softmax_q2_filter(nn, model), input, filter_size, scale, model->softmax_bias_q2[nn]);
 435
 436             for (int nn = 0; nn < nns; nn++)
 437                 activation[nns + nn] = dot_dsp(s, elliott_q2_filter(nn, model), input, filter_size, scale, model->elliott_bias_q2[nn]);
 438
 439             transform_softmax_exp(activation, nns);
 440             wae5(activation, activation + nns, nns, mstd);
 441         }
 442
 443         dst_p[i] = mstd[3] * (use_q2 ? 0.5f : 1.f);
 444     }
 445 }
 446
 447 static void read_bytes(const uint8_t *src, float *dst,
 448                        int src_stride, int dst_stride,
 449                        int width, int height, float scale)
 450 {
 451     for (int y = 0; y < height; y++) {
 452         for (int x = 0; x < 32; x++)
 453             dst[-x - 1] = src[x];
 454
 455         for (int x = 0; x < width; x++)
 456             dst[x] = src[x];
 457
 458         for (int x = 0; x < 32; x++)
 459             dst[width + x] = src[width - x - 1];
 460
 461         dst += dst_stride;
 462         src += src_stride;
 463     }
 464 }
 465
 466 static void read_words(const uint8_t *srcp, float *dst,
 467                        int src_stride, int dst_stride,
 468                        int width, int height, float scale)
 469 {
 470     const uint16_t *src = (const uint16_t *)srcp;
 471
 472     src_stride /= 2;
 473
 474     for (int y = 0; y < height; y++) {
 475         for (int x = 0; x < 32; x++)
 476             dst[-x - 1] = src[x] * scale;
 477
 478         for (int x = 0; x < width; x++)
 479             dst[x] = src[x] * scale;
 480
 481         for (int x = 0; x < 32; x++)
 482             dst[width + x] = src[width - x - 1] * scale;
 483
 484         dst += dst_stride;
 485         src += src_stride;
 486     }
 487 }
 488
 489 static void write_bytes(const float *src, uint8_t *dst,
 490                         int src_stride, int dst_stride,
 491                         int width, int height, int depth,
 492                         float scale)
 493 {
 494     for (int y = 0; y < height; y++) {
 495         for (int x = 0; x < width; x++)
 496             dst[x] = av_clip_uint8(src[x]);
 497
 498         dst += dst_stride;
 499         src += src_stride;
 500     }
 501 }
 502
 503 static void write_words(const float *src, uint8_t *dstp,
 504                         int src_stride, int dst_stride,
 505                         int width, int height, int depth,
 506                         float scale)
 507 {
 508     uint16_t *dst = (uint16_t *)dstp;
 509
 510     dst_stride /= 2;
 511
 512     for (int y = 0; y < height; y++) {
 513         for (int x = 0; x < width; x++)
 514             dst[x] = av_clip_uintp2_c(src[x] * scale, depth);
 515
 516         dst += dst_stride;
 517         src += src_stride;
 518     }
 519 }
 520
 521 static void interpolation(const void *src, ptrdiff_t src_stride,
 522                           void *dst, const uint8_t *prescreen, int n)
 523 {
 524     const float *src_p = src;
 525     float *dst_p = dst;
 526     const float *window = src_p - 2 * src_stride;
 527
 528     for (int i = 0; i < n; i++) {
 529         float accum = 0.0f;
 530
 531         if (!prescreen[i])
 532             continue;
 533
 534         accum += (-3.0f / 32.0f) * window[0 * src_stride + i];
 535         accum += (19.0f / 32.0f) * window[1 * src_stride + i];
 536         accum += (19.0f / 32.0f) * window[2 * src_stride + i];
 537         accum += (-3.0f / 32.0f) * window[3 * src_stride + i];
 538
 539         dst_p[i] = accum;
 540     }
 541 }
 542
 543 static int filter_slice(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 544 {
 545     const NNEDIContext *const s = ctx->priv;
 546     AVFrame *out = arg;
 547     AVFrame *in = s->prev;
 548     const float in_scale = s->in_scale;
 549     const float out_scale = s->out_scale;
 550     const int depth = s->depth;
 551     const int interlaced = in->interlaced_frame;
 552     const int tff = s->field_n == (s->field < 0 ? interlaced ? in->top_field_first : 1 :
 553                                   (s->field & 1) ^ 1);
 554
 555
 556     for (int p = 0; p < s->nb_planes; p++) {
 557         const int height = s->planeheight[p];
 558         const int width = s->planewidth[p];
 559         const int slice_start = 2 * ((height / 2 * jobnr) / nb_jobs);
 560         const int slice_end = 2 * ((height / 2 * (jobnr+1)) / nb_jobs);
 561         const uint8_t *src_data = in->data[p];
 562         uint8_t *dst_data = out->data[p];
 563         uint8_t *dst = out->data[p] + slice_start * out->linesize[p];
 564         const int src_linesize = in->linesize[p];
 565         const int dst_linesize = out->linesize[p];
 566         uint8_t *prescreen_buf = s->prescreen_buf[jobnr];
 567         float *srcbuf = s->input_buf[jobnr];
 568         const int srcbuf_stride = width + 64;
 569         float *dstbuf = s->output_buf[jobnr];
 570         const int dstbuf_stride = width;
 571         const int slice_height = (slice_end - slice_start) / 2;
 572         const int last_slice = slice_end == height;
 573         const uint8_t *in_line;
 574         uint8_t *out_line;
 575         int y_out;
 576
 577         if (!(s->process_plane & (1 << p))) {
 578             av_image_copy_plane(dst, out->linesize[p],
 579                                 in->data[p] + slice_start * in->linesize[p],
 580                                 in->linesize[p],
 581                                 s->linesize[p], slice_end - slice_start);
 582             continue;
 583         }
 584
 585         y_out    = slice_start + (tff ^ (slice_start & 1));
 586         in_line  = src_data + (y_out * src_linesize);
 587         out_line = dst_data + (y_out * dst_linesize);
 588
 589         while (y_out < slice_end) {
 590             memcpy(out_line, in_line, s->linesize[p]);
 591             y_out += 2;
 592             in_line  += src_linesize * 2;
 593             out_line += dst_linesize * 2;
 594         }
 595
 596         y_out = slice_start + ((!tff) ^ (slice_start & 1));
 597
 598         s->read(src_data + FFMAX(y_out - 5, tff) * src_linesize,
 599                 srcbuf + 32,
 600                 src_linesize * 2, srcbuf_stride,
 601                 width, 1, in_scale);
 602         srcbuf += srcbuf_stride;
 603
 604         s->read(src_data + FFMAX(y_out - 3, tff) * src_linesize,
 605                 srcbuf + 32,
 606                 src_linesize * 2, srcbuf_stride,
 607                 width, 1, in_scale);
 608         srcbuf += srcbuf_stride;
 609
 610         s->read(src_data + FFMAX(y_out - 1, tff) * src_linesize,
 611                 srcbuf + 32,
 612                 src_linesize * 2, srcbuf_stride,
 613                 width, 1, in_scale);
 614         srcbuf += srcbuf_stride;
 615
 616         in_line  = src_data + FFMIN(y_out + 1, height - 1 - !tff) * src_linesize;
 617         out_line = dst_data + (y_out * dst_linesize);
 618
 619         s->read(in_line, srcbuf + 32, src_linesize * 2, srcbuf_stride,
 620                 width, slice_height - last_slice, in_scale);
 621
 622         y_out += (slice_height - last_slice) * 2;
 623
 624         s->read(src_data + FFMIN(y_out + 1, height - 1 - !tff) * src_linesize,
 625                 srcbuf + 32 + srcbuf_stride * (slice_height - last_slice),
 626                 src_linesize * 2, srcbuf_stride,
 627                 width, 1, in_scale);
 628
 629         s->read(src_data + FFMIN(y_out + 3, height - 1 - !tff) * src_linesize,
 630                 srcbuf + 32 + srcbuf_stride * (slice_height + 1 - last_slice),
 631                 src_linesize * 2, srcbuf_stride,
 632                 width, 1, in_scale);
 633
 634         s->read(src_data + FFMIN(y_out + 5, height - 1 - !tff) * src_linesize,
 635                 srcbuf + 32 + srcbuf_stride * (slice_height + 2 - last_slice),
 636                 src_linesize * 2, srcbuf_stride,
 637                 width, 1, in_scale);
 638
 639         for (int y = 0; y < slice_end - slice_start; y += 2) {
 640             if (s->pscrn > 0)
 641                 s->prescreen[s->pscrn > 1](ctx, srcbuf + (y / 2) * srcbuf_stride + 32,
 642                              srcbuf_stride, prescreen_buf, width,
 643                              &s->prescreener[s->pscrn - 1]);
 644
 645             predictor(ctx,
 646                       srcbuf + (y / 2) * srcbuf_stride + 32,
 647                       srcbuf_stride,
 648                       dstbuf + (y / 2) * dstbuf_stride,
 649                       prescreen_buf, width,
 650                       &s->coeffs[s->etype][s->nnsparam][s->nsize], s->qual == 2);
 651
 652             if (s->pscrn > 0)
 653                 interpolation(srcbuf + (y / 2) * srcbuf_stride + 32,
 654                               srcbuf_stride,
 655                               dstbuf + (y / 2) * dstbuf_stride,
 656                               prescreen_buf, width);
 657         }
 658
 659         s->write(dstbuf, out_line, dstbuf_stride, dst_linesize * 2,
 660                  width, slice_height, depth, out_scale);
 661     }
 662
 663     return 0;
 664 }
 665
 666 static int get_frame(AVFilterContext *ctx, int is_second)
 667 {
 668     NNEDIContext *s = ctx->priv;
 669     AVFilterLink *outlink = ctx->outputs[0];
 670     AVFrame *dst;
 671
 672     dst = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 673     if (!dst)
 674         return AVERROR(ENOMEM);
 675     av_frame_copy_props(dst, s->prev);
 676     dst->interlaced_frame = 0;
 677     dst->pts = s->pts;
 678
 679     ctx->internal->execute(ctx, filter_slice, dst, NULL, FFMIN(s->planeheight[1] / 2, s->nb_threads));
 680
 681     if (s->field == -2 || s->field > 1)
 682         s->field_n = !s->field_n;
 683
 684     return ff_filter_frame(outlink, dst);
 685 }
 686
 687 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 688 {
 689     AVFilterContext *ctx = inlink->dst;
 690     NNEDIContext *s = ctx->priv;
 691     int ret;
 692
 693     if (!s->prev) {
 694         s->prev = in;
 695         return 0;
 696     }
 697
 698     if ((s->deint && !in->interlaced_frame) || ctx->is_disabled) {
 699         s->prev->pts *= 2;
 700         ret = ff_filter_frame(ctx->outputs[0], s->prev);
 701         s->prev = in;
 702         return ret;
 703     }
 704
 705     s->pts = s->prev->pts * 2;
 706     ret = get_frame(ctx, 0);
 707     if (ret < 0 || (s->field > -2 && s->field < 2)) {
 708         av_frame_free(&s->prev);
 709         s->prev = in;
 710         return ret;
 711     }
 712
 713     s->pts = s->prev->pts + in->pts;
 714     ret = get_frame(ctx, 1);
 715     av_frame_free(&s->prev);
 716     s->prev = in;
 717     return ret;
 718 }
 719
 720 static int request_frame(AVFilterLink *link)
 721 {
 722     AVFilterContext *ctx = link->src;
 723     NNEDIContext *s = ctx->priv;
 724     int ret;
 725
 726     if (s->eof)
 727         return AVERROR_EOF;
 728
 729     ret  = ff_request_frame(ctx->inputs[0]);
 730
 731     if (ret == AVERROR_EOF && s->prev) {
 732         AVFrame *next = av_frame_clone(s->prev);
 733
 734         if (!next)
 735             return AVERROR(ENOMEM);
 736
 737         next->pts = s->prev->pts + av_rescale_q(1, av_inv_q(ctx->outputs[0]->frame_rate),
 738                                                 ctx->outputs[0]->time_base);
 739         s->eof = 1;
 740
 741         ret = filter_frame(ctx->inputs[0], next);
 742     } else if (ret < 0) {
 743         return ret;
 744     }
 745
 746     return ret;
 747 }
 748
 749 static void copy_weights(float *dst, int n, const float **data)
 750 {
 751     memcpy(dst, *data, n * sizeof(float));
 752     *data += n;
 753 }
 754
 755 static float *allocate(float **ptr, int size)
 756 {
 757     float *ret = *ptr;
 758
 759     *ptr += size;
 760
 761     return ret;
 762 }
 763
 764 static int allocate_model(PredictorCoefficients *coeffs, int xdim, int ydim, int nns)
 765 {
 766     int filter_size = nns * xdim * ydim;
 767     int bias_size = nns;
 768     float *data;
 769
 770     data = av_calloc(filter_size + bias_size, 4 * sizeof(float));
 771     if (!data)
 772         return AVERROR(ENOMEM);
 773
 774     coeffs->data = data;
 775     coeffs->xdim = xdim;
 776     coeffs->ydim = ydim;
 777     coeffs->nsize = xdim * ydim;
 778     coeffs->nns  = nns;
 779
 780     coeffs->softmax_q1 = allocate(&data, filter_size);
 781     coeffs->elliott_q1 = allocate(&data, filter_size);
 782     coeffs->softmax_bias_q1 = allocate(&data, bias_size);
 783     coeffs->elliott_bias_q1 = allocate(&data, bias_size);
 784
 785     coeffs->softmax_q2 = allocate(&data, filter_size);
 786     coeffs->elliott_q2 = allocate(&data, filter_size);
 787     coeffs->softmax_bias_q2 = allocate(&data, bias_size);
 788     coeffs->elliott_bias_q2 = allocate(&data, bias_size);
 789
 790     return 0;
 791 }
 792
 793 static int read_weights(AVFilterContext *ctx, const float *bdata)
 794 {
 795     NNEDIContext *s = ctx->priv;
 796     int ret;
 797
 798     copy_weights(&s->prescreener[0].kernel_l0[0][0], 4 * 48, &bdata);
 799     copy_weights(s->prescreener[0].bias_l0, 4, &bdata);
 800
 801     copy_weights(&s->prescreener[0].kernel_l1[0][0], 4 * 4, &bdata);
 802     copy_weights(s->prescreener[0].bias_l1, 4, &bdata);
 803
 804     copy_weights(&s->prescreener[0].kernel_l2[0][0], 4 * 8, &bdata);
 805     copy_weights(s->prescreener[0].bias_l2, 4, &bdata);
 806
 807     for (int i = 0; i < 3; i++) {
 808         PrescreenerCoefficients *data = &s->prescreener[i + 1];
 809         float kernel_l0_shuffled[4 * 64];
 810         float kernel_l1_shuffled[4 * 4];
 811
 812         copy_weights(kernel_l0_shuffled, 4 * 64, &bdata);
 813         copy_weights(data->bias_l0, 4, &bdata);
 814
 815         copy_weights(kernel_l1_shuffled, 4 * 4, &bdata);
 816         copy_weights(data->bias_l1, 4, &bdata);
 817
 818         for (int n = 0; n < 4; n++) {
 819             for (int k = 0; k < 64; k++)
 820                 data->kernel_l0[n][k] = kernel_l0_shuffled[(k / 8) * 32 + n * 8 + k % 8];
 821             for (int k = 0; k < 4; k++)
 822                 data->kernel_l1[n][k] = kernel_l1_shuffled[k * 4 + n];
 823         }
 824     }
 825
 826     for (int m = 0; m < 2; m++) {
 827         // Grouping by neuron count.
 828         for (int i = 0; i < 5; i++) {
 829             const int nns = NNEDI_NNS[i];
 830
 831             // Grouping by window size.
 832             for (int j = 0; j < 7; j++) {
 833                 PredictorCoefficients *model = &s->coeffs[m][i][j];
 834                 const int xdim = NNEDI_XDIM[j];
 835                 const int ydim = NNEDI_YDIM[j];
 836                 const int filter_size = xdim * ydim;
 837
 838                 ret = allocate_model(model, xdim, ydim, nns);
 839                 if (ret < 0)
 840                     return ret;
 841
 842                 // Quality 1 model. NNS[i] * (XDIM[j] * YDIM[j]) * 2 coefficients.
 843                 copy_weights(model->softmax_q1, nns * filter_size, &bdata);
 844                 copy_weights(model->elliott_q1, nns * filter_size, &bdata);
 845
 846                 // Quality 1 model bias. NNS[i] * 2 coefficients.
 847                 copy_weights(model->softmax_bias_q1, nns, &bdata);
 848                 copy_weights(model->elliott_bias_q1, nns, &bdata);
 849
 850                 // Quality 2 model. NNS[i] * (XDIM[j] * YDIM[j]) * 2 coefficients.
 851                 copy_weights(model->softmax_q2, nns * filter_size, &bdata);
 852                 copy_weights(model->elliott_q2, nns * filter_size, &bdata);
 853
 854                 // Quality 2 model bias. NNS[i] * 2 coefficients.
 855                 copy_weights(model->softmax_bias_q2, nns, &bdata);
 856                 copy_weights(model->elliott_bias_q2, nns, &bdata);
 857             }
 858         }
 859     }
 860
 861     return 0;
 862 }
 863
 864 static float mean(const float *input, int size)
 865 {
 866     float sum = 0.f;
 867
 868     for (int i = 0; i < size; i++)
 869         sum += input[i];
 870
 871     return sum / size;
 872 }
 873
 874 static void transform(float *input, int size, float mean, float half)
 875 {
 876     for (int i = 0; i < size; i++)
 877         input[i] = (input[i] - mean) / half;
 878 }
 879
 880 static void subtract_mean_old(PrescreenerCoefficients *coeffs, float half)
 881 {
 882     for (int n = 0; n < 4; n++) {
 883         float m = mean(coeffs->kernel_l0[n], 48);
 884
 885         transform(coeffs->kernel_l0[n], 48, m, half);
 886     }
 887 }
 888
 889 static void subtract_mean_new(PrescreenerCoefficients *coeffs, float half)
 890 {
 891     for (int n = 0; n < 4; n++) {
 892         float m = mean(coeffs->kernel_l0[n], 64);
 893
 894         transform(coeffs->kernel_l0[n], 64, m, half);
 895     }
 896 }
 897
 898 static void subtract_mean_predictor(PredictorCoefficients *model)
 899 {
 900     const int filter_size = model->nsize;
 901     const int nns = model->nns;
 902     const float scale = 1.f / nns;
 903
 904     double softmax_means[256]; // Average of individual softmax filters.
 905     double elliott_means[256]; // Average of individual elliott filters.
 906     double mean_filter[48 * 6] = { 0 }; // Pointwise average of all softmax filters.
 907     double mean_bias;
 908
 909     // Quality 1.
 910     for (int nn = 0; nn < nns; nn++) {
 911         softmax_means[nn] = mean(model->softmax_q1 + nn * filter_size, filter_size);
 912         elliott_means[nn] = mean(model->elliott_q1 + nn * filter_size, filter_size);
 913
 914         for (int k = 0; k < filter_size; k++)
 915             mean_filter[k] += model->softmax_q1[nn * filter_size + k] - softmax_means[nn];
 916     }
 917
 918     for (int k = 0; k < filter_size; k++)
 919         mean_filter[k] *= scale;
 920
 921     mean_bias = mean(model->softmax_bias_q1, nns);
 922
 923     for (int nn = 0; nn < nns; nn++) {
 924         for (int k = 0; k < filter_size; k++) {
 925             model->softmax_q1[nn * filter_size + k] -= softmax_means[nn] + mean_filter[k];
 926             model->elliott_q1[nn * filter_size + k] -= elliott_means[nn];
 927         }
 928         model->softmax_bias_q1[nn] -= mean_bias;
 929     }
 930
 931     // Quality 2.
 932     memset(mean_filter, 0, sizeof(mean_filter));
 933
 934     for (int nn = 0; nn < nns; nn++) {
 935         softmax_means[nn] = mean(model->softmax_q2 + nn * filter_size, filter_size);
 936         elliott_means[nn] = mean(model->elliott_q2 + nn * filter_size, filter_size);
 937
 938         for (int k = 0; k < filter_size; k++) {
 939             mean_filter[k] += model->softmax_q2[nn * filter_size + k] - softmax_means[nn];
 940         }
 941     }
 942
 943     for (int k = 0; k < filter_size; k++)
 944         mean_filter[k] *= scale;
 945
 946     mean_bias = mean(model->softmax_bias_q2, nns);
 947
 948     for (int nn = 0; nn < nns; nn++) {
 949         for (int k = 0; k < filter_size; k++) {
 950             model->softmax_q2[nn * filter_size + k] -= softmax_means[nn] + mean_filter[k];
 951             model->elliott_q2[nn * filter_size + k] -= elliott_means[nn];
 952         }
 953
 954         model->softmax_bias_q2[nn] -= mean_bias;
 955     }
 956 }
 957
 958 static av_cold int init(AVFilterContext *ctx)
 959 {
 960     NNEDIContext *s = ctx->priv;
 961     FILE *weights_file = NULL;
 962     int64_t weights_size;
 963     float *bdata;
 964     size_t bytes_read;
 965     int ret = 0;
 966
 967     weights_file = av_fopen_utf8(s->weights_file, "rb");
 968     if (!weights_file) {
 969         av_log(ctx, AV_LOG_ERROR, "No weights file provided, aborting!\n");
 970         return AVERROR(EINVAL);
 971     }
 972
 973     if (fseek(weights_file, 0, SEEK_END)) {
 974         av_log(ctx, AV_LOG_ERROR, "Couldn't seek to the end of weights file.\n");
 975         fclose(weights_file);
 976         return AVERROR(EINVAL);
 977     }
 978
 979     weights_size = ftell(weights_file);
 980
 981     if (weights_size == -1) {
 982         fclose(weights_file);
 983         av_log(ctx, AV_LOG_ERROR, "Couldn't get size of weights file.\n");
 984         return AVERROR(EINVAL);
 985     } else if (weights_size != NNEDI_WEIGHTS_SIZE) {
 986         fclose(weights_file);
 987         av_log(ctx, AV_LOG_ERROR, "Unexpected weights file size.\n");
 988         return AVERROR(EINVAL);
 989     }
 990
 991     if (fseek(weights_file, 0, SEEK_SET)) {
 992         fclose(weights_file);
 993         av_log(ctx, AV_LOG_ERROR, "Couldn't seek to the start of weights file.\n");
 994         return AVERROR(EINVAL);
 995     }
 996
 997     bdata = av_malloc(NNEDI_WEIGHTS_SIZE);
 998     if (!bdata) {
 999         fclose(weights_file);
1000         return AVERROR(ENOMEM);
1001     }
1002
1003     bytes_read = fread(bdata, 1, NNEDI_WEIGHTS_SIZE, weights_file);
1004     if (bytes_read != NNEDI_WEIGHTS_SIZE) {
1005         fclose(weights_file);
1006         ret = AVERROR_INVALIDDATA;
1007         av_log(ctx, AV_LOG_ERROR, "Couldn't read weights file.\n");
1008         goto fail;
1009     }
1010
1011     fclose(weights_file);
1012
1013     s->fdsp = avpriv_float_dsp_alloc(0);
1014     if (!s->fdsp) {
1015         ret = AVERROR(ENOMEM);
1016         goto fail;
1017     }
1018
1019     ret = read_weights(ctx, bdata);
1020     if (ret < 0)
1021         goto fail;
1022
1023 fail:
1024     av_free(bdata);
1025     return ret;
1026 }
1027
1028 static int config_input(AVFilterLink *inlink)
1029 {
1030     AVFilterContext *ctx = inlink->dst;
1031     NNEDIContext *s = ctx->priv;
1032     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
1033     int ret;
1034
1035     s->depth = desc->comp[0].depth;
1036     s->nb_threads = ff_filter_get_nb_threads(ctx);
1037     s->nb_planes = av_pix_fmt_count_planes(inlink->format);
1038     if ((ret = av_image_fill_linesizes(s->linesize, inlink->format, inlink->w)) < 0)
1039         return ret;
1040
1041     s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
1042     s->planewidth[0] = s->planewidth[3] = inlink->w;
1043     s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
1044     s->planeheight[0] = s->planeheight[3] = inlink->h;
1045
1046     s->half = ((1 << 8) - 1) / 2.f;
1047     s->out_scale = 1 << (s->depth - 8);
1048     s->in_scale = 1.f / s->out_scale;
1049
1050     switch (s->depth) {
1051     case 8:
1052         s->read  = read_bytes;
1053         s->write = write_bytes;
1054         break;
1055     default:
1056         s->read  = read_words;
1057         s->write = write_words;
1058         break;
1059     }
1060
1061     subtract_mean_old(&s->prescreener[0], s->half);
1062     subtract_mean_new(&s->prescreener[1], s->half);
1063     subtract_mean_new(&s->prescreener[2], s->half);
1064     subtract_mean_new(&s->prescreener[3], s->half);
1065
1066     s->prescreen[0] = process_old;
1067     s->prescreen[1] = process_new;
1068
1069     for (int i = 0; i < 2; i++) {
1070         for (int j = 0; j < 5; j++) {
1071             for (int k = 0; k < 7; k++)
1072                 subtract_mean_predictor(&s->coeffs[i][j][k]);
1073         }
1074     }
1075
1076     s->input_size = (s->planewidth[0] + 64) * (s->planeheight[0] + 6);
1077     s->input_buf = av_calloc(s->nb_threads, sizeof(*s->input_buf));
1078     if (!s->input_buf)
1079         return AVERROR(ENOMEM);
1080
1081     for (int i = 0; i < s->nb_threads; i++) {
1082         s->input_buf[i] = av_calloc(s->input_size, sizeof(**s->input_buf));
1083         if (!s->input_buf[i])
1084             return AVERROR(ENOMEM);
1085     }
1086
1087     s->output_buf = av_calloc(s->nb_threads, sizeof(*s->output_buf));
1088     if (!s->output_buf)
1089         return AVERROR(ENOMEM);
1090
1091     for (int i = 0; i < s->nb_threads; i++) {
1092         s->output_buf[i] = av_calloc(s->input_size, sizeof(**s->output_buf));
1093         if (!s->output_buf[i])
1094             return AVERROR(ENOMEM);
1095     }
1096
1097     s->prescreen_buf = av_calloc(s->nb_threads, sizeof(*s->prescreen_buf));
1098     if (!s->prescreen_buf)
1099         return AVERROR(ENOMEM);
1100
1101     for (int i = 0; i < s->nb_threads; i++) {
1102         s->prescreen_buf[i] = av_calloc(s->planewidth[0], sizeof(**s->prescreen_buf));
1103         if (!s->prescreen_buf[i])
1104             return AVERROR(ENOMEM);
1105     }
1106
1107     return 0;
1108 }
1109
1110 static av_cold void uninit(AVFilterContext *ctx)
1111 {
1112     NNEDIContext *s = ctx->priv;
1113
1114     for (int i = 0; i < s->nb_threads && s->prescreen_buf; i++)
1115         av_freep(&s->prescreen_buf[i]);
1116
1117     av_freep(&s->prescreen_buf);
1118
1119     for (int i = 0; i < s->nb_threads && s->input_buf; i++)
1120         av_freep(&s->input_buf[i]);
1121
1122     av_freep(&s->input_buf);
1123
1124     for (int i = 0; i < s->nb_threads && s->output_buf; i++)
1125         av_freep(&s->output_buf[i]);
1126
1127     av_freep(&s->output_buf);
1128     av_freep(&s->fdsp);
1129
1130     for (int i = 0; i < 2; i++) {
1131         for (int j = 0; j < 5; j++) {
1132             for (int k = 0; k < 7; k++) {
1133                 av_freep(&s->coeffs[i][j][k].data);
1134             }
1135         }
1136     }
1137
1138     av_frame_free(&s->prev);
1139 }
1140
1141 static const AVFilterPad inputs[] = {
1142     {
1143         .name          = "default",
1144         .type          = AVMEDIA_TYPE_VIDEO,
1145         .filter_frame  = filter_frame,
1146         .config_props  = config_input,
1147     },
1148     { NULL }
1149 };
1150
1151 static const AVFilterPad outputs[] = {
1152     {
1153         .name          = "default",
1154         .type          = AVMEDIA_TYPE_VIDEO,
1155         .config_props  = config_output,
1156         .request_frame = request_frame,
1157     },
1158     { NULL }
1159 };
1160
1161 const AVFilter ff_vf_nnedi = {
1162     .name          = "nnedi",
1163     .description   = NULL_IF_CONFIG_SMALL("Apply neural network edge directed interpolation intra-only deinterlacer."),
1164     .priv_size     = sizeof(NNEDIContext),
1165     .priv_class    = &nnedi_class,
1166     .init          = init,
1167     .uninit        = uninit,
1168     .query_formats = query_formats,
1169     .inputs        = inputs,
1170     .outputs       = outputs,
1171     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL | AVFILTER_FLAG_SLICE_THREADS,
1172     .process_command = ff_filter_process_command,
1173 };