TF_Status *status;
TF_Output input, output;
TF_Tensor *input_tensor;
+ TF_Tensor *output_tensor;
} TFModel;
static void free_buffer(void *data, size_t length)
return NULL;
}
- tf_model = av_malloc(sizeof(TFModel));
+ tf_model = av_mallocz(sizeof(TFModel));
if (!tf_model){
av_freep(&model);
return NULL;
}
- tf_model->session = NULL;
- tf_model->input_tensor = NULL;
if (load_tf_model(tf_model, model_filename) != DNN_SUCCESS){
if (load_native_model(tf_model, model_filename) != DNN_SUCCESS){
DNNReturnType ff_dnn_execute_model_tf(const DNNModel *model, DNNData *output)
{
TFModel *tf_model = (TFModel *)model->model;
- TF_Tensor *output_tensor;
- uint64_t count;
- uint64_t old_count = output->height * output->width * output->channels * sizeof(float);
+ if (tf_model->output_tensor)
+ TF_DeleteTensor(tf_model->output_tensor);
TF_SessionRun(tf_model->session, NULL,
&tf_model->input, &tf_model->input_tensor, 1,
- &tf_model->output, &output_tensor, 1,
+ &tf_model->output, &tf_model->output_tensor, 1,
NULL, 0, NULL, tf_model->status);
if (TF_GetCode(tf_model->status) != TF_OK){
return DNN_ERROR;
}
- output->height = TF_Dim(output_tensor, 1);
- output->width = TF_Dim(output_tensor, 2);
- output->channels = TF_Dim(output_tensor, 3);
- count = output->height * output->width * output->channels * sizeof(float);
- if (output->data) {
- if (count > old_count) {
- av_freep(&output->data);
- }
- }
- if (!output->data) {
- output->data = av_malloc(count);
- if (!output->data){
- return DNN_ERROR;
- }
- }
- memcpy(output->data, TF_TensorData(output_tensor), count);
- TF_DeleteTensor(output_tensor);
+ output->height = TF_Dim(tf_model->output_tensor, 1);
+ output->width = TF_Dim(tf_model->output_tensor, 2);
+ output->channels = TF_Dim(tf_model->output_tensor, 3);
+ output->data = TF_TensorData(tf_model->output_tensor);
return DNN_SUCCESS;
}
if (tf_model->input_tensor){
TF_DeleteTensor(tf_model->input_tensor);
}
+ if (tf_model->output_tensor){
+ TF_DeleteTensor(tf_model->output_tensor);
+ }
av_freep(&tf_model);
av_freep(model);
}