pt download

Author: Jacobsolawetz

roboflow/core/version.py

@@ -432,6 +432,34 @@ 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: