Fixing format export for CLI/python sdk

Author: digaobarbosa

roboflow/__init__.py

@@ -15,7 +15,7 @@
 from roboflow.models import CLIPModel, GazeModel  # noqa: F401
 from roboflow.util.general import write_line
 
-__version__ = "1.2.4"
+__version__ = "1.2.6"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

roboflow/core/version.py

@@ -33,7 +33,7 @@
 from roboflow.util.annotations import amend_data_yaml
 from roboflow.util.general import write_line
 from roboflow.util.model_processor import process
-from roboflow.util.versions import get_wrong_dependencies_versions, normalize_yolo_model_type
+from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions, normalize_yolo_model_type
 
 if TYPE_CHECKING:
     import numpy as np
@@ -235,7 +235,7 @@ def download(self, model_format=None, location=None, overwrite: bool = False):
 
         return Dataset(self.name, self.version, model_format, os.path.abspath(location))
 
-    def export(self, model_format=None):
+    def export(self, model_format=None) -> str:
         """
         Ask the Roboflow API to generate a version's dataset in a given format so that it can be downloaded via the `download()` method.
 
@@ -245,7 +245,7 @@ def export(self, model_format=None):
             model_format (str): A format to use for downloading
 
         Returns:
-            True
+            The URL of the exported dataset.
 
         Raises:
             RuntimeError: If the Roboflow API returns an error with a helpful JSON body
@@ -283,7 +283,7 @@ def export(self, model_format=None):
             sys.stdout.write("\n")
             print("\r" + "Version export complete for " + model_format + " format")
             sys.stdout.flush()
-            return True
+            return url
         else:
             try:
                 raise RuntimeError(response.json())
@@ -310,26 +310,17 @@ def train(self, speed=None, model_type=None, checkpoint=None, plot_in_notebook=F
 
         self.__wait_if_generating()
 
-        train_model_format = "yolov5pytorch"
-
-        if self.type == TYPE_CLASSICATION:
-            train_model_format = "folder"
-
-        if self.type == TYPE_INSTANCE_SEGMENTATION:
-            train_model_format = "yolov5pytorch"
-
-        if self.type == TYPE_SEMANTIC_SEGMENTATION:
-            train_model_format = "png-mask-semantic"
-
-        # if classification
-        if train_model_format not in self.exports:
-            self.export(train_model_format)
+        train_model_format = get_model_format(model_type)
 
         workspace, project, *_ = self.id.rsplit("/")
         url = f"{API_URL}/{workspace}/{project}/{self.version}/train"
+        link = self.export(train_model_format)
 
         data = {}
 
+        if link:
+            data["link"] = link
+
         if speed:
             data["speed"] = speed
 
@@ -341,6 +332,7 @@ def train(self, speed=None, model_type=None, checkpoint=None, plot_in_notebook=F
             data["modelType"] = model_type
 
         write_line("Reaching out to Roboflow to start training...")
+        print(data)
 
         response = requests.post(url, json=data, params={"api_key": self.__api_key})
         if not response.ok:

roboflow/util/versions.py

@@ -95,3 +95,40 @@ def normalize_yolo_model_type(model_type: str) -> str:
     model_type = model_type.replace("yolo11", "yolov11")
     model_type = model_type.replace("yolo12", "yolov12")
     return model_type
+
+
+def get_model_format(model_type: str) -> str:
+    """
+    Get the model format for a given model type.
+    Args:
+        model_type (str): The model type to get the format for.
+
+    Returns:
+        str: The model format.
+
+    Example:
+        >>> get_model_format("yolov5v6n")
+        "yolov5pytorch"
+        >>> get_model_format("rfdetr-nano")
+        "coco"
+        >>> get_model_format("yolov11n")
+        "yolov5pytorch"
+    """
+    # Prefixes extrated from modelRegistry.js in roboflow.
+    model_formats = {
+        "yolo": "yolov5pytorch",
+        "pali": "jsonl",
+        "flor": "jsonl",
+        "qwen": "jsonl",
+        "smol": "jsonl",
+        "vit-b": "folder",
+        "resn": "folder",
+        "rfdetr": "coco",
+        "rf-detr": "coco",
+        "deep": "png-mask-semantic",
+    }
+
+    for prefix, format in model_formats.items():
+        if prefix in model_type:
+            return format
+    return "yolov5pytorch"

tests/util/test_versions.py

@@ -1,7 +1,7 @@
 import unittest
 from importlib import import_module
 
-from roboflow.util.versions import get_wrong_dependencies_versions
+from roboflow.util.versions import get_wrong_dependencies_versions, get_model_format
 
 
 class TestVersions(unittest.TestCase):
@@ -23,3 +23,21 @@ def test_wrong_dependencies_versions(self):
             wrong_dependencies_versions = get_wrong_dependencies_versions([test])
             is_correct_dep = len(wrong_dependencies_versions) == 0
             self.assertEqual(is_correct_dep, expected_result)
+
+
+class TestGetModelFormat(unittest.TestCase):
+    def test_get_model_format_with_various_ids(self):
+        cases = [
+            ("yolov5v2s", "yolov5pytorch"),
+            ("yolov11n", "yolov5pytorch"),
+            ("rf-detr-nas-parent", "coco"),
+            ("rfdetr-nano", "coco"),
+            ("vit-base-patch16-224-in21k", "folder"),
+            ("resnet14", "folder"),
+            ("resenet38", "yolov5pytorch"),
+            ("invlid-type", "yolov5pytorch"),
+        ]
+
+        for model_type, expected_format in cases:
+            with self.subTest(model_type=model_type):
+                self.assertEqual(get_model_format(model_type), expected_format)