diff --git a/model_server/extensions/ilastik/models.py b/model_server/extensions/ilastik/models.py
index 7229b43eb601dc0703499392403e42e1bce3bf98..0cb6370cf217746b40f30bc32eb6ea329e4f74e9 100644
--- a/model_server/extensions/ilastik/models.py
+++ b/model_server/extensions/ilastik/models.py
@@ -36,6 +36,7 @@ class IlastikModel(Model):
             raise ParameterExpectedError('Ilastik model expects a project (*.ilp) file')
 
         self.shell = None
+        self.axes = None
         super().__init__(autoload, params)
 
     def load(self):
@@ -53,18 +54,20 @@ class IlastikModel(Model):
         shell = app.main(args, init_logging=False)
 
         # validate if inputs are embedded in project file
-        input_groups = shell.projectManager.currentProjectFile['Input Data']['infos']
-        lanes = input_groups.keys()
-        for ll in lanes:
-            input_types = input_groups[ll]
-            for tt in input_types:
-                ds_loc = input_groups[ll][tt].get('location', False)
-                if self.enforce_embedded and ds_loc and ds_loc[()] == b'FileSystem':
+        h5 = shell.projectManager.currentProjectFile
+        for lane in h5['Input Data/infos'].keys():
+            for role in h5[f'Input Data/infos/{lane}'].keys():
+                grp = h5[f'Input Data/infos/{lane}/{role}']
+                if self.enforce_embedded and ('location' in grp.keys()) and grp['location'][()] != b'ProjectInternal':
                     raise IlastikInputEmbedding('Cannot load ilastik project file where inputs are on filesystem')
+            assert True
         if not isinstance(shell.workflow, self.get_workflow()):
             raise ParameterExpectedError(
                 f'Ilastik project file {self.project_file} does not describe an instance of {shell.workflow.__class__}'
             )
+        self.axes = [
+            a['key'].upper() for a in json.loads(h5[f'Input Data/infos/lane0000/Raw Data/axistags'][()])['axes']
+        ]
         self.shell = shell
 
         return True