n_images bug fix

supervisely-ecosystem · Jul 18, 2023 · 2ed1399 · 2ed1399
1 parent e7a0e31
commit 2ed1399
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train/src/main.py b/train/src/main.py
@@ -1393,6 +1393,10 @@ def auto_train(request: Request):
             dataset_ids.append(dataset_info.id)
     else:
         dataset_ids = state["dataset_ids"]
+        n_images = 0
+        for dataset_id in dataset_ids:
+            dataset_info = api.dataset.get_info_by_id(dataset_id)
+            n_images += dataset_info.images_count
     # download dataset
     if os.path.exists(g.project_dir):
         sly.fs.clean_dir(g.project_dir)