diff --git a/ml_utils/ml_utils_cli/cli/apps/datasets.py b/ml_utils/ml_utils_cli/cli/apps/datasets.py index b5e6fd69..f6dc6701 100644 --- a/ml_utils/ml_utils_cli/cli/apps/datasets.py +++ b/ml_utils/ml_utils_cli/cli/apps/datasets.py @@ -158,6 +158,12 @@ def export( "provided (typically, if the source is Label Studio)" ), ] = 0.8, + error_raise: Annotated[ + bool, + typer.Option( + help="Raise an error if an image download fails, only for Ultralytics" + ), + ] = True, ): """Export Label Studio annotation, either to Hugging Face Datasets or local files (ultralytics format).""" @@ -204,6 +210,7 @@ def export( label_names_list, typing.cast(int, project_id), train_ratio=train_ratio, + error_raise=error_raise, ) elif from_ == ExportSource.hf: @@ -212,6 +219,7 @@ def export( typing.cast(str, repo_id), typing.cast(Path, output_dir), download_images=download_images, + error_raise=error_raise, ) else: raise typer.BadParameter("Unsupported export format") diff --git a/ml_utils/ml_utils_cli/cli/export.py b/ml_utils/ml_utils_cli/cli/export.py index 59afb269..abe0ee08 100644 --- a/ml_utils/ml_utils_cli/cli/export.py +++ b/ml_utils/ml_utils_cli/cli/export.py @@ -64,6 +64,7 @@ def export_from_ls_to_ultralytics( category_names: list[str], project_id: int, train_ratio: float = 0.8, + error_raise: bool = True, ): """Export annotations from a Label Studio project to the Ultralytics format. @@ -161,7 +162,7 @@ def export_from_ls_to_ultralytics( has_valid_annotation = True if has_valid_annotation: - download_output = download_image(image_url, return_bytes=True) + download_output = download_image(image_url, return_bytes=True, error_raise=error_raise) if download_output is None: logger.error("Failed to download image: %s", image_url) continue @@ -182,7 +183,9 @@ def export_from_ls_to_ultralytics( def export_from_hf_to_ultralytics( - repo_id: str, output_dir: Path, download_images: bool = True + repo_id: str, output_dir: Path, + download_images: bool = True, + error_raise: bool = True, ): """Export annotations from a Hugging Face dataset project to the Ultralytics format. @@ -207,7 +210,7 @@ def export_from_hf_to_ultralytics( image_url = sample["meta"]["image_url"] if download_images: - download_output = download_image(image_url, return_bytes=True) + download_output = download_image(image_url, return_bytes=True, error_raise=error_raise) if download_output is None: logger.error("Failed to download image: %s", image_url) continue