LuxonisParser - RoboFlow URL Support

.github/workflows/ci.yaml

@@ -41,6 +41,9 @@ jobs:
       with:
         ref: ${{ github.head_ref }}
 
+    - name: Install pre-commit
+      run: python3 -m pip install 'pre-commit<4.0.0'
+
     - name: Run pre-commit
       uses: pre-commit/[email protected]
 

luxonis_ml/data/parsers/luxonis_parser.py

@@ -13,10 +13,12 @@
     overload,
 )
 
+from roboflow import Roboflow
+
 from luxonis_ml.data import DATASETS_REGISTRY, BaseDataset, LuxonisDataset
 from luxonis_ml.data.utils.enums import LabelType
 from luxonis_ml.enums import DatasetType
-from luxonis_ml.utils import LuxonisFileSystem
+from luxonis_ml.utils import LuxonisFileSystem, environ
 
 from .base_parser import BaseParser
 from .classification_directory_parser import ClassificationDirectoryParser
@@ -72,8 +74,15 @@
         appropriate parser.
 
         @type dataset_dir: str
-        @param dataset_dir: Path to the dataset directory or zip file.
-            Can also be a remote URL supported by L{LuxonisFileSystem}.
+        @param dataset_dir: Identifier of the dataset directory.
+            Can be one of:
+                - Local path to the dataset directory.
+                - Remote URL supported by L{LuxonisFileSystem}.
+                  - C{gcs://} for Google Cloud Storage
+                  - C{s3://} for Amazon S3
+                - C{roboflow://} for Roboflow datasets.
+                  - Expected format: C{roboflow://workspace/project/version/format}.
+            Can be a remote URL supported by L{LuxonisFileSystem}.
         @type dataset_name: Optional[str]
         @param dataset_name: Name of the dataset. If C{None}, the name
             is derived from the name of the dataset directory.
@@ -97,9 +106,16 @@
             names.
         """
         save_dir = Path(save_dir) if save_dir else None
-        name = Path(dataset_dir).name
-        local_path = (save_dir or Path.cwd()) / name
-        self.dataset_dir = LuxonisFileSystem.download(dataset_dir, local_path)
+        if dataset_dir.startswith("roboflow://"):
+            self.dataset_dir, name = self._download_roboflow_dataset(
+                dataset_dir, save_dir
+            )
+        else:
+            name = dataset_dir.split("/")[-1]
+            local_path = (save_dir or Path.cwd()) / name
+            self.dataset_dir = LuxonisFileSystem.download(
+                dataset_dir, local_path
+            )
         if self.dataset_dir.suffix == ".zip":
             with zipfile.ZipFile(self.dataset_dir, "r") as zip_ref:
                 unzip_dir = self.dataset_dir.parent / self.dataset_dir.stem
@@ -237,3 +253,38 @@
         return self.parser.parse_split(
             split, random_split, split_ratios, **parsed_kwargs, **kwargs
         )
+
+    def _download_roboflow_dataset(
+        self, dataset_dir: str, local_path: Optional[Path]
+    ) -> Tuple[Path, str]:
+        if environ.ROBOFLOW_API_KEY is None:
+            raise RuntimeError(
+                "ROBOFLOW_API_KEY environment variable is not set. "
+                "Please set it to your Roboflow API key."
+            )
+
+        rf = Roboflow(api_key=environ.ROBOFLOW_API_KEY)
+        parts = dataset_dir.split("roboflow://")[1].split("/")
+        if len(parts) != 4:
+            raise ValueError(
+                f"Incorrect Roboflow dataset URL: `{dataset_dir}`. "
+                "Expected format: `roboflow://workspace/project/version/format`."
+            )
+        workspace, project, version, format = dataset_dir.split("roboflow://")[
+            1
+        ].split("/")
+        try:
+            version = int(version)
+        except ValueError as e:
+            raise ValueError(
+                f"Roboflow version must be an integer, got `{version}`."
+            ) from e
+
+        local_path = local_path or Path.cwd() / f"{project}_{format}"
+        dataset = (
+            rf.workspace(workspace)
+            .project(project)
+            .version(int(version))
+            .download(format, str(local_path / project))
+        )
+        return Path(dataset.location), project

luxonis_ml/data/requirements.txt

@@ -12,3 +12,4 @@ pycocotools>=2.0.7
 typeguard>=4.1.0
 polars[timezone]>=0.20.31
 ordered-set>=4.0.0
+roboflow>=0.1.1

luxonis_ml/utils/environ.py

@@ -34,6 +34,8 @@ class Environ(BaseSettings):
     LUXONISML_BASE_PATH: Path = Path.home() / "luxonis_ml"
     LUXONISML_TEAM_ID: str = "offline"
 
+    ROBOFLOW_API_KEY: Optional[str] = None
+
     GOOGLE_APPLICATION_CREDENTIALS: Optional[str] = None
 
     LOG_LEVEL: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = (

tests/test_data/test_parsers.py

@@ -4,6 +4,7 @@
 
 from luxonis_ml.data import LabelType, LuxonisLoader, LuxonisParser
 from luxonis_ml.enums import DatasetType
+from luxonis_ml.utils import environ
 
 URL_PREFIX: Final[str] = "gs://luxonis-test-bucket/luxonis-ml-test-data"
 WORK_DIR: Final[str] = "tests/data/parser_datasets"
@@ -82,13 +83,24 @@ def prepare_dir():
             "D1_ParkingSlot-solo.zip",
             [LabelType.BOUNDINGBOX, LabelType.SEGMENTATION],
         ),
+        (
+            DatasetType.COCO,
+            "roboflow://team-roboflow/coco-128/2/coco",
+            [LabelType.BOUNDINGBOX, LabelType.CLASSIFICATION],
+        ),
     ],
 )
 def test_dir_parser(
     dataset_type: DatasetType, url: str, expected_label_types: List[LabelType]
 ):
+    if not url.startswith("roboflow://"):
+        url = f"{URL_PREFIX}/{url}"
+
+    elif environ.ROBOFLOW_API_KEY is None:
+        pytest.skip("Roboflow API key is not set")
+
     parser = LuxonisParser(
-        f"{URL_PREFIX}/{url}",
+        url,
         dataset_name=f"test-{dataset_type}",
         delete_existing=True,
         save_dir=WORK_DIR,