Merge pull request #123 from roboflow/feature/yolo-model-uploads

YOLO Model Deploys

Author: paulguerrie

roboflow/__init__.py

@@ -11,7 +11,7 @@
 from roboflow.core.workspace import Workspace
 from roboflow.util.general import write_line
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

roboflow/core/version.py

@@ -441,16 +441,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-seg", "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
@@ -464,7 +466,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:
@@ -483,16 +485,22 @@ def deploy(self, model_type: str, model_path: str) -> None:
             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__,
@@ -502,33 +510,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",
@@ -537,11 +551,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,
                     )
@@ -554,7 +570,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:
@@ -569,7 +585,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()