update train

supervisely-ecosystem · Jun 7, 2024 · 69ccc4d · 69ccc4d
1 parent 0094d89
commit 69ccc4d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/supervisely/train/src/sly_train_globals.py b/supervisely/train/src/sly_train_globals.py
@@ -3,7 +3,7 @@
 import sys
 import yaml
 import supervisely as sly
-from supervisely.nn.checkpoints.yolov5 import YOLOv5Checkpoint
+from supervisely.nn.models.yolov5 import YOLOv5
 from supervisely.app.v1.app_service import AppService
 from dotenv import load_dotenv
 
@@ -56,13 +56,13 @@
 local_artifacts_dir = os.path.join(runs_dir, experiment_name)
 sly.logger.info(f"All training artifacts will be saved to local directory {local_artifacts_dir}")
 
-checkpoint = YOLOv5Checkpoint(team_id)
-model_dir = checkpoint.get_model_dir()
+sly_yolov5 = YOLOv5(team_id)
+model_dir = sly_yolov5.framework_dir
 
 remote_artifacts_dir = os.path.join(model_dir, project_info.name, experiment_name)
 remote_artifacts_dir = api.file.get_free_dir_name(team_id, remote_artifacts_dir)
 
-remote_weights_dir = os.path.join(remote_artifacts_dir, checkpoint.weights_dir)
+remote_weights_dir = os.path.join(remote_artifacts_dir, sly_yolov5.weights_dir)
 remote_weights_dir = api.file.get_free_dir_name(team_id, remote_artifacts_dir)
 
 sly.logger.info(f"After training artifacts will be uploaded to Team Files: {remote_artifacts_dir}")
diff --git a/supervisely/train/src/sly_utils.py b/supervisely/train/src/sly_utils.py
@@ -69,12 +69,12 @@ def _gen_message(current, total):
         time.sleep(0.5)
 
     # generate metadata
-    globals.checkpoint.generate_sly_metadata(
-        app_name=globals.checkpoint.app_name,
+    globals.sly_yolov5.generate_sly_metadata(
+        app_name=globals.sly_yolov5.app_name,
         session_id=globals.experiment_name,
         session_path=globals.remote_artifacts_dir,
         weights_path=globals.remote_weights_dir,
         training_project_name=globals.project_info.name,
-        task_type=globals.checkpoint.task_type,
+        task_type=globals.sly_yolov5.task_type,
         config_path=None,
     )