]> git.sesse.net Git - ffmpeg/blobdiff - libavfilter/dnn/dnn_backend_tf.c
avfilter/dnn/dnn_backend_tf: simplify the code with ff_hex_to_data
[ffmpeg] / libavfilter / dnn / dnn_backend_tf.c
index 7923e1db694bcd08de4a0adf1d096b8b58bc76b8..076dd3d6a90ae94948e5d64e6c05e372666e29d5 100644 (file)
 #include "dnn_backend_native_layer_conv2d.h"
 #include "dnn_backend_native_layer_depth2space.h"
 #include "libavformat/avio.h"
+#include "libavformat/internal.h"
 #include "libavutil/avassert.h"
+#include "../internal.h"
 #include "dnn_backend_native_layer_pad.h"
 #include "dnn_backend_native_layer_maximum.h"
 #include "dnn_io_proc.h"
 
 #include <tensorflow/c/c_api.h>
 
+typedef struct TFOptions{
+    char *sess_config;
+} TFOptions;
+
 typedef struct TFContext {
     const AVClass *class;
+    TFOptions options;
 } TFContext;
 
 typedef struct TFModel{
@@ -47,14 +54,15 @@ typedef struct TFModel{
     TF_Status *status;
 } TFModel;
 
-static const AVClass dnn_tensorflow_class = {
-    .class_name = "dnn_tensorflow",
-    .item_name  = av_default_item_name,
-    .option     = NULL,
-    .version    = LIBAVUTIL_VERSION_INT,
-    .category   = AV_CLASS_CATEGORY_FILTER,
+#define OFFSET(x) offsetof(TFContext, x)
+#define FLAGS AV_OPT_FLAG_FILTERING_PARAM
+static const AVOption dnn_tensorflow_options[] = {
+    { "sess_config", "config for SessionOptions", OFFSET(options.sess_config), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS },
+    { NULL }
 };
 
+AVFILTER_DEFINE_CLASS(dnn_tensorflow);
+
 static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_name, AVFrame *in_frame,
                                       const char **output_names, uint32_t nb_output, AVFrame *out_frame,
                                       int do_ioproc);
@@ -90,7 +98,7 @@ static TF_Buffer *read_graph(const char *model_filename)
     }
 
     graph_buf = TF_NewBuffer();
-    graph_buf->data = (void *)graph_data;
+    graph_buf->data = graph_data;
     graph_buf->length = size;
     graph_buf->data_deallocator = free_buffer;
 
@@ -121,7 +129,7 @@ static TF_Tensor *allocate_input_tensor(const DNNData *input)
 
 static DNNReturnType get_input_tf(void *model, DNNData *input, const char *input_name)
 {
-    TFModel *tf_model = (TFModel *)model;
+    TFModel *tf_model = model;
     TFContext *ctx = &tf_model->ctx;
     TF_Status *status;
     int64_t dims[4];
@@ -158,7 +166,7 @@ static DNNReturnType get_output_tf(void *model, const char *input_name, int inpu
                                    const char *output_name, int *output_width, int *output_height)
 {
     DNNReturnType ret;
-    TFModel *tf_model = (TFModel *)model;
+    TFModel *tf_model = model;
     TFContext *ctx = &tf_model->ctx;
     AVFrame *in_frame = av_frame_alloc();
     AVFrame *out_frame = NULL;
@@ -194,10 +202,37 @@ static DNNReturnType load_tf_model(TFModel *tf_model, const char *model_filename
     TF_ImportGraphDefOptions *graph_opts;
     TF_SessionOptions *sess_opts;
     const TF_Operation *init_op;
+    uint8_t *sess_config = NULL;
+    int sess_config_length = 0;
+
+    // prepare the sess config data
+    if (tf_model->ctx.options.sess_config != NULL) {
+        const char *config;
+        /*
+        tf_model->ctx.options.sess_config is hex to present the serialized proto
+        required by TF_SetConfig below, so we need to first generate the serialized
+        proto in a python script, tools/python/tf_sess_config.py is a script example
+        to generate the configs of sess_config.
+        */
+        if (strncmp(tf_model->ctx.options.sess_config, "0x", 2) != 0) {
+            av_log(ctx, AV_LOG_ERROR, "sess_config should start with '0x'\n");
+            return DNN_ERROR;
+        }
+        config = tf_model->ctx.options.sess_config + 2;
+        sess_config_length = ff_hex_to_data(NULL, config);
+
+        sess_config = av_mallocz(sess_config_length + AV_INPUT_BUFFER_PADDING_SIZE);
+        if (!sess_config) {
+            av_log(ctx, AV_LOG_ERROR, "failed to allocate memory\n");
+            return DNN_ERROR;
+        }
+        ff_hex_to_data(sess_config, config);
+    }
 
     graph_def = read_graph(model_filename);
     if (!graph_def){
         av_log(ctx, AV_LOG_ERROR, "Failed to read model \"%s\" graph\n", model_filename);
+        av_freep(&sess_config);
         return DNN_ERROR;
     }
     tf_model->graph = TF_NewGraph();
@@ -210,15 +245,32 @@ static DNNReturnType load_tf_model(TFModel *tf_model, const char *model_filename
         TF_DeleteGraph(tf_model->graph);
         TF_DeleteStatus(tf_model->status);
         av_log(ctx, AV_LOG_ERROR, "Failed to import serialized graph to model graph\n");
+        av_freep(&sess_config);
         return DNN_ERROR;
     }
 
     init_op = TF_GraphOperationByName(tf_model->graph, "init");
     sess_opts = TF_NewSessionOptions();
+
+    if (sess_config) {
+        TF_SetConfig(sess_opts, sess_config, sess_config_length,tf_model->status);
+        av_freep(&sess_config);
+        if (TF_GetCode(tf_model->status) != TF_OK) {
+            TF_DeleteGraph(tf_model->graph);
+            TF_DeleteStatus(tf_model->status);
+            TF_DeleteSessionOptions(sess_opts);
+            av_log(ctx, AV_LOG_ERROR, "Failed to set config for sess options with %s\n",
+                                      tf_model->ctx.options.sess_config);
+            return DNN_ERROR;
+        }
+    }
+
     tf_model->session = TF_NewSession(tf_model->graph, sess_opts, tf_model->status);
     TF_DeleteSessionOptions(sess_opts);
     if (TF_GetCode(tf_model->status) != TF_OK)
     {
+        TF_DeleteGraph(tf_model->graph);
+        TF_DeleteStatus(tf_model->status);
         av_log(ctx, AV_LOG_ERROR, "Failed to create new session with model graph\n");
         return DNN_ERROR;
     }
@@ -231,6 +283,9 @@ static DNNReturnType load_tf_model(TFModel *tf_model, const char *model_filename
                       &init_op, 1, NULL, tf_model->status);
         if (TF_GetCode(tf_model->status) != TF_OK)
         {
+            TF_DeleteSession(tf_model->session, tf_model->status);
+            TF_DeleteGraph(tf_model->graph);
+            TF_DeleteStatus(tf_model->status);
             av_log(ctx, AV_LOG_ERROR, "Failed to run session when initializing\n");
             return DNN_ERROR;
         }
@@ -249,7 +304,7 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_OperationDescription *op_desc;
     TF_Output input;
     int64_t strides[] = {1, 1, 1, 1};
-    TF_Tensor *tensor;
+    TF_Tensor *kernel_tensor = NULL, *biases_tensor = NULL;
     int64_t dims[4];
     int dims_len;
     char name_buffer[NAME_BUFFER_SIZE];
@@ -266,17 +321,15 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     dims[2] = params->kernel_size;
     dims[3] = params->input_num;
     dims_len = 4;
-    tensor = TF_AllocateTensor(TF_FLOAT, dims, dims_len, size * sizeof(float));
-    memcpy(TF_TensorData(tensor), params->kernel, size * sizeof(float));
-    TF_SetAttrTensor(op_desc, "value", tensor, tf_model->status);
+    kernel_tensor = TF_AllocateTensor(TF_FLOAT, dims, dims_len, size * sizeof(float));
+    memcpy(TF_TensorData(kernel_tensor), params->kernel, size * sizeof(float));
+    TF_SetAttrTensor(op_desc, "value", kernel_tensor, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to set value for kernel of conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
     op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add kernel to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     snprintf(name_buffer, NAME_BUFFER_SIZE, "transpose%d", layer);
@@ -289,8 +342,7 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_SetAttrType(op_desc, "Tperm", TF_INT32);
     op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add transpose to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     snprintf(name_buffer, NAME_BUFFER_SIZE, "conv2d%d", layer);
@@ -304,8 +356,7 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_SetAttrString(op_desc, "padding", "VALID", 5);
     *cur_op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add conv2d to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     snprintf(name_buffer, NAME_BUFFER_SIZE, "conv_biases%d", layer);
@@ -313,17 +364,15 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_SetAttrType(op_desc, "dtype", TF_FLOAT);
     dims[0] = params->output_num;
     dims_len = 1;
-    tensor = TF_AllocateTensor(TF_FLOAT, dims, dims_len, params->output_num * sizeof(float));
-    memcpy(TF_TensorData(tensor), params->biases, params->output_num * sizeof(float));
-    TF_SetAttrTensor(op_desc, "value", tensor, tf_model->status);
+    biases_tensor = TF_AllocateTensor(TF_FLOAT, dims, dims_len, params->output_num * sizeof(float));
+    memcpy(TF_TensorData(biases_tensor), params->biases, params->output_num * sizeof(float));
+    TF_SetAttrTensor(op_desc, "value", biases_tensor, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to set value for conv_biases of conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
     op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add conv_biases to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     snprintf(name_buffer, NAME_BUFFER_SIZE, "bias_add%d", layer);
@@ -335,8 +384,7 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_SetAttrType(op_desc, "T", TF_FLOAT);
     *cur_op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add bias_add to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     snprintf(name_buffer, NAME_BUFFER_SIZE, "activation%d", layer);
@@ -351,7 +399,7 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
         op_desc = TF_NewOperation(tf_model->graph, "Sigmoid", name_buffer);
         break;
     default:
-        av_log(ctx, AV_LOG_ERROR, "Unsupported convolutional activation function\n");
+        avpriv_report_missing_feature(ctx, "convolutional activation function %d", params->activation);
         return DNN_ERROR;
     }
     input.oper = *cur_op;
@@ -359,11 +407,15 @@ static DNNReturnType add_conv_layer(TFModel *tf_model, TF_Operation *transpose_o
     TF_SetAttrType(op_desc, "T", TF_FLOAT);
     *cur_op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
-        av_log(ctx, AV_LOG_ERROR, "Failed to add activation function to conv layer %d\n", layer);
-        return DNN_ERROR;
+        goto err;
     }
 
     return DNN_SUCCESS;
+err:
+    TF_DeleteTensor(kernel_tensor);
+    TF_DeleteTensor(biases_tensor);
+    av_log(ctx, AV_LOG_ERROR, "Failed to add conv layer %d\n", layer);
+    return DNN_ERROR;
 }
 
 static DNNReturnType add_depth_to_space_layer(TFModel *tf_model, TF_Operation **cur_op,
@@ -418,11 +470,13 @@ static DNNReturnType add_pad_layer(TFModel *tf_model, TF_Operation **cur_op,
     pads[7] = params->paddings[3][1];
     TF_SetAttrTensor(op_desc, "value", tensor, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to set value for pad of layer %d\n", layer);
         return DNN_ERROR;
     }
     op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to add pad to layer %d\n", layer);
         return DNN_ERROR;
     }
@@ -438,6 +492,7 @@ static DNNReturnType add_pad_layer(TFModel *tf_model, TF_Operation **cur_op,
     TF_SetAttrString(op_desc, "mode", "SYMMETRIC", 9);
     *cur_op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to add mirror_pad to layer %d\n", layer);
         return DNN_ERROR;
     }
@@ -465,11 +520,13 @@ static DNNReturnType add_maximum_layer(TFModel *tf_model, TF_Operation **cur_op,
     *y = params->val.y;
     TF_SetAttrTensor(op_desc, "value", tensor, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to set value for maximum/y of layer %d", layer);
         return DNN_ERROR;
     }
     op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to add maximum/y to layer %d\n", layer);
         return DNN_ERROR;
     }
@@ -484,6 +541,7 @@ static DNNReturnType add_maximum_layer(TFModel *tf_model, TF_Operation **cur_op,
     TF_SetAttrType(op_desc, "T", TF_FLOAT);
     *cur_op = TF_FinishOperation(op_desc, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK){
+        TF_DeleteTensor(tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to add maximum to layer %d\n", layer);
         return DNN_ERROR;
     }
@@ -498,7 +556,7 @@ static DNNReturnType load_native_model(TFModel *tf_model, const char *model_file
     TF_OperationDescription *op_desc;
     TF_Operation *op;
     TF_Operation *transpose_op;
-    TF_Tensor *tensor;
+    TF_Tensor *tensor = NULL;
     TF_Output input;
     int32_t *transpose_perm;
     int64_t transpose_perm_shape[] = {4};
@@ -507,18 +565,19 @@ static DNNReturnType load_native_model(TFModel *tf_model, const char *model_file
     DNNModel *model = NULL;
     NativeModel *native_model;
 
-    model = ff_dnn_load_model_native(model_filename, NULL, NULL);
+    model = ff_dnn_load_model_native(model_filename, DFT_PROCESS_FRAME, NULL, NULL);
     if (!model){
         av_log(ctx, AV_LOG_ERROR, "Failed to load native model\n");
         return DNN_ERROR;
     }
 
-    native_model = (NativeModel *)model->model;
+    native_model = model->model;
     tf_model->graph = TF_NewGraph();
     tf_model->status = TF_NewStatus();
 
 #define CLEANUP_ON_ERROR(tf_model) \
     { \
+        TF_DeleteTensor(tensor); \
         TF_DeleteGraph(tf_model->graph); \
         TF_DeleteStatus(tf_model->status); \
         av_log(ctx, AV_LOG_ERROR, "Failed to set value or add operator to layer\n"); \
@@ -546,6 +605,9 @@ static DNNReturnType load_native_model(TFModel *tf_model, const char *model_file
         CLEANUP_ON_ERROR(tf_model);
     }
     transpose_op = TF_FinishOperation(op_desc, tf_model->status);
+    if (TF_GetCode(tf_model->status) != TF_OK){
+        CLEANUP_ON_ERROR(tf_model);
+    }
 
     for (layer = 0; layer < native_model->layers_num; ++layer){
         switch (native_model->layers[layer].type){
@@ -591,7 +653,7 @@ static DNNReturnType load_native_model(TFModel *tf_model, const char *model_file
     return DNN_SUCCESS;
 }
 
-DNNModel *ff_dnn_load_model_tf(const char *model_filename, const char *options, void *userdata)
+DNNModel *ff_dnn_load_model_tf(const char *model_filename, DNNFunctionType func_type, const char *options, AVFilterContext *filter_ctx)
 {
     DNNModel *model = NULL;
     TFModel *tf_model = NULL;
@@ -609,6 +671,15 @@ DNNModel *ff_dnn_load_model_tf(const char *model_filename, const char *options,
     tf_model->ctx.class = &dnn_tensorflow_class;
     tf_model->model = model;
 
+    //parse options
+    av_opt_set_defaults(&tf_model->ctx);
+    if (av_opt_set_from_string(&tf_model->ctx, options, NULL, "=", "&") < 0) {
+        av_log(&tf_model->ctx, AV_LOG_ERROR, "Failed to parse options \"%s\"\n", options);
+        av_freep(&tf_model);
+        av_freep(&model);
+        return NULL;
+    }
+
     if (load_tf_model(tf_model, model_filename) != DNN_SUCCESS){
         if (load_native_model(tf_model, model_filename) != DNN_SUCCESS){
             av_freep(&tf_model);
@@ -618,11 +689,12 @@ DNNModel *ff_dnn_load_model_tf(const char *model_filename, const char *options,
         }
     }
 
-    model->model = (void *)tf_model;
+    model->model = tf_model;
     model->get_input = &get_input_tf;
     model->get_output = &get_output_tf;
     model->options = options;
-    model->userdata = userdata;
+    model->filter_ctx = filter_ctx;
+    model->func_type = func_type;
 
     return model;
 }
@@ -632,7 +704,7 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
                                       int do_ioproc)
 {
     TF_Output *tf_outputs;
-    TFModel *tf_model = (TFModel *)model->model;
+    TFModel *tf_model = model->model;
     TFContext *ctx = &tf_model->ctx;
     DNNData input, output;
     TF_Tensor **output_tensors;
@@ -658,28 +730,31 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
     input.data = (float *)TF_TensorData(input_tensor);
 
     if (do_ioproc) {
-        if (tf_model->model->pre_proc != NULL) {
-            tf_model->model->pre_proc(in_frame, &input, tf_model->model->userdata);
+        if (tf_model->model->frame_pre_proc != NULL) {
+            tf_model->model->frame_pre_proc(in_frame, &input, tf_model->model->filter_ctx);
         } else {
-            proc_from_frame_to_dnn(in_frame, &input, ctx);
+            ff_proc_from_frame_to_dnn(in_frame, &input, tf_model->model->func_type, ctx);
         }
     }
 
     if (nb_output != 1) {
         // currently, the filter does not need multiple outputs,
         // so we just pending the support until we really need it.
-        av_log(ctx, AV_LOG_ERROR, "do not support multiple outputs\n");
+        TF_DeleteTensor(input_tensor);
+        avpriv_report_missing_feature(ctx, "multiple outputs");
         return DNN_ERROR;
     }
 
     tf_outputs = av_malloc_array(nb_output, sizeof(*tf_outputs));
     if (tf_outputs == NULL) {
+        TF_DeleteTensor(input_tensor);
         av_log(ctx, AV_LOG_ERROR, "Failed to allocate memory for *tf_outputs\n"); \
         return DNN_ERROR;
     }
 
     output_tensors = av_mallocz_array(nb_output, sizeof(*output_tensors));
     if (!output_tensors) {
+        TF_DeleteTensor(input_tensor);
         av_freep(&tf_outputs);
         av_log(ctx, AV_LOG_ERROR, "Failed to allocate memory for output tensor\n"); \
         return DNN_ERROR;
@@ -688,6 +763,7 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
     for (int i = 0; i < nb_output; ++i) {
         tf_outputs[i].oper = TF_GraphOperationByName(tf_model->graph, output_names[i]);
         if (!tf_outputs[i].oper) {
+            TF_DeleteTensor(input_tensor);
             av_freep(&tf_outputs);
             av_freep(&output_tensors);
             av_log(ctx, AV_LOG_ERROR, "Could not find output \"%s\" in model\n", output_names[i]); \
@@ -701,6 +777,7 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
                   tf_outputs, output_tensors, nb_output,
                   NULL, 0, NULL, tf_model->status);
     if (TF_GetCode(tf_model->status) != TF_OK) {
+        TF_DeleteTensor(input_tensor);
         av_freep(&tf_outputs);
         av_freep(&output_tensors);
         av_log(ctx, AV_LOG_ERROR, "Failed to run session when executing model\n");
@@ -715,10 +792,10 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
         output.dt = TF_TensorType(output_tensors[i]);
 
         if (do_ioproc) {
-            if (tf_model->model->post_proc != NULL) {
-                tf_model->model->post_proc(out_frame, &output, tf_model->model->userdata);
+            if (tf_model->model->frame_post_proc != NULL) {
+                tf_model->model->frame_post_proc(out_frame, &output, tf_model->model->filter_ctx);
             } else {
-                proc_from_dnn_to_frame(out_frame, &output, ctx);
+                ff_proc_from_dnn_to_frame(out_frame, &output, ctx);
             }
         } else {
             out_frame->width = output.width;
@@ -740,7 +817,7 @@ static DNNReturnType execute_model_tf(const DNNModel *model, const char *input_n
 DNNReturnType ff_dnn_execute_model_tf(const DNNModel *model, const char *input_name, AVFrame *in_frame,
                                       const char **output_names, uint32_t nb_output, AVFrame *out_frame)
 {
-    TFModel *tf_model = (TFModel *)model->model;
+    TFModel *tf_model = model->model;
     TFContext *ctx = &tf_model->ctx;
 
     if (!in_frame) {
@@ -761,7 +838,7 @@ void ff_dnn_free_model_tf(DNNModel **model)
     TFModel *tf_model;
 
     if (*model){
-        tf_model = (TFModel *)(*model)->model;
+        tf_model = (*model)->model;
         if (tf_model->graph){
             TF_DeleteGraph(tf_model->graph);
         }