Merge pull request #329 from roboflow/yolov11-model-upload

Add yolov11 model deploy support
roboflow · Oct 2, 2024 · 2f4c88b · 2f4c88b
2 parents 7d4357f + 2db07cc
commit 2f4c88b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/roboflow/__init__.py b/roboflow/__init__.py
@@ -15,7 +15,7 @@
 from roboflow.models import CLIPModel, GazeModel  # noqa: F401
 from roboflow.util.general import write_line
 
-__version__ = "1.1.45"
+__version__ = "1.1.46"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

diff --git a/roboflow/core/version.py b/roboflow/core/version.py
@@ -459,8 +459,20 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best
             model_path (str): File path to the model weights to be uploaded.
             filename (str, optional): The name of the weights file. Defaults to "weights/best.pt".
         """
-
-        supported_models = ["yolov5", "yolov7-seg", "yolov8", "yolov9", "yolonas", "paligemma", "yolov10", "florence-2"]
+        if model_type.startswith("yolo11"):
+            model_type = model_type.replace("yolo11", "yolov11")
+
+        supported_models = [
+            "yolov5",
+            "yolov7-seg",
+            "yolov8",
+            "yolov9",
+            "yolonas",
+            "paligemma",
+            "yolov10",
+            "florence-2",
+            "yolov11",
+        ]
 
         if not any(supported_model in model_type for supported_model in supported_models):
             raise (ValueError(f"Model type {model_type} not supported. Supported models are" f" {supported_models}"))
@@ -519,6 +531,19 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best
                     " Please install it with `pip install torch`"
                 )
 
+        elif "yolov11" in model_type:
+            try:
+                import torch
+                import ultralytics
+
+            except ImportError:
+                raise RuntimeError(
+                    "The ultralytics python package is required to deploy yolov10"
+                    " models. Please install it with `pip install ultralytics`"
+                )
+
+            print_warn_for_wrong_dependencies_versions([("ultralytics", ">=", "8.3.0")], ask_to_continue=True)
+
         model = torch.load(os.path.join(model_path, filename))
 
         if isinstance(model["model"].names, list):
@@ -530,9 +555,9 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best
             class_names.sort(key=lambda x: x[0])
             class_names = [x[1] for x in class_names]
 
-        if "yolov8" in model_type or "yolov10" in model_type:
+        if "yolov8" in model_type or "yolov10" in model_type or "yolov11" in model_type:
             # try except for backwards compatibility with older versions of ultralytics
-            if "-cls" in model_type or model_type.startswith("yolov10"):
+            if "-cls" in model_type or model_type.startswith("yolov10") or model_type.startswith("yolov11"):
                 nc = model["model"].yaml["nc"]
                 args = model["train_args"]
             else: