Merge pull request #138 from roboflow/ptDownload2

Pt download2

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.4"
+__version__ = "1.0.5"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

roboflow/core/version.py

@@ -433,32 +433,6 @@ 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

roboflow/models/object_detection.py

@@ -4,16 +4,18 @@
 import json
 import os
 import random
+import sys
 import urllib
 from pathlib import Path
 
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
 import requests
+import wget
 from PIL import Image
 
-from roboflow.config import OBJECT_DETECTION_MODEL
+from roboflow.config import API_URL, OBJECT_DETECTION_MODEL
 from roboflow.util.image_utils import check_image_url
 from roboflow.util.prediction import PredictionGroup
 from roboflow.util.versions import (
@@ -459,6 +461,38 @@ def view(button):
         else:
             view(stopButton)
 
+    def download(self, format="pt", location="."):
+        supported_formats = ["pt"]
+        if format not in supported_formats:
+            raise Exception(
+                f"Unsupported format {format}. Must be one of {supported_formats}"
+            )
+
+        workspace, project, version = 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
+
     def __exception_check(self, image_path_check=None):
         # Check if Image path exists exception check (for both hosted URL and local image)
         if image_path_check is not None: