Merge pull request #89 from roboflow/pipGenerate

Jacobsolawetz · web-flow · commit f790b660a6dd · 2022-12-30T11:11:57.000-06:00
Code Actions V1
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-certifi==2021.5.30
+certifi==2022.12.7
 chardet==4.0.0
 cycler==0.10.0
 idna==2.10
diff --git a/roboflow/__init__.py b/roboflow/__init__.py
@@ -8,7 +8,7 @@
 from roboflow.core.project import Project
 from roboflow.core.workspace import Workspace
 
-__version__ = "0.2.21"
+__version__ = "0.2.22"
 
 
 def check_key(api_key, model, notebook, num_retries=0):
diff --git a/roboflow/core/project.py b/roboflow/core/project.py
@@ -2,6 +2,7 @@
 import io
 import json
 import os
+import sys
 import urllib
 import warnings
 
@@ -96,6 +97,119 @@ def versions(self):
             version_array.append(version_object)
         return version_array
 
+    def generate_version(self, settings):
+
+        """
+        Settings, a python dict with augmentation and preprocessing keys and specifications for generation.
+        These settings mirror capabilities available via the Roboflow UI.
+        For example:
+             {
+                "augmentation": {
+                    "bbblur": { "pixels": 1.5 },
+                    "bbbrightness": { "brighten": true, "darken": false, "percent": 91 },
+                    "bbcrop": { "min": 12, "max": 71 },
+                    "bbexposure": { "percent": 30 },
+                    "bbflip": { "horizontal": true, "vertical": false },
+                    "bbnoise": { "percent": 50 },
+                    "bbninety": { "clockwise": true, "counter-clockwise": false, "upside-down": false },
+                    "bbrotate": { "degrees": 45 },
+                    "bbshear": { "horizontal": 45, "vertical": 45 },
+                    "blur": { "pixels": 1.5 },
+                    "brightness": { "brighten": true, "darken": false, "percent": 91 },
+                    "crop": { "min": 12, "max": 71 },
+                    "cutout": { "count": 26, "percent": 71 },
+                    "exposure": { "percent": 30 },
+                    "flip": { "horizontal": true, "vertical": false },
+                    "hue": { "degrees": 180 },
+                    "image": { "versions": 32 },
+                    "mosaic": true,
+                    "ninety": { "clockwise": true, "counter-clockwise": false, "upside-down": false },
+                    "noise": { "percent": 50 },
+                    "rgrayscale": { "percent": 50 },
+                    "rotate": { "degrees": 45 },
+                    "saturation": { "percent": 50 },
+                    "shear": { "horizontal": 45, "vertical": 45 }
+                },
+                "preprocessing": {
+                    "auto-orient": true,
+                    "contrast": { "type": "Contrast Stretching" },
+                    "filter-null": { "percent": 50 },
+                    "grayscale": true,
+                    "isolate": true,
+                    "remap": { "original_class_name": "new_class_name" },
+                    "resize": { "width": 200, "height": 200, "format": "Stretch to" },
+                    "static-crop": { "x_min": 10, "x_max": 90, "y_min": 10, "y_max": 90 },
+                    "tile": { "rows": 2, "columns": 2 }
+                }
+            }
+
+        Returns: The version number that is being generated
+        """
+
+        if not {"augmentation", "preprocessing"} <= settings.keys():
+            raise (
+                RuntimeError(
+                    "augmentation and preprocessing keys are required to generate. If none are desired specify empty dict associated with that key."
+                )
+            )
+
+        r = requests.post(
+            f"{API_URL}/{self.__workspace}/{self.__project_name}/generate?api_key={self.__api_key}",
+            json=settings,
+        )
+
+        try:
+            r_json = r.json()
+        except:
+            raise ("Error when requesting to generate a new version for project.")
+
+        # if the generation succeeds, return the version that is being generated
+        if r.status_code == 200:
+            sys.stdout.write(
+                "\r"
+                + r_json["message"]
+                + " for new version "
+                + str(r_json["version"])
+                + "."
+            )
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+            return int(r_json["version"])
+        else:
+            if "error" in r_json.keys():
+                raise RuntimeError(r_json["error"])
+            else:
+                raise RuntimeError(json.dumps(r_json))
+
+    def train(
+        self,
+        new_version_settings={
+            "preprocessing": {
+                "auto-orient": True,
+                "resize": {"width": 640, "height": 640, "format": "Stretch to"},
+            },
+            "augmentation": {},
+        },
+        speed=None,
+        checkpoint=None,
+    ) -> bool:
+        """
+        Ask the Roboflow API to train a previously exported version's dataset.
+        Args:
+            speed: Whether to train quickly or accurately. Note: accurate training is a paid feature. Default speed is `fast`.
+            checkpoint: A string representing the checkpoint to use while training
+        Returns:
+            True
+            RuntimeError: If the Roboflow API returns an error with a helpful JSON body
+            HTTPError: If the Network/Roboflow API fails and does not return JSON
+        """
+
+        new_version = self.generate_version(settings=new_version_settings)
+        new_version = self.version(new_version)
+        new_version.train(speed=speed, checkpoint=checkpoint)
+
+        return True
+
     def version(self, version_number, local=None):
         """Retrieves information about a specific version, and throws it into an object.
         :param version_number: the version number that you want to retrieve
diff --git a/roboflow/core/version.py b/roboflow/core/version.py
@@ -1,6 +1,7 @@
 import json
 import os
 import sys
+import time
 import zipfile
 
 import requests
@@ -69,6 +70,10 @@ def __init__(
             self.model_format = model_format
             self.workspace = workspace
             self.project = project
+            if "exports" in version_dict.keys():
+                self.exports = version_dict["exports"]
+            else:
+                self.exports = []
 
             version_without_workspace = os.path.basename(str(version))
 
@@ -101,6 +106,42 @@ def __init__(
             else:
                 self.model = None
 
+    def __check_if_generating(self):
+        # check Roboflow API to see if this version is still generating
+
+        url = f"{API_URL}/{self.workspace}/{self.project}/{self.version}?nocache=true"
+        response = requests.get(url, params={"api_key": self.__api_key})
+        response.raise_for_status()
+
+        if response.json()["version"]["progress"] == None:
+            progress = 0.0
+        else:
+            progress = float(response.json()["version"]["progress"])
+
+        return response.json()["version"]["generating"], progress
+
+    def __wait_if_generating(self, recurse=False):
+        # checks if a given version is still in the progress of generating
+
+        still_generating, progress = self.__check_if_generating()
+
+        if still_generating:
+            progress_message = (
+                "Generating version still in progress. Progress: "
+                + str(round(progress * 100, 2))
+                + "%"
+            )
+            sys.stdout.write("\r" + progress_message)
+            sys.stdout.flush()
+            time.sleep(5)
+            return self.__wait_if_generating(recurse=True)
+
+        else:
+            if recurse:
+                sys.stdout.write("\n")
+                sys.stdout.flush()
+            return
+
     def download(self, model_format=None, location=None):
         """
         Download and extract a ZIP of a version's dataset in a given format
@@ -110,11 +151,19 @@ def download(self, model_format=None, location=None):
 
         :return: Dataset
         """
-        if location is None:
-            location = self.__get_download_location()
+
+        self.__wait_if_generating()
 
         model_format = self.__get_format_identifier(model_format)
 
+        if model_format not in self.exports:
+            self.export(model_format)
+
+        # if model_format is not in
+
+        if location is None:
+            location = self.__get_download_location()
+
         if self.__api_key == "coco-128-sample":
             link = "https://app.roboflow.com/ds/n9QwXwUK42?key=NnVCe2yMxP"
         else:
@@ -144,14 +193,92 @@ def export(self, model_format=None):
         :return: True
         :raises RuntimeError / HTTPError:
         """
+
+        model_format = self.__get_format_identifier(model_format)
+
+        self.__wait_if_generating()
+
         url = self.__get_download_url(model_format)
-        response = requests.post(url, params={"api_key": self.__api_key})
+        response = requests.get(url, params={"api_key": self.__api_key})
         if not response.ok:
             try:
                 raise RuntimeError(response.json())
             except requests.exceptions.JSONDecodeError:
                 response.raise_for_status()
 
+        # the rest api returns 202 if the export is still in progress
+        if response.status_code == 202:
+            status_code_check = 202
+            while status_code_check == 202:
+                time.sleep(1)
+                response = requests.get(url, params={"api_key": self.__api_key})
+                status_code_check = response.status_code
+                if status_code_check == 202:
+                    progress = response.json()["progress"]
+                    progress_message = (
+                        "Exporting format "
+                        + model_format
+                        + " in progress : "
+                        + str(round(progress * 100, 2))
+                        + "%"
+                    )
+                    sys.stdout.write("\r" + progress_message)
+                    sys.stdout.flush()
+
+        if response.status_code == 200:
+            sys.stdout.write("\n")
+            print("\r" + "Version export complete for " + model_format + " format")
+            sys.stdout.flush()
+            return True
+        else:
+            try:
+                raise RuntimeError(response.json())
+            except requests.exceptions.JSONDecodeError:
+                response.raise_for_status()
+
+    def train(self, speed=None, checkpoint=None) -> bool:
+        """
+        Ask the Roboflow API to train a previously exported version's dataset.
+        Args:
+            speed: Whether to train quickly or accurately. Note: accurate training is a paid feature. Default speed is `fast`.
+            checkpoint: A string representing the checkpoint to use while training
+        Returns:
+            True
+            RuntimeError: If the Roboflow API returns an error with a helpful JSON body
+            HTTPError: If the Network/Roboflow API fails and does not return JSON
+        """
+
+        self.__wait_if_generating()
+
+        train_model_format = "yolov5pytorch"
+
+        if train_model_format not in self.exports:
+            self.export(train_model_format)
+
+        workspace, project, *_ = self.id.rsplit("/")
+        url = f"{API_URL}/{workspace}/{project}/{self.version}/train"
+
+        data = {}
+        if speed:
+            data["speed"] = speed
+
+        if checkpoint:
+            data["checkpoint"] = checkpoint
+
+        sys.stdout.write("\r" + "Reaching out to Roboflow to start training...")
+        sys.stdout.write("\n")
+        sys.stdout.flush()
+
+        response = requests.post(url, json=data, params={"api_key": self.__api_key})
+        if not response.ok:
+            try:
+                raise RuntimeError(response.json())
+            except requests.exceptions.JSONDecodeError:
+                response.raise_for_status()
+
+        sys.stdout.write("\r" + "Training model in progress...")
+        sys.stdout.flush()
+
         return True
 
     def upload_model(self, model_path: str) -> None:
diff --git a/setup.py b/setup.py
@@ -21,7 +21,7 @@
     long_description_content_type="text/markdown",
     url="https://github.com/roboflow-ai/roboflow-python",
     install_requires=[
-        "certifi==2021.5.30",
+        "certifi==2022.12.7",
         "chardet==4.0.0",
         "cycler==0.10.0",
         "glob2",
diff --git a/tests/test_version.py b/tests/test_version.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-certifi==2021.5.30`
	`1`	`+certifi==2022.12.7`
`2`	`2`	`chardet==4.0.0`
`3`	`3`	`cycler==0.10.0`
`4`	`4`	`idna==2.10`