Merge pull request #101 from roboflow/yolov5Upload

Upload YOLOv5 Models to Roboflow Deploy

Author: Jacobsolawetz

roboflow/__init__.py

@@ -8,7 +8,7 @@
 from roboflow.core.project import Project
 from roboflow.core.workspace import Workspace
 
-__version__ = "0.2.27"
+__version__ = "0.2.28"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

roboflow/core/version.py

@@ -25,6 +25,10 @@
 from roboflow.models.instance_segmentation import InstanceSegmentationModel
 from roboflow.models.object_detection import ObjectDetectionModel
 from roboflow.models.semantic_segmentation import SemanticSegmentationModel
+from roboflow.util.versions import (
+    print_warn_for_wrong_dependencies_versions,
+    warn_for_wrong_dependencies_versions,
+)
 
 load_dotenv()
 
@@ -153,6 +157,12 @@ def download(self, model_format=None, location=None, overwrite: bool = True):
 
         self.__wait_if_generating()
 
+        if model_format == "yolov8":
+            # we assume the user will want to use yolov8, for now we only support ultralytics=="8.11.0"
+            print_warn_for_wrong_dependencies_versions(
+                [("ultralytics", "<=", "8.0.20")]
+            )
+
         model_format = self.__get_format_identifier(model_format)
 
         if model_format not in self.exports:
@@ -284,14 +294,15 @@ def train(self, speed=None, checkpoint=None) -> bool:
 
         return True
 
+    # @warn_for_wrong_dependencies_versions([("ultralytics", "<=", "8.0.20")])
     def deploy(self, model_type: str, model_path: str) -> None:
         """Uploads provided weights file to Roboflow
 
         Args:
             model_path (str): File path to model weights to be uploaded
         """
 
-        supported_models = ["yolov8"]
+        supported_models = ["yolov8", "yolov5"]
 
         if model_type not in supported_models:
             raise (
@@ -300,24 +311,38 @@ def deploy(self, model_type: str, model_path: str) -> None:
                 )
             )
 
-        try:
-            import torch
-            import ultralytics
-        except ImportError as e:
-            raise (
-                "The ultralytics python package is required to deploy yolov8 models. Please install it with `pip install ultralytics`"
+        if model_type == "yolov8":
+            try:
+                import torch
+                import ultralytics
+
+            except ImportError as e:
+                raise (
+                    "The ultralytics python package is required to deploy yolov8 models. Please install it with `pip install ultralytics`"
+                )
+
+            print_warn_for_wrong_dependencies_versions(
+                [("ultralytics", "<=", "8.0.20")]
             )
 
-        # add logic to save torch state dict safely
-        if model_type == "yolov8":
-            model = torch.load(os.path.join(model_path + "weights/best.pt"))
+        elif model_type == "yolov5":
+            try:
+                import torch
+            except ImportError as e:
+                raise (
+                    "The torch python package is required to deploy yolov5 models. Please install it with `pip install torch`"
+                )
+
+        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]))
-            class_names.sort(key=lambda x: x[0])
-            class_names = [x[1] for x in class_names]
+        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":
+            # try except for backwards compatibility with older versions of ultralytics
             try:
                 model_artifacts = {
                     "names": class_names,
@@ -344,29 +369,39 @@ def deploy(self, model_type: str, model_path: str) -> None:
                     "ultralytics_version": ultralytics.__version__,
                     "model_type": model_type,
                 }
-
-            with open(os.path.join(model_path + "model_artifacts.json", "w")) as fp:
-                json.dump(model_artifacts, fp)
-
-            torch.save(
-                model["model"].state_dict(), os.path.join(model_path + "state_dict.pt")
-            )
-
-            lista_files = [
-                "results.csv",
-                "results.png",
-                "model_artifacts.json",
-                "state_dict.pt",
-            ]
-            with zipfile.ZipFile(
-                os.path.join(model_path + "roboflow_deploy.zip", "w")
-            ) as zipMe:
-                for file in lista_files:
-                    zipMe.write(
-                        os.path.join(model_path + file),
-                        arcname=file,
-                        compress_type=zipfile.ZIP_DEFLATED,
-                    )
+        elif model_type == "yolov5":
+            # 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"]},
+                "model_type": model_type,
+            }
+
+        with open(model_path + "model_artifacts.json", "w") as fp:
+            json.dump(model_artifacts, fp)
+
+        torch.save(model["model"].state_dict(), model_path + "state_dict.pt")
+
+        lista_files = [
+            "results.csv",
+            "results.png",
+            "model_artifacts.json",
+            "state_dict.pt",
+        ]
+
+        with zipfile.ZipFile(model_path + "roboflow_deploy.zip", "w") as zipMe:
+            for file in lista_files:
+                zipMe.write(
+                    model_path + file,
+                    arcname=file,
+                    compress_type=zipfile.ZIP_DEFLATED,
+                )
 
         res = requests.get(
             f"{API_URL}/{self.workspace}/{self.project}/{self.version}/uploadModel?api_key={self.__api_key}"
@@ -384,7 +419,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()
@@ -404,8 +439,6 @@ def deploy(self, model_type: str, model_path: str) -> None:
         except Exception as e:
             print(f"An error occured when uploading the model: {e}")
 
-    # torch.load("state_dict.pt", weights_only=True)
-
     def __download_zip(self, link, location, format):
         """
         Download a dataset's zip file from the given URL and save it in the desired location

roboflow/util/versions.py

@@ -0,0 +1,85 @@
+from importlib import import_module
+from typing import List, Tuple
+
+
+def get_wrong_dependencies_versions(
+    dependencies_versions: List[Tuple[str, str, str]]
+) -> List[Tuple[str, str, str, str]]:
+    """
+    Get a list of missmatching dependencies with current version installed.
+    E.g., assuming we pass `get_wrong_dependencies_versions([("torch", "==", "1.2.0")]), we will check if the current version of `torch` is `==1.2.0`. If not, we will return `[("torch", "==", "1.2.0", "<current_installed_version>")]
+
+    We support `<=`, `==`, `>=`
+
+    Args:
+        dependencies_versions (List[Tuple[str, str]]): List of dependencies we want to check, [("<package_name>", "<version_number_to_check")]
+
+    Returns:
+        List[Tuple[str, str, str]]: List of dependencies with wrong version, [("<package_name>", "<version_number_to_check", "<current_version>")]
+    """
+    wrong_dependencies_versions = []
+    # from e.g. "1.12.0" -> 1120
+    parse_version_to_int = lambda x: int(x.replace(".", ""))
+    order_funcs = {
+        "==": lambda x, y: x == y,
+        ">=": lambda x, y: x >= y,
+        "<=": lambda x, y: x <= y,
+    }
+    for dependency, order, version in dependencies_versions:
+        module = import_module(dependency)
+        module_version = module.__version__
+        if order not in order_funcs:
+            raise ValueError(
+                f"order={order} not supported, please use `{', '.join(order_funcs.keys())}`"
+            )
+        is_okay = order_funcs[order](
+            parse_version_to_int(module_version), parse_version_to_int(version)
+        )
+        if not is_okay:
+            wrong_dependencies_versions.append(
+                (dependency, order, version, module_version)
+            )
+    return wrong_dependencies_versions
+
+
+def print_warn_for_wrong_dependencies_versions(
+    dependencies_versions: List[Tuple[str, str, str]]
+):
+    wrong_dependencies_versions = get_wrong_dependencies_versions(dependencies_versions)
+    for (dependency, order, version, module_version) in wrong_dependencies_versions:
+        print(
+            f"Dependency {dependency}{order}{version} is required but found version={module_version}, to fix: `pip install {dependency}{order}{version}`"
+        )
+
+
+def warn_for_wrong_dependencies_versions(
+    dependencies_versions: List[Tuple[str, str, str]]
+):
+    """
+    Decorator to print a warning based on dependencies versions. E.g.
+
+    ```python
+    @warn_for_wrong_dependencies_versions([("torch", "==", "1.2.0")])
+    def foo(x):
+        # I only work with torch `1.2.0` but another one is installed
+        print(f"foo {x}")
+    ```
+
+    prints:
+
+    ```
+    Dependency torch==1.2.0 is required but found version=1.13.1, to fix: `pip install torch==1.2.0`
+    ```
+
+    Args:
+        dependencies_versions (List[Tuple[str, str]]): List of dependencies we want to check, [("<package_name>", "<version_number_to_check")]
+    """
+
+    def _inner(func):
+        def _wrapper(*args, **kwargs):
+            print_warn_for_wrong_dependencies_versions(dependencies_versions)
+            func(*args, **kwargs)
+
+        return _wrapper
+
+    return _inner

setup.py

@@ -11,6 +11,11 @@
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
+# we are using the packages in `requirements.txt` for now, 
+# not 100% ideal but will do
+with open("requirements.txt", "r") as fh:
+    install_requires = fh.read().split('\n')
+
 setuptools.setup(
     name="roboflow",
     version=version,
@@ -20,28 +25,7 @@
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/roboflow-ai/roboflow-python",
-    install_requires=[
-        "certifi==2022.12.7",
-        "chardet==4.0.0",
-        "cycler==0.10.0",
-        "glob2",
-        "idna==2.10",
-        "kiwisolver>=1.3.1",
-        "matplotlib",
-        "numpy>=1.18.5",
-        "opencv-python-headless>=4.5.1.48",
-        "Pillow>=7.1.2",
-        "pyparsing==2.4.7",
-        "python-dateutil",
-        "python-dotenv",
-        "requests",
-        "requests_toolbelt",
-        "six",
-        "urllib3==1.26.6",
-        "tqdm>=4.41.0",
-        "PyYAML>=5.3.1",
-        "wget",
-    ],
+    install_requires=install_requires,
     packages=find_packages(exclude=("tests",)),
     extras_require={
         "dev": ["flake8", "black==22.3.0", "isort", "responses", "twine", "wheel"],

tests/util/dummy_module/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.3.0"

tests/util/test_versions.py

@@ -0,0 +1,26 @@
+from importlib import import_module
+from roboflow.util.versions import get_wrong_dependencies_versions
+import unittest
+import sys
+from pathlib import Path
+
+class TestVersions(unittest.TestCase):
+
+    def test_wrong_dependencies_versions(self):
+        module_path = "tests.util.dummy_module"
+        module_version = import_module(module_path).__version__
+        tests = [
+            ("tests.util.dummy_module", "==", module_version),
+            ("tests.util.dummy_module", "<=", "0.2.0"),
+            ("tests.util.dummy_module", "<=", "1.0.0"),
+            ("tests.util.dummy_module", ">=", "0.1.0"),
+            ("tests.util.dummy_module", ">=", "0.6.0")
+
+        ]
+        # true if dep is correc
+        expected_results = [True, False, True, True, False]
+
+        for (test, expected_result) in zip(tests, expected_results):
+            wrong_dependencies_versions = get_wrong_dependencies_versions([test])
+            is_correct_dep = len(wrong_dependencies_versions) == 0
+            self.assertEqual(is_correct_dep, expected_result)