fix KeyError

supervisely-ecosystem · Nov 13, 2023 · ab3fd27 · ab3fd27
1 parent 9885833
commit ab3fd27
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/supervisely/train/src/sly_train.py b/supervisely/train/src/sly_train.py
@@ -132,7 +132,7 @@ def train(api: sly.Api, task_id, context, state, app_logger):
         my_app.show_modal_window(f"{msg} {repr(e)}", level="error", log_message=False)
         sly.logger.error(repr(e), exc_info=True, extra={ 'exc_str': str(e)})
         try:
-            api.task.set_output_error(task_id, repr(e), msg)
+            api.task.set_output_error(task_id, repr(e), "Find more info in the app logs.")
             api.app.set_field(task_id, "state.started", False)
         except:
             pass

diff --git a/supervisely/train/src/sly_train_globals.py b/supervisely/train/src/sly_train_globals.py
@@ -40,6 +40,7 @@
 remote_artifacts_dir = None
 project_info = api.project.get_info_by_id(project_id)
 project_meta = sly.ProjectMeta.from_json(api.project.get_meta(project_id))
+project_stats = api.project.get_stats(project_id)
 
 with open(os.path.join(root_source_dir, "data/hyp.scratch.yaml"), 'r') as file:
     scratch_str = file.read()  # yaml.safe_load(

diff --git a/supervisely/train/src/ui/classes.py b/supervisely/train/src/ui/classes.py
@@ -1,8 +1,7 @@
 import supervisely as sly
 
 
-def init(api: sly.Api, data, state, project_id, project_meta: sly.ProjectMeta):
-    stats = api.project.get_stats(project_id)
+def init(api: sly.Api, data, state, project_meta: sly.ProjectMeta, stats):
     class_images = {}
     for item in stats["images"]["objectClasses"]:
         class_images[item["objectClass"]["name"]] = item["total"]
@@ -12,8 +11,13 @@ def init(api: sly.Api, data, state, project_id, project_meta: sly.ProjectMeta):
 
     classes_json = project_meta.obj_classes.to_json()
     for obj_class in classes_json:
-        obj_class["imagesCount"] = class_images[obj_class["title"]]
-        obj_class["objectsCount"] = class_objects[obj_class["title"]]
+        cls_title = obj_class["title"]
+        if cls_title not in class_images or cls_title not in class_objects:
+            raise Exception(
+                f"Class '{cls_title}' not found in the project. Please check your data."
+            )
+        obj_class["imagesCount"] = class_images[cls_title]
+        obj_class["objectsCount"] = class_objects[cls_title]
 
     unlabeled_count = 0
     for ds_counter in stats["images"]["datasets"]:

diff --git a/supervisely/train/src/ui/ui.py b/supervisely/train/src/ui/ui.py
@@ -10,7 +10,7 @@
 
 def init(data, state):
     input_project.init(data)
-    training_classes.init(g.api, data, state, g.project_id, g.project_meta)
+    training_classes.init(g.api, data, state, g.project_meta, g.project_stats)
     train_val_split.init(g.project_info, g.project_meta, data, state)
     model_architectures.init(data, state)
     hyperparameters.init(state)