training tweaks

Author: Jacobsolawetz

roboflow/core/project.py

@@ -209,11 +209,11 @@ def train(
 
         new_version = self.generate_version(settings=new_version_settings)
         new_version = self.version(new_version)
-        new_version.train(
+        new_model = new_version.train(
             speed=speed, checkpoint=checkpoint, plot_in_notebook=plot_in_notebook
         )
 
-        return True
+        return new_model
 
     def version(self, version_number, local=None):
         """Retrieves information about a specific version, and throws it into an object.
@@ -348,7 +348,7 @@ def __annotation_upload(self, annotation_path, image_id):
         annotation_string = open(annotation_path, "r").read()
         # Set annotation upload url
 
-        project_name = self.id.rsplit("/")[1]   
+        project_name = self.id.rsplit("/")[1]
 
         self.annotation_upload_url = "".join(
             [
@@ -528,14 +528,18 @@ def single_upload(
                 if "success" in response_data.keys():
                     annotation_success = True
                 elif "error" in response_data.keys():
-                    warnings.warn(f"Uploading annotation data for image failed: {str(response_data['error'])}")
+                    warnings.warn(
+                        f"Uploading annotation data for image failed: {str(response_data['error'])}"
+                    )
                     annotation_success = False
                 else:
-                    warnings.warn(f"Uploading annotation data for image failed: {str(response_data)}")
-                    annotation_success = False  
+                    warnings.warn(
+                        f"Uploading annotation data for image failed: {str(response_data)}"
+                    )
+                    annotation_success = False
             except:
                 warnings.warn(f"Bad response: {response.status_code}")
-                annotation_success = False  
+                annotation_success = False
 
             # Give user warning that annotation failed to upload
             if not annotation_success:

roboflow/core/version.py

@@ -349,7 +349,7 @@ def live_plot(epochs, mAP, loss, title=""):
         previous_epochs = []
         num_machine_spin_dots = []
 
-        while status == "training":
+        while status == "training" or status == "running":
             url = (
                 f"{API_URL}/{self.workspace}/{self.project}/{self.version}?nocache=true"
             )
@@ -366,8 +366,7 @@ def live_plot(epochs, mAP, loss, title=""):
                     status = "finished"
                     break
                 if "status" in version["train"].keys():
-                    status = version["train"]["status"]
-                    if status == "failed":
+                    if version["train"]["status"] == "failed":
                         write_line(line="Training failed")
                         break
 
@@ -396,7 +395,7 @@ def live_plot(epochs, mAP, loss, title=""):
                     ]
                 )
 
-                title = "Training in Progress: "
+                title = "Training in Progress"
                 # plottling logic
             else:
                 num_machine_spin_dots.append(".")
@@ -415,7 +414,7 @@ def live_plot(epochs, mAP, loss, title=""):
                     if len(epochs) > 0:
                         title = (
                             title
-                            + " Epoch: "
+                            + ": Epoch: "
                             + str(epochs[-1])
                             + " mAP: "
                             + str(mAP[-1])