libavfilter/dnn_interface.h

   1 /*
   2  * Copyright (c) 2018 Sergey Lavrushkin
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Lesser General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2.1 of the License, or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 /**
  22  * @file
  23  * DNN inference engine interface.
  24  */
  25
  26 #ifndef AVFILTER_DNN_INTERFACE_H
  27 #define AVFILTER_DNN_INTERFACE_H
  28
  29 #include <stdint.h>
  30 #include "libavutil/frame.h"
  31 #include "avfilter.h"
  32
  33 #define DNN_GENERIC_ERROR FFERRTAG('D','N','N','!')
  34
  35 typedef enum {DNN_NATIVE, DNN_TF, DNN_OV} DNNBackendType;
  36
  37 typedef enum {DNN_FLOAT = 1, DNN_UINT8 = 4} DNNDataType;
  38
  39 typedef enum {
  40     DCO_NONE,
  41     DCO_BGR,
  42     DCO_RGB,
  43 } DNNColorOrder;
  44
  45 typedef enum {
  46     DAST_FAIL,              // something wrong
  47     DAST_EMPTY_QUEUE,       // no more inference result to get
  48     DAST_NOT_READY,         // all queued inferences are not finished
  49     DAST_SUCCESS            // got a result frame successfully
  50 } DNNAsyncStatusType;
  51
  52 typedef enum {
  53     DFT_NONE,
  54     DFT_PROCESS_FRAME,      // process the whole frame
  55     DFT_ANALYTICS_DETECT,   // detect from the whole frame
  56     DFT_ANALYTICS_CLASSIFY, // classify for each bounding box
  57 }DNNFunctionType;
  58
  59 typedef struct DNNData{
  60     void *data;
  61     int width, height, channels;
  62     // dt and order together decide the color format
  63     DNNDataType dt;
  64     DNNColorOrder order;
  65 } DNNData;
  66
  67 typedef struct DNNExecBaseParams {
  68     const char *input_name;
  69     const char **output_names;
  70     uint32_t nb_output;
  71     AVFrame *in_frame;
  72     AVFrame *out_frame;
  73 } DNNExecBaseParams;
  74
  75 typedef struct DNNExecClassificationParams {
  76     DNNExecBaseParams base;
  77     const char *target;
  78 } DNNExecClassificationParams;
  79
  80 typedef int (*FramePrePostProc)(AVFrame *frame, DNNData *model, AVFilterContext *filter_ctx);
  81 typedef int (*DetectPostProc)(AVFrame *frame, DNNData *output, uint32_t nb, AVFilterContext *filter_ctx);
  82 typedef int (*ClassifyPostProc)(AVFrame *frame, DNNData *output, uint32_t bbox_index, AVFilterContext *filter_ctx);
  83
  84 typedef struct DNNModel{
  85     // Stores model that can be different for different backends.
  86     void *model;
  87     // Stores options when the model is executed by the backend
  88     const char *options;
  89     // Stores FilterContext used for the interaction between AVFrame and DNNData
  90     AVFilterContext *filter_ctx;
  91     // Stores function type of the model
  92     DNNFunctionType func_type;
  93     // Gets model input information
  94     // Just reuse struct DNNData here, actually the DNNData.data field is not needed.
  95     int (*get_input)(void *model, DNNData *input, const char *input_name);
  96     // Gets model output width/height with given input w/h
  97     int (*get_output)(void *model, const char *input_name, int input_width, int input_height,
  98                                 const char *output_name, int *output_width, int *output_height);
  99     // set the pre process to transfer data from AVFrame to DNNData
 100     // the default implementation within DNN is used if it is not provided by the filter
 101     FramePrePostProc frame_pre_proc;
 102     // set the post process to transfer data from DNNData to AVFrame
 103     // the default implementation within DNN is used if it is not provided by the filter
 104     FramePrePostProc frame_post_proc;
 105     // set the post process to interpret detect result from DNNData
 106     DetectPostProc detect_post_proc;
 107     // set the post process to interpret classify result from DNNData
 108     ClassifyPostProc classify_post_proc;
 109 } DNNModel;
 110
 111 // Stores pointers to functions for loading, executing, freeing DNN models for one of the backends.
 112 typedef struct DNNModule{
 113     // Loads model and parameters from given file. Returns NULL if it is not possible.
 114     DNNModel *(*load_model)(const char *model_filename, DNNFunctionType func_type, const char *options, AVFilterContext *filter_ctx);
 115     // Executes model with specified input and output. Returns the error code otherwise.
 116     int (*execute_model)(const DNNModel *model, DNNExecBaseParams *exec_params);
 117     // Retrieve inference result.
 118     DNNAsyncStatusType (*get_result)(const DNNModel *model, AVFrame **in, AVFrame **out);
 119     // Flush all the pending tasks.
 120     int (*flush)(const DNNModel *model);
 121     // Frees memory allocated for model.
 122     void (*free_model)(DNNModel **model);
 123 } DNNModule;
 124
 125 // Initializes DNNModule depending on chosen backend.
 126 DNNModule *ff_get_dnn_module(DNNBackendType backend_type);
 127
 128 #endif