avfilter/sr: process and output message when load_model is NULL
[ffmpeg.git] / libavfilter / dnn_interface.c
index e981829..86fc283 100644 (file)
 
 #include "dnn_interface.h"
 #include "dnn_backend_native.h"
+#include "dnn_backend_tf.h"
 #include "libavutil/mem.h"
 
-DNNModuleff_get_dnn_module(DNNBackendType backend_type)
+DNNModule *ff_get_dnn_module(DNNBackendType backend_type)
 {
-    DNNModulednn_module;
+    DNNModule *dnn_module;
 
     dnn_module = av_malloc(sizeof(DNNModule));
     if(!dnn_module){
@@ -39,9 +40,23 @@ DNNModule* ff_get_dnn_module(DNNBackendType backend_type)
     switch(backend_type){
     case DNN_NATIVE:
         dnn_module->load_model = &ff_dnn_load_model_native;
-        dnn_module->load_default_model = &ff_dnn_load_default_model_native;
         dnn_module->execute_model = &ff_dnn_execute_model_native;
         dnn_module->free_model = &ff_dnn_free_model_native;
+        break;
+    case DNN_TF:
+    #if (CONFIG_LIBTENSORFLOW == 1)
+        dnn_module->load_model = &ff_dnn_load_model_tf;
+        dnn_module->execute_model = &ff_dnn_execute_model_tf;
+        dnn_module->free_model = &ff_dnn_free_model_tf;
+    #else
+        av_freep(&dnn_module);
+        return NULL;
+    #endif
+        break;
+    default:
+        av_log(NULL, AV_LOG_ERROR, "Module backend_type is not native or tensorflow\n");
+        av_freep(&dnn_module);
+        return NULL;
     }
 
     return dnn_module;