set task type for runtimes

supervisely-ecosystem · Nov 26, 2024 · 757c08f · 757c08f
1 parent f9bdaff
commit 757c08f
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/serve/src/yolov8.py b/serve/src/yolov8.py
@@ -578,7 +578,14 @@ def _load_runtime(self, weights_path: str, format: str, **kwargs):
                 self._dump_yaml_checkpoint_info(model, os.path.dirname(weights_path))
         else:
             exported_weights_path = weights_path
-        model = YOLO(exported_weights_path)
+
+        task_type_map = {
+            "object detection": "detect",
+            "instance segmentation": "segment",
+            "pose estimation": "pose",
+        }
+
+        model = YOLO(exported_weights_path, task=task_type_map[self.task_type])
         return model
 
     def _load_onnx(self, weights_path: str):