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

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include "dnn_filter_common.h"
  20
  21 int ff_dnn_init(DnnContext *ctx, DNNFunctionType func_type, AVFilterContext *filter_ctx)
  22 {
  23     if (!ctx->model_filename) {
  24         av_log(filter_ctx, AV_LOG_ERROR, "model file for network is not specified\n");
  25         return AVERROR(EINVAL);
  26     }
  27     if (!ctx->model_inputname) {
  28         av_log(filter_ctx, AV_LOG_ERROR, "input name of the model network is not specified\n");
  29         return AVERROR(EINVAL);
  30     }
  31     if (!ctx->model_outputname) {
  32         av_log(filter_ctx, AV_LOG_ERROR, "output name of the model network is not specified\n");
  33         return AVERROR(EINVAL);
  34     }
  35
  36     ctx->dnn_module = ff_get_dnn_module(ctx->backend_type);
  37     if (!ctx->dnn_module) {
  38         av_log(filter_ctx, AV_LOG_ERROR, "could not create DNN module for requested backend\n");
  39         return AVERROR(ENOMEM);
  40     }
  41     if (!ctx->dnn_module->load_model) {
  42         av_log(filter_ctx, AV_LOG_ERROR, "load_model for network is not specified\n");
  43         return AVERROR(EINVAL);
  44     }
  45
  46     ctx->model = (ctx->dnn_module->load_model)(ctx->model_filename, func_type, ctx->backend_options, filter_ctx);
  47     if (!ctx->model) {
  48         av_log(filter_ctx, AV_LOG_ERROR, "could not load DNN model\n");
  49         return AVERROR(EINVAL);
  50     }
  51
  52     if (!ctx->dnn_module->execute_model_async && ctx->async) {
  53         ctx->async = 0;
  54         av_log(filter_ctx, AV_LOG_WARNING, "this backend does not support async execution, roll back to sync.\n");
  55     }
  56
  57 #if !HAVE_PTHREAD_CANCEL
  58     if (ctx->async) {
  59         ctx->async = 0;
  60         av_log(filter_ctx, AV_LOG_WARNING, "pthread is not supported, roll back to sync.\n");
  61     }
  62 #endif
  63
  64     return 0;
  65 }
  66
  67 int ff_dnn_set_frame_proc(DnnContext *ctx, FramePrePostProc pre_proc, FramePrePostProc post_proc)
  68 {
  69     ctx->model->frame_pre_proc = pre_proc;
  70     ctx->model->frame_post_proc = post_proc;
  71     return 0;
  72 }
  73
  74 int ff_dnn_set_detect_post_proc(DnnContext *ctx, DetectPostProc post_proc)
  75 {
  76     ctx->model->detect_post_proc = post_proc;
  77     return 0;
  78 }
  79
  80 DNNReturnType ff_dnn_get_input(DnnContext *ctx, DNNData *input)
  81 {
  82     return ctx->model->get_input(ctx->model->model, input, ctx->model_inputname);
  83 }
  84
  85 DNNReturnType ff_dnn_get_output(DnnContext *ctx, int input_width, int input_height, int *output_width, int *output_height)
  86 {
  87     return ctx->model->get_output(ctx->model->model, ctx->model_inputname, input_width, input_height,
  88                                     ctx->model_outputname, output_width, output_height);
  89 }
  90
  91 DNNReturnType ff_dnn_execute_model(DnnContext *ctx, AVFrame *in_frame, AVFrame *out_frame)
  92 {
  93     return (ctx->dnn_module->execute_model)(ctx->model, ctx->model_inputname, in_frame,
  94                                             (const char **)&ctx->model_outputname, 1, out_frame);
  95 }
  96
  97 DNNReturnType ff_dnn_execute_model_async(DnnContext *ctx, AVFrame *in_frame, AVFrame *out_frame)
  98 {
  99     return (ctx->dnn_module->execute_model_async)(ctx->model, ctx->model_inputname, in_frame,
 100                                                   (const char **)&ctx->model_outputname, 1, out_frame);
 101 }
 102
 103 DNNAsyncStatusType ff_dnn_get_async_result(DnnContext *ctx, AVFrame **in_frame, AVFrame **out_frame)
 104 {
 105     return (ctx->dnn_module->get_async_result)(ctx->model, in_frame, out_frame);
 106 }
 107
 108 DNNReturnType ff_dnn_flush(DnnContext *ctx)
 109 {
 110     return (ctx->dnn_module->flush)(ctx->model);
 111 }
 112
 113 void ff_dnn_uninit(DnnContext *ctx)
 114 {
 115     if (ctx->dnn_module) {
 116         (ctx->dnn_module->free_model)(&ctx->model);
 117         av_freep(&ctx->dnn_module);
 118     }
 119 }