Updated deploy method of version class to support more model types.

Author: paulguerrie

roboflow/core/version.py

@@ -312,16 +312,18 @@ def deploy(self, model_type: str, model_path: str) -> None:
             model_path (str): File path to model weights to be uploaded
         """
 
-        supported_models = ["yolov8", "yolov5", "yolov7-seg"]
+        supported_models = ["yolov5", "yolov7", "yolov8"]
 
-        if model_type not in supported_models:
+        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 {supported_models}"
                 )
             )
 
-        if model_type == "yolov8":
+        if "yolov8" in model_type:
             try:
                 import torch
                 import ultralytics
@@ -335,7 +337,7 @@ def deploy(self, model_type: str, model_path: str) -> None:
                 [("ultralytics", "<=", "8.0.20")]
             )
 
-        elif model_type in ["yolov5", "yolov7-seg"]:
+        elif "yolov5" in model_type or "yolov7" in model_type:
             try:
                 import torch
             except ImportError as e:
@@ -345,22 +347,31 @@ def deploy(self, model_type: str, model_path: str) -> None:
 
         model = torch.load(os.path.join(model_path, "weights/best.pt"))
 
-        class_names = []
-        for i, val in enumerate(model["model"].names):
-            class_names.append((val, model["model"].names[val]))
+        if isinstance(model["model"].names, list):
+            class_names = model["model"].names
+        else:
+            class_names = []
+            for i, val in enumerate(model["model"].names):
+                class_names.append((val, model["model"].names[val]))
         class_names.sort(key=lambda x: x[0])
         class_names = [x[1] for x in class_names]
 
-        if model_type == "yolov8":
+        if "yolov8" in model_type:
             # try except for backwards compatibility with older versions of ultralytics
+            if "-cls" in model_type:
+                nc = model["model"].yaml["nc"]
+                args = model["train_args"]
+            else:
+                nc = model["model"].nc
+                args = model["model"].args
             try:
                 model_artifacts = {
                     "names": class_names,
                     "yaml": model["model"].yaml,
-                    "nc": model["model"].nc,
+                    "nc": nc,
                     "args": {
                         k: val
-                        for k, val in model["model"].args.items()
+                        for k, val in args.items()
                         if ((k == "model") or (k == "imgsz") or (k == "batch"))
                     },
                     "ultralytics_version": ultralytics.__version__,
@@ -370,33 +381,39 @@ def deploy(self, model_type: str, model_path: str) -> None:
                 model_artifacts = {
                     "names": class_names,
                     "yaml": model["model"].yaml,
-                    "nc": model["model"].nc,
+                    "nc": nc,
                     "args": {
                         k: val
-                        for k, val in model["model"].args.__dict__.items()
+                        for k, val in args.__dict__.items()
                         if ((k == "model") or (k == "imgsz") or (k == "batch"))
                     },
                     "ultralytics_version": ultralytics.__version__,
                     "model_type": model_type,
                 }
-        elif model_type in ["yolov5", "yolov7-seg"]:
+        elif "yolov5" in model_type or "yolov7" in model_type:
             # parse from yaml for yolov5
 
             with open(os.path.join(model_path, "opt.yaml"), "r") as stream:
                 opts = yaml.safe_load(stream)
 
             model_artifacts = {
                 "names": class_names,
-                "yaml": model["model"].yaml,
                 "nc": model["model"].nc,
-                "args": {"imgsz": opts["imgsz"], "batch": opts["batch_size"]},
+                "args": {
+                    "imgsz": opts["imgsz"] if "imgsz" in opts else opts["img_size"],
+                    "batch": opts["batch_size"],
+                },
                 "model_type": model_type,
             }
+            if hasattr(model["model"], "yaml"):
+                model_artifacts["yaml"] = model["model"].yaml
 
-        with open(model_path + "model_artifacts.json", "w") as fp:
+        with open(os.path.join(model_path, "model_artifacts.json"), "w") as fp:
             json.dump(model_artifacts, fp)
 
-        torch.save(model["model"].state_dict(), model_path + "state_dict.pt")
+        torch.save(
+            model["model"].state_dict(), os.path.join(model_path, "state_dict.pt")
+        )
 
         lista_files = [
             "results.csv",
@@ -405,11 +422,13 @@ def deploy(self, model_type: str, model_path: str) -> None:
             "state_dict.pt",
         ]
 
-        with zipfile.ZipFile(model_path + "roboflow_deploy.zip", "w") as zipMe:
+        with zipfile.ZipFile(
+            os.path.join(model_path, "roboflow_deploy.zip"), "w"
+        ) as zipMe:
             for file in lista_files:
-                if os.path.exists(model_path + file):
+                if os.path.exists(os.path.join(model_path, file)):
                     zipMe.write(
-                        model_path + file,
+                        os.path.join(model_path, file),
                         arcname=file,
                         compress_type=zipfile.ZIP_DEFLATED,
                     )
@@ -422,7 +441,7 @@ def deploy(self, model_type: str, model_path: str) -> None:
                         )
 
         res = requests.get(
-            f"{API_URL}/{self.workspace}/{self.project}/{self.version}/uploadModel?api_key={self.__api_key}"
+            f"{API_URL}/{self.workspace}/{self.project}/{self.version}/uploadModel?api_key={self.__api_key}&modelType={model_type}"
         )
         try:
             if res.status_code == 429:
@@ -437,7 +456,7 @@ def deploy(self, model_type: str, model_path: str) -> None:
 
         res = requests.put(
             res.json()["url"],
-            data=open(os.path.join(model_path + "roboflow_deploy.zip"), "rb"),
+            data=open(os.path.join(model_path, "roboflow_deploy.zip"), "rb"),
         )
         try:
             res.raise_for_status()
@@ -590,7 +609,8 @@ def callback(content: dict) -> dict:
                 pass
             return content
 
-        amend_data_yaml(path=data_path, callback=callback)
+        if os.path.exists(data_path):
+            amend_data_yaml(path=data_path, callback=callback)
 
     def __str__(self):
         """string representation of version object."""