Skip to content

Commit

Permalink
roboflow url support
Browse files Browse the repository at this point in the history
  • Loading branch information
kozlov721 committed Oct 7, 2024
1 parent 4fa0a02 commit 162a4da
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 6 deletions.
58 changes: 52 additions & 6 deletions luxonis_ml/data/parsers/luxonis_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -72,8 +74,15 @@ def __init__(
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.
Expand All @@ -97,9 +106,16 @@ def __init__(
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)
name = dataset_dir.split("/")[-1]
if dataset_dir.startswith("roboflow://"):
self.dataset_dir = self._download_roboflow_dataset(
dataset_dir, save_dir
)
else:
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
Expand Down Expand Up @@ -237,3 +253,33 @@ def _parse_split(
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]
) -> Path:
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))
)
print(dataset.location)
return Path(dataset.location)
2 changes: 2 additions & 0 deletions luxonis_ml/utils/environ.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"] = (
Expand Down

0 comments on commit 162a4da

Please sign in to comment.