Merge pull request #137 from roboflow/ptDownload

PT Weights Download

Author: tonylampada

roboflow/__init__.py

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

roboflow/core/version.py

@@ -433,6 +433,32 @@ def live_plot(epochs, mAP, loss, title=""):
         # return the model object
         return self.model
 
+    def get_pt_weights(self, location="."):
+        workspace, project, *_ = self.id.rsplit("/")
+
+        # get pt url
+        pt_api_url = f"{API_URL}/{workspace}/{project}/{self.version}/ptFile"
+
+        r = requests.get(pt_api_url, params={"api_key": self.__api_key})
+
+        r.raise_for_status()
+
+        pt_weights_url = r.json()["weightsUrl"]
+
+        def bar_progress(current, total, width=80):
+            progress_message = (
+                "Downloading weights to "
+                + location
+                + "/weights.pt"
+                + ": %d%% [%d / %d] bytes" % (current / total * 100, current, total)
+            )
+            sys.stdout.write("\r" + progress_message)
+            sys.stdout.flush()
+
+        wget.download(pt_weights_url, out=location + "/weights.pt", bar=bar_progress)
+
+        return
+
     # @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