1 files changed, 22 insertions, 2 deletions
diff --git a/source/use_case/img_class/src/MainLoop.cc b/source/use_case/img_class/src/MainLoop.cc
index d9fb925..de3779f 100644
--- a/source/use_case/img_class/src/MainLoop.cc
+++ b/source/use_case/img_class/src/MainLoop.cc
@@ -21,7 +21,16 @@
 #include "MobileNetModel.hpp"       /* Model class for running inference. */
 #include "UseCaseHandler.hpp"       /* Handlers for different user options. */
 #include "UseCaseCommonUtils.hpp"   /* Utils functions. */
-#include "log_macros.h"
+#include "BufAttributes.hpp"        /* Buffer attributes to be applied */
+
+namespace arm {
+    namespace app {
+        static uint8_t  tensorArena[ACTIVATION_BUF_SZ] ACTIVATION_BUF_ATTRIBUTE;
+    } /* namespace app */
+} /* namespace arm */
+
+extern uint8_t* GetModelPointer();
+extern size_t GetModelLen();
 
 using ImgClassClassifier = arm::app::Classifier;
 
@@ -30,11 +39,22 @@ void main_loop()
     arm::app::MobileNetModel model;  /* Model wrapper object. */
 
     /* Load the model. */
-    if (!model.Init()) {
+    if (!model.Init(arm::app::tensorArena,
+                    sizeof(arm::app::tensorArena),
+                    GetModelPointer(),
+                    GetModelLen())) {
         printf_err("Failed to initialise model\n");
         return;
     }
 
+#if !defined(ARM_NPU)
+    /* If it is not a NPU build check if the model contains a NPU operator */
+    if (model.ContainsEthosUOperator()) {
+        printf_err("No driver support for Ethos-U operator found in the model.\n");
+        return;
+    }
+#endif /* ARM_NPU */
+
     /* Instantiate application context. */
     arm::app::ApplicationContext caseContext;