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