Machine learning demo improvements (#178)

NTNU-IHB · Apr 18, 2023 · 3a00eb4 · 3a00eb4
1 parent b3bbafd
commit 3a00eb4
Showing 1 changed file with 22 additions and 13 deletions.
diff --git a/examples/mldemo.py b/examples/mldemo.py
@@ -1,11 +1,12 @@
 from pythonfmu import Fmi2Causality, Fmi2Slave, Real
 
+import pathlib  # pathlib is included in the Python Standard Library
+
 try:
-    import os
     import tensorflow as tf
     import numpy as np
-except ImportError:
-    os, np, tf = None, None, None
+except ImportError: # Trick to be able to generate the FMU without TensorFlow and NumPy installed
+    np, tf = None, None
 
 
 class MLDemo(Fmi2Slave):
@@ -14,19 +15,27 @@ class MLDemo(Fmi2Slave):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        # Import Tensorflow model directory created by model.save([some_directory_name])
-        try:
-            # Fetch model from directory included in the FMU. Included files are not available until after
-            # the FMU has been built. Included files and directories will have the same root as this file by default
-            self.model = tf.keras.models.load_model(os.path.join(os.path.dirname(__file__), "stored-model"))
-        except OSError:
-            # Fetch model from local directory when building the FMU. Remember to include this directory
-            # in the FMU itself as a project file. Path assumes this command is run from repo root folder
-            self.model = tf.keras.models.load_model("./examples/tensorflow-model-export/stored-model")
+        # Import Tensorflow model directory created by model.save(["stored-model"])
+        parent_path = pathlib.Path(__file__).parent
+        # Check if building or initializing
+        if parent_path.name == "resources":
+            # __init__ called from within FMU (probably)
+            # Path relative to "resources" directory root within the FMU
+            model_dir_path = parent_path / "stored-model"
+            if model_dir_path.exists():
+                try:
+                    # Fetch saved model from directory included in the FMU
+                    self.model = tf.keras.models.load_model(model_dir_path)
+                except AttributeError:
+                    print("Unable to load model from directory. Has TensorFlow been included in the environment?")
+                except OSError:
+                    print("Unable to load model from directory. Has the correct directory been specified?")
+            else:
+                print("No model directory found. Has the directory been included in the FMU when building?")
 
         self.sin_input = 0.
         self.sin_output_tf = 0.
-        self.sin_output_ref = np.sin([self.sin_input])[0]
+        self.sin_output_ref = 0.
 
         self.register_variable(Real("sin_input", causality=Fmi2Causality.input))
         self.register_variable(Real("sin_output_tf", causality=Fmi2Causality.output))