Signed-off-by: Pedro Arthur <bygrandao@gmail.com>pull/349/head
Sergey Lavrushkin
7 years ago
committed by
Pedro Arthur
parent
d24c9e55f6
commit
575b718990
10 changed files with 13249 additions and 379 deletions
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,354 @@ |
|||||||
|
/*
|
||||||
|
* Copyright (c) 2018 Sergey Lavrushkin |
||||||
|
* |
||||||
|
* This file is part of FFmpeg. |
||||||
|
* |
||||||
|
* FFmpeg is free software; you can redistribute it and/or |
||||||
|
* modify it under the terms of the GNU Lesser General Public |
||||||
|
* License as published by the Free Software Foundation; either |
||||||
|
* version 2.1 of the License, or (at your option) any later version. |
||||||
|
* |
||||||
|
* FFmpeg is distributed in the hope that it will be useful, |
||||||
|
* but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||||
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
||||||
|
* Lesser General Public License for more details. |
||||||
|
* |
||||||
|
* You should have received a copy of the GNU Lesser General Public |
||||||
|
* License along with FFmpeg; if not, write to the Free Software |
||||||
|
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
||||||
|
*/ |
||||||
|
|
||||||
|
/**
|
||||||
|
* @file |
||||||
|
* Filter implementing image super-resolution using deep convolutional networks. |
||||||
|
* https://arxiv.org/abs/1501.00092
|
||||||
|
* https://arxiv.org/abs/1609.05158
|
||||||
|
*/ |
||||||
|
|
||||||
|
#include "avfilter.h" |
||||||
|
#include "formats.h" |
||||||
|
#include "internal.h" |
||||||
|
#include "libavutil/opt.h" |
||||||
|
#include "libavformat/avio.h" |
||||||
|
#include "libswscale/swscale.h" |
||||||
|
#include "dnn_interface.h" |
||||||
|
|
||||||
|
typedef enum {SRCNN, ESPCN} SRModel; |
||||||
|
|
||||||
|
typedef struct SRContext { |
||||||
|
const AVClass *class; |
||||||
|
|
||||||
|
SRModel model_type; |
||||||
|
char* model_filename; |
||||||
|
DNNBackendType backend_type; |
||||||
|
DNNModule* dnn_module; |
||||||
|
DNNModel* model; |
||||||
|
DNNData input, output; |
||||||
|
int scale_factor; |
||||||
|
struct SwsContext* sws_context; |
||||||
|
int sws_slice_h; |
||||||
|
} SRContext; |
||||||
|
|
||||||
|
#define OFFSET(x) offsetof(SRContext, x) |
||||||
|
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM |
||||||
|
static const AVOption sr_options[] = { |
||||||
|
{ "model", "specifies what DNN model to use", OFFSET(model_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "model_type" }, |
||||||
|
{ "srcnn", "Super-Resolution Convolutional Neural Network model (scale factor should be specified for custom SRCNN model)", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "model_type" }, |
||||||
|
{ "espcn", "Efficient Sub-Pixel Convolutional Neural Network model", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "model_type" }, |
||||||
|
{ "dnn_backend", "DNN backend used for model execution", OFFSET(backend_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "backend" }, |
||||||
|
{ "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" }, |
||||||
|
#if (CONFIG_LIBTENSORFLOW == 1) |
||||||
|
{ "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" }, |
||||||
|
#endif |
||||||
|
{"scale_factor", "scale factor for SRCNN model", OFFSET(scale_factor), AV_OPT_TYPE_INT, { .i64 = 2 }, 2, 4, FLAGS}, |
||||||
|
{ "model_filename", "path to model file specifying network architecture and its parameters", OFFSET(model_filename), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS }, |
||||||
|
{ NULL } |
||||||
|
}; |
||||||
|
|
||||||
|
AVFILTER_DEFINE_CLASS(sr); |
||||||
|
|
||||||
|
static av_cold int init(AVFilterContext* context) |
||||||
|
{ |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
|
||||||
|
sr_context->dnn_module = ff_get_dnn_module(sr_context->backend_type); |
||||||
|
if (!sr_context->dnn_module){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not create DNN module for requested backend\n"); |
||||||
|
return AVERROR(ENOMEM); |
||||||
|
} |
||||||
|
if (!sr_context->model_filename){ |
||||||
|
av_log(context, AV_LOG_VERBOSE, "model file for network was not specified, using default network for x2 upsampling\n"); |
||||||
|
sr_context->scale_factor = 2; |
||||||
|
switch (sr_context->model_type){ |
||||||
|
case SRCNN: |
||||||
|
sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_SRCNN); |
||||||
|
break; |
||||||
|
case ESPCN: |
||||||
|
sr_context->model = (sr_context->dnn_module->load_default_model)(DNN_ESPCN); |
||||||
|
} |
||||||
|
} |
||||||
|
else{ |
||||||
|
sr_context->model = (sr_context->dnn_module->load_model)(sr_context->model_filename); |
||||||
|
} |
||||||
|
if (!sr_context->model){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not load DNN model\n"); |
||||||
|
return AVERROR(EIO); |
||||||
|
} |
||||||
|
|
||||||
|
return 0; |
||||||
|
} |
||||||
|
|
||||||
|
static int query_formats(AVFilterContext* context) |
||||||
|
{ |
||||||
|
const enum AVPixelFormat pixel_formats[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P, |
||||||
|
AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_GRAY8, |
||||||
|
AV_PIX_FMT_NONE}; |
||||||
|
AVFilterFormats* formats_list; |
||||||
|
|
||||||
|
formats_list = ff_make_format_list(pixel_formats); |
||||||
|
if (!formats_list){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not create formats list\n"); |
||||||
|
return AVERROR(ENOMEM); |
||||||
|
} |
||||||
|
return ff_set_common_formats(context, formats_list); |
||||||
|
} |
||||||
|
|
||||||
|
static int config_props(AVFilterLink* inlink) |
||||||
|
{ |
||||||
|
AVFilterContext* context = inlink->dst; |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
AVFilterLink* outlink = context->outputs[0]; |
||||||
|
DNNReturnType result; |
||||||
|
int sws_src_h, sws_src_w, sws_dst_h, sws_dst_w; |
||||||
|
|
||||||
|
switch (sr_context->model_type){ |
||||||
|
case SRCNN: |
||||||
|
sr_context->input.width = inlink->w * sr_context->scale_factor; |
||||||
|
sr_context->input.height = inlink->h * sr_context->scale_factor; |
||||||
|
break; |
||||||
|
case ESPCN: |
||||||
|
sr_context->input.width = inlink->w; |
||||||
|
sr_context->input.height = inlink->h; |
||||||
|
} |
||||||
|
sr_context->input.channels = 1; |
||||||
|
|
||||||
|
result = (sr_context->model->set_input_output)(sr_context->model->model, &sr_context->input, &sr_context->output); |
||||||
|
if (result != DNN_SUCCESS){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not set input and output for the model\n"); |
||||||
|
return AVERROR(EIO); |
||||||
|
} |
||||||
|
else{ |
||||||
|
outlink->h = sr_context->output.height; |
||||||
|
outlink->w = sr_context->output.width; |
||||||
|
switch (sr_context->model_type){ |
||||||
|
case SRCNN: |
||||||
|
sr_context->sws_context = sws_getContext(inlink->w, inlink->h, inlink->format, |
||||||
|
outlink->w, outlink->h, outlink->format, SWS_BICUBIC, NULL, NULL, NULL); |
||||||
|
if (!sr_context->sws_context){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not create SwsContext\n"); |
||||||
|
return AVERROR(ENOMEM); |
||||||
|
} |
||||||
|
sr_context->sws_slice_h = inlink->h; |
||||||
|
break; |
||||||
|
case ESPCN: |
||||||
|
if (inlink->format == AV_PIX_FMT_GRAY8){ |
||||||
|
sr_context->sws_context = NULL; |
||||||
|
} |
||||||
|
else{ |
||||||
|
sws_src_h = sr_context->input.height; |
||||||
|
sws_src_w = sr_context->input.width; |
||||||
|
sws_dst_h = sr_context->output.height; |
||||||
|
sws_dst_w = sr_context->output.width; |
||||||
|
|
||||||
|
switch (inlink->format){ |
||||||
|
case AV_PIX_FMT_YUV420P: |
||||||
|
sws_src_h = (sws_src_h >> 1) + (sws_src_h % 2 != 0 ? 1 : 0); |
||||||
|
sws_src_w = (sws_src_w >> 1) + (sws_src_w % 2 != 0 ? 1 : 0); |
||||||
|
sws_dst_h = (sws_dst_h >> 1) + (sws_dst_h % 2 != 0 ? 1 : 0); |
||||||
|
sws_dst_w = (sws_dst_w >> 1) + (sws_dst_w % 2 != 0 ? 1 : 0); |
||||||
|
break; |
||||||
|
case AV_PIX_FMT_YUV422P: |
||||||
|
sws_src_w = (sws_src_w >> 1) + (sws_src_w % 2 != 0 ? 1 : 0); |
||||||
|
sws_dst_w = (sws_dst_w >> 1) + (sws_dst_w % 2 != 0 ? 1 : 0); |
||||||
|
break; |
||||||
|
case AV_PIX_FMT_YUV444P: |
||||||
|
break; |
||||||
|
case AV_PIX_FMT_YUV410P: |
||||||
|
sws_src_h = (sws_src_h >> 2) + (sws_src_h % 4 != 0 ? 1 : 0); |
||||||
|
sws_src_w = (sws_src_w >> 2) + (sws_src_w % 4 != 0 ? 1 : 0); |
||||||
|
sws_dst_h = (sws_dst_h >> 2) + (sws_dst_h % 4 != 0 ? 1 : 0); |
||||||
|
sws_dst_w = (sws_dst_w >> 2) + (sws_dst_w % 4 != 0 ? 1 : 0); |
||||||
|
break; |
||||||
|
case AV_PIX_FMT_YUV411P: |
||||||
|
sws_src_w = (sws_src_w >> 2) + (sws_src_w % 4 != 0 ? 1 : 0); |
||||||
|
sws_dst_w = (sws_dst_w >> 2) + (sws_dst_w % 4 != 0 ? 1 : 0); |
||||||
|
break; |
||||||
|
default: |
||||||
|
av_log(context, AV_LOG_ERROR, "could not create SwsContext for input pixel format"); |
||||||
|
return AVERROR(EIO); |
||||||
|
} |
||||||
|
sr_context->sws_context = sws_getContext(sws_src_w, sws_src_h, AV_PIX_FMT_GRAY8, |
||||||
|
sws_dst_w, sws_dst_h, AV_PIX_FMT_GRAY8, SWS_BICUBIC, NULL, NULL, NULL); |
||||||
|
if (!sr_context->sws_context){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not create SwsContext\n"); |
||||||
|
return AVERROR(ENOMEM); |
||||||
|
} |
||||||
|
sr_context->sws_slice_h = sws_src_h; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
return 0; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
typedef struct ThreadData{ |
||||||
|
uint8_t* data; |
||||||
|
int data_linesize, height, width; |
||||||
|
} ThreadData; |
||||||
|
|
||||||
|
static int uint8_to_float(AVFilterContext* context, void* arg, int jobnr, int nb_jobs) |
||||||
|
{ |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
const ThreadData* td = arg; |
||||||
|
const int slice_start = (td->height * jobnr ) / nb_jobs; |
||||||
|
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs; |
||||||
|
const uint8_t* src = td->data + slice_start * td->data_linesize; |
||||||
|
float* dst = sr_context->input.data + slice_start * td->width; |
||||||
|
int y, x; |
||||||
|
|
||||||
|
for (y = slice_start; y < slice_end; ++y){ |
||||||
|
for (x = 0; x < td->width; ++x){ |
||||||
|
dst[x] = (float)src[x] / 255.0f; |
||||||
|
} |
||||||
|
src += td->data_linesize; |
||||||
|
dst += td->width; |
||||||
|
} |
||||||
|
|
||||||
|
return 0; |
||||||
|
} |
||||||
|
|
||||||
|
static int float_to_uint8(AVFilterContext* context, void* arg, int jobnr, int nb_jobs) |
||||||
|
{ |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
const ThreadData* td = arg; |
||||||
|
const int slice_start = (td->height * jobnr ) / nb_jobs; |
||||||
|
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs; |
||||||
|
const float* src = sr_context->output.data + slice_start * td->width; |
||||||
|
uint8_t* dst = td->data + slice_start * td->data_linesize; |
||||||
|
int y, x; |
||||||
|
|
||||||
|
for (y = slice_start; y < slice_end; ++y){ |
||||||
|
for (x = 0; x < td->width; ++x){ |
||||||
|
dst[x] = (uint8_t)(255.0f * FFMIN(src[x], 1.0f)); |
||||||
|
} |
||||||
|
src += td->width; |
||||||
|
dst += td->data_linesize; |
||||||
|
} |
||||||
|
|
||||||
|
return 0; |
||||||
|
} |
||||||
|
|
||||||
|
static int filter_frame(AVFilterLink* inlink, AVFrame* in) |
||||||
|
{ |
||||||
|
AVFilterContext* context = inlink->dst; |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
AVFilterLink* outlink = context->outputs[0]; |
||||||
|
AVFrame* out = ff_get_video_buffer(outlink, outlink->w, outlink->h); |
||||||
|
ThreadData td; |
||||||
|
int nb_threads; |
||||||
|
DNNReturnType dnn_result; |
||||||
|
|
||||||
|
if (!out){ |
||||||
|
av_log(context, AV_LOG_ERROR, "could not allocate memory for output frame\n"); |
||||||
|
av_frame_free(&in); |
||||||
|
return AVERROR(ENOMEM); |
||||||
|
} |
||||||
|
av_frame_copy_props(out, in); |
||||||
|
out->height = sr_context->output.height; |
||||||
|
out->width = sr_context->output.width; |
||||||
|
switch (sr_context->model_type){ |
||||||
|
case SRCNN: |
||||||
|
sws_scale(sr_context->sws_context, in->data, in->linesize, |
||||||
|
0, sr_context->sws_slice_h, out->data, out->linesize); |
||||||
|
td.data = out->data[0]; |
||||||
|
td.data_linesize = out->linesize[0]; |
||||||
|
td.height = out->height; |
||||||
|
td.width = out->width; |
||||||
|
break; |
||||||
|
case ESPCN: |
||||||
|
if (sr_context->sws_context){ |
||||||
|
sws_scale(sr_context->sws_context, in->data + 1, in->linesize + 1, |
||||||
|
0, sr_context->sws_slice_h, out->data + 1, out->linesize + 1); |
||||||
|
sws_scale(sr_context->sws_context, in->data + 2, in->linesize + 2, |
||||||
|
0, sr_context->sws_slice_h, out->data + 2, out->linesize + 2); |
||||||
|
} |
||||||
|
td.data = in->data[0]; |
||||||
|
td.data_linesize = in->linesize[0]; |
||||||
|
td.height = in->height; |
||||||
|
td.width = in->width; |
||||||
|
} |
||||||
|
|
||||||
|
nb_threads = ff_filter_get_nb_threads(context); |
||||||
|
context->internal->execute(context, uint8_to_float, &td, NULL, FFMIN(td.height, nb_threads)); |
||||||
|
av_frame_free(&in); |
||||||
|
|
||||||
|
dnn_result = (sr_context->dnn_module->execute_model)(sr_context->model); |
||||||
|
if (dnn_result != DNN_SUCCESS){ |
||||||
|
av_log(context, AV_LOG_ERROR, "failed to execute loaded model\n"); |
||||||
|
return AVERROR(EIO); |
||||||
|
} |
||||||
|
|
||||||
|
td.data = out->data[0]; |
||||||
|
td.data_linesize = out->linesize[0]; |
||||||
|
td.height = out->height; |
||||||
|
td.width = out->width; |
||||||
|
context->internal->execute(context, float_to_uint8, &td, NULL, FFMIN(td.height, nb_threads)); |
||||||
|
|
||||||
|
return ff_filter_frame(outlink, out); |
||||||
|
} |
||||||
|
|
||||||
|
static av_cold void uninit(AVFilterContext* context) |
||||||
|
{ |
||||||
|
SRContext* sr_context = context->priv; |
||||||
|
|
||||||
|
if (sr_context->dnn_module){ |
||||||
|
(sr_context->dnn_module->free_model)(&sr_context->model); |
||||||
|
av_freep(&sr_context->dnn_module); |
||||||
|
} |
||||||
|
|
||||||
|
if (sr_context->sws_context){ |
||||||
|
sws_freeContext(sr_context->sws_context); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
static const AVFilterPad sr_inputs[] = { |
||||||
|
{ |
||||||
|
.name = "default", |
||||||
|
.type = AVMEDIA_TYPE_VIDEO, |
||||||
|
.config_props = config_props, |
||||||
|
.filter_frame = filter_frame, |
||||||
|
}, |
||||||
|
{ NULL } |
||||||
|
}; |
||||||
|
|
||||||
|
static const AVFilterPad sr_outputs[] = { |
||||||
|
{ |
||||||
|
.name = "default", |
||||||
|
.type = AVMEDIA_TYPE_VIDEO, |
||||||
|
}, |
||||||
|
{ NULL } |
||||||
|
}; |
||||||
|
|
||||||
|
AVFilter ff_vf_sr = { |
||||||
|
.name = "sr", |
||||||
|
.description = NULL_IF_CONFIG_SMALL("Apply DNN-based image super resolution to the input."), |
||||||
|
.priv_size = sizeof(SRContext), |
||||||
|
.init = init, |
||||||
|
.uninit = uninit, |
||||||
|
.query_formats = query_formats, |
||||||
|
.inputs = sr_inputs, |
||||||
|
.outputs = sr_outputs, |
||||||
|
.priv_class = &sr_class, |
||||||
|
.flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS, |
||||||
|
}; |
||||||
|
|
||||||
@ -1,250 +0,0 @@ |
|||||||
/*
|
|
||||||
* Copyright (c) 2018 Sergey Lavrushkin |
|
||||||
* |
|
||||||
* This file is part of FFmpeg. |
|
||||||
* |
|
||||||
* FFmpeg is free software; you can redistribute it and/or |
|
||||||
* modify it under the terms of the GNU Lesser General Public |
|
||||||
* License as published by the Free Software Foundation; either |
|
||||||
* version 2.1 of the License, or (at your option) any later version. |
|
||||||
* |
|
||||||
* FFmpeg is distributed in the hope that it will be useful, |
|
||||||
* but WITHOUT ANY WARRANTY; without even the implied warranty of |
|
||||||
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
|
||||||
* Lesser General Public License for more details. |
|
||||||
* |
|
||||||
* You should have received a copy of the GNU Lesser General Public |
|
||||||
* License along with FFmpeg; if not, write to the Free Software |
|
||||||
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
|
||||||
*/ |
|
||||||
|
|
||||||
/**
|
|
||||||
* @file |
|
||||||
* Filter implementing image super-resolution using deep convolutional networks. |
|
||||||
* https://arxiv.org/abs/1501.00092
|
|
||||||
*/ |
|
||||||
|
|
||||||
#include "avfilter.h" |
|
||||||
#include "formats.h" |
|
||||||
#include "internal.h" |
|
||||||
#include "libavutil/opt.h" |
|
||||||
#include "libavformat/avio.h" |
|
||||||
#include "dnn_interface.h" |
|
||||||
|
|
||||||
typedef struct SRCNNContext { |
|
||||||
const AVClass *class; |
|
||||||
|
|
||||||
char* model_filename; |
|
||||||
float* input_output_buf; |
|
||||||
DNNBackendType backend_type; |
|
||||||
DNNModule* dnn_module; |
|
||||||
DNNModel* model; |
|
||||||
DNNData input_output; |
|
||||||
} SRCNNContext; |
|
||||||
|
|
||||||
#define OFFSET(x) offsetof(SRCNNContext, x) |
|
||||||
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM |
|
||||||
static const AVOption srcnn_options[] = { |
|
||||||
{ "dnn_backend", "DNN backend used for model execution", OFFSET(backend_type), AV_OPT_TYPE_FLAGS, { .i64 = 0 }, 0, 1, FLAGS, "backend" }, |
|
||||||
{ "native", "native backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 0 }, 0, 0, FLAGS, "backend" }, |
|
||||||
#if (CONFIG_LIBTENSORFLOW == 1) |
|
||||||
{ "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0, FLAGS, "backend" }, |
|
||||||
#endif |
|
||||||
{ "model_filename", "path to model file specifying network architecture and its parameters", OFFSET(model_filename), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS }, |
|
||||||
{ NULL } |
|
||||||
}; |
|
||||||
|
|
||||||
AVFILTER_DEFINE_CLASS(srcnn); |
|
||||||
|
|
||||||
static av_cold int init(AVFilterContext* context) |
|
||||||
{ |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
|
|
||||||
srcnn_context->dnn_module = ff_get_dnn_module(srcnn_context->backend_type); |
|
||||||
if (!srcnn_context->dnn_module){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not create DNN module for requested backend\n"); |
|
||||||
return AVERROR(ENOMEM); |
|
||||||
} |
|
||||||
if (!srcnn_context->model_filename){ |
|
||||||
av_log(context, AV_LOG_VERBOSE, "model file for network was not specified, using default network for x2 upsampling\n"); |
|
||||||
srcnn_context->model = (srcnn_context->dnn_module->load_default_model)(DNN_SRCNN); |
|
||||||
} |
|
||||||
else{ |
|
||||||
srcnn_context->model = (srcnn_context->dnn_module->load_model)(srcnn_context->model_filename); |
|
||||||
} |
|
||||||
if (!srcnn_context->model){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not load DNN model\n"); |
|
||||||
return AVERROR(EIO); |
|
||||||
} |
|
||||||
|
|
||||||
return 0; |
|
||||||
} |
|
||||||
|
|
||||||
static int query_formats(AVFilterContext* context) |
|
||||||
{ |
|
||||||
const enum AVPixelFormat pixel_formats[] = {AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P, |
|
||||||
AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_GRAY8, |
|
||||||
AV_PIX_FMT_NONE}; |
|
||||||
AVFilterFormats* formats_list; |
|
||||||
|
|
||||||
formats_list = ff_make_format_list(pixel_formats); |
|
||||||
if (!formats_list){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not create formats list\n"); |
|
||||||
return AVERROR(ENOMEM); |
|
||||||
} |
|
||||||
return ff_set_common_formats(context, formats_list); |
|
||||||
} |
|
||||||
|
|
||||||
static int config_props(AVFilterLink* inlink) |
|
||||||
{ |
|
||||||
AVFilterContext* context = inlink->dst; |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
DNNReturnType result; |
|
||||||
|
|
||||||
srcnn_context->input_output_buf = av_malloc(inlink->h * inlink->w * sizeof(float)); |
|
||||||
if (!srcnn_context->input_output_buf){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not allocate memory for input/output buffer\n"); |
|
||||||
return AVERROR(ENOMEM); |
|
||||||
} |
|
||||||
|
|
||||||
srcnn_context->input_output.data = srcnn_context->input_output_buf; |
|
||||||
srcnn_context->input_output.width = inlink->w; |
|
||||||
srcnn_context->input_output.height = inlink->h; |
|
||||||
srcnn_context->input_output.channels = 1; |
|
||||||
|
|
||||||
result = (srcnn_context->model->set_input_output)(srcnn_context->model->model, &srcnn_context->input_output, &srcnn_context->input_output); |
|
||||||
if (result != DNN_SUCCESS){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not set input and output for the model\n"); |
|
||||||
return AVERROR(EIO); |
|
||||||
} |
|
||||||
else{ |
|
||||||
return 0; |
|
||||||
} |
|
||||||
} |
|
||||||
|
|
||||||
typedef struct ThreadData{ |
|
||||||
uint8_t* out; |
|
||||||
int out_linesize, height, width; |
|
||||||
} ThreadData; |
|
||||||
|
|
||||||
static int uint8_to_float(AVFilterContext* context, void* arg, int jobnr, int nb_jobs) |
|
||||||
{ |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
const ThreadData* td = arg; |
|
||||||
const int slice_start = (td->height * jobnr ) / nb_jobs; |
|
||||||
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs; |
|
||||||
const uint8_t* src = td->out + slice_start * td->out_linesize; |
|
||||||
float* dst = srcnn_context->input_output_buf + slice_start * td->width; |
|
||||||
int y, x; |
|
||||||
|
|
||||||
for (y = slice_start; y < slice_end; ++y){ |
|
||||||
for (x = 0; x < td->width; ++x){ |
|
||||||
dst[x] = (float)src[x] / 255.0f; |
|
||||||
} |
|
||||||
src += td->out_linesize; |
|
||||||
dst += td->width; |
|
||||||
} |
|
||||||
|
|
||||||
return 0; |
|
||||||
} |
|
||||||
|
|
||||||
static int float_to_uint8(AVFilterContext* context, void* arg, int jobnr, int nb_jobs) |
|
||||||
{ |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
const ThreadData* td = arg; |
|
||||||
const int slice_start = (td->height * jobnr ) / nb_jobs; |
|
||||||
const int slice_end = (td->height * (jobnr + 1)) / nb_jobs; |
|
||||||
const float* src = srcnn_context->input_output_buf + slice_start * td->width; |
|
||||||
uint8_t* dst = td->out + slice_start * td->out_linesize; |
|
||||||
int y, x; |
|
||||||
|
|
||||||
for (y = slice_start; y < slice_end; ++y){ |
|
||||||
for (x = 0; x < td->width; ++x){ |
|
||||||
dst[x] = (uint8_t)(255.0f * FFMIN(src[x], 1.0f)); |
|
||||||
} |
|
||||||
src += td->width; |
|
||||||
dst += td->out_linesize; |
|
||||||
} |
|
||||||
|
|
||||||
return 0; |
|
||||||
} |
|
||||||
|
|
||||||
static int filter_frame(AVFilterLink* inlink, AVFrame* in) |
|
||||||
{ |
|
||||||
AVFilterContext* context = inlink->dst; |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
AVFilterLink* outlink = context->outputs[0]; |
|
||||||
AVFrame* out = ff_get_video_buffer(outlink, outlink->w, outlink->h); |
|
||||||
ThreadData td; |
|
||||||
int nb_threads; |
|
||||||
DNNReturnType dnn_result; |
|
||||||
|
|
||||||
if (!out){ |
|
||||||
av_log(context, AV_LOG_ERROR, "could not allocate memory for output frame\n"); |
|
||||||
av_frame_free(&in); |
|
||||||
return AVERROR(ENOMEM); |
|
||||||
} |
|
||||||
av_frame_copy_props(out, in); |
|
||||||
av_frame_copy(out, in); |
|
||||||
av_frame_free(&in); |
|
||||||
td.out = out->data[0]; |
|
||||||
td.out_linesize = out->linesize[0]; |
|
||||||
td.height = out->height; |
|
||||||
td.width = out->width; |
|
||||||
|
|
||||||
nb_threads = ff_filter_get_nb_threads(context); |
|
||||||
context->internal->execute(context, uint8_to_float, &td, NULL, FFMIN(td.height, nb_threads)); |
|
||||||
|
|
||||||
dnn_result = (srcnn_context->dnn_module->execute_model)(srcnn_context->model); |
|
||||||
if (dnn_result != DNN_SUCCESS){ |
|
||||||
av_log(context, AV_LOG_ERROR, "failed to execute loaded model\n"); |
|
||||||
return AVERROR(EIO); |
|
||||||
} |
|
||||||
|
|
||||||
context->internal->execute(context, float_to_uint8, &td, NULL, FFMIN(td.height, nb_threads)); |
|
||||||
|
|
||||||
return ff_filter_frame(outlink, out); |
|
||||||
} |
|
||||||
|
|
||||||
static av_cold void uninit(AVFilterContext* context) |
|
||||||
{ |
|
||||||
SRCNNContext* srcnn_context = context->priv; |
|
||||||
|
|
||||||
if (srcnn_context->dnn_module){ |
|
||||||
(srcnn_context->dnn_module->free_model)(&srcnn_context->model); |
|
||||||
av_freep(&srcnn_context->dnn_module); |
|
||||||
} |
|
||||||
av_freep(&srcnn_context->input_output_buf); |
|
||||||
} |
|
||||||
|
|
||||||
static const AVFilterPad srcnn_inputs[] = { |
|
||||||
{ |
|
||||||
.name = "default", |
|
||||||
.type = AVMEDIA_TYPE_VIDEO, |
|
||||||
.config_props = config_props, |
|
||||||
.filter_frame = filter_frame, |
|
||||||
}, |
|
||||||
{ NULL } |
|
||||||
}; |
|
||||||
|
|
||||||
static const AVFilterPad srcnn_outputs[] = { |
|
||||||
{ |
|
||||||
.name = "default", |
|
||||||
.type = AVMEDIA_TYPE_VIDEO, |
|
||||||
}, |
|
||||||
{ NULL } |
|
||||||
}; |
|
||||||
|
|
||||||
AVFilter ff_vf_srcnn = { |
|
||||||
.name = "srcnn", |
|
||||||
.description = NULL_IF_CONFIG_SMALL("Apply super resolution convolutional neural network to the input. Use bicubic upsamping with corresponding scaling factor before."), |
|
||||||
.priv_size = sizeof(SRCNNContext), |
|
||||||
.init = init, |
|
||||||
.uninit = uninit, |
|
||||||
.query_formats = query_formats, |
|
||||||
.inputs = srcnn_inputs, |
|
||||||
.outputs = srcnn_outputs, |
|
||||||
.priv_class = &srcnn_class, |
|
||||||
.flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS, |
|
||||||
}; |
|
||||||
|
|
||||||
Loading…
Reference in new issue