From efd86879c018c2cd4dcdbde2ddc48a83d498f1a3 Mon Sep 17 00:00:00 2001 From: mikhailbovt Date: Mon, 2 Nov 2026 00:50:19 +0300 Subject: [PATCH] Fix SDK data handling, error reporting, and transfers Preserve box coordinates when images are not resized, resolve COCO image references using contextual paths, and propagate upload and training failures. Refresh Vision Events endpoints after region changes. Stream video uploads and buffer downloads with byte progress and strict Content-Length validation. Replace weights only after a successful download, and distinguish completed dataset caches from interrupted downloads and unrelated directories. Keep upload_dataset's default None result and expose counts on request for the CLI. Reject packaging mismatches without prompting in CLI model uploads while retaining the SDK's interactive default. Discard SDK logs without retaining them in memory or failing on Unicode output. Preserve HTTP status through upload failures and map CLI errors to the documented authentication and not-found exit codes. --- CHANGELOG.md | 25 ++++ roboflow/__init__.py | 2 +- roboflow/adapters/rfapi.py | 15 ++- roboflow/adapters/vision_events_api.py | 22 ++-- roboflow/cli/_output.py | 10 +- roboflow/cli/handlers/image.py | 52 ++++---- roboflow/cli/handlers/model.py | 64 +++++----- roboflow/core/version.py | 139 +++++++++----------- roboflow/core/workspace.py | 99 +++++++-------- roboflow/models/inference.py | 34 ++--- roboflow/models/object_detection.py | 22 +--- roboflow/util/folderparser.py | 73 ++++++----- roboflow/util/general.py | 34 +++++ tests/cli/test_image_handler.py | 46 ++++++- tests/cli/test_model_handler.py | 127 ++++++++++++++++++- tests/cli/test_train_handler.py | 4 +- tests/models/test_inference.py | 131 +++++++++++++++++++ tests/models/test_object_detection.py | 38 +++++- tests/test_model_upload.py | 59 +++++++++ tests/test_project.py | 3 +- tests/test_version.py | 169 ++++++++++++++++++++++++- tests/test_vision_events.py | 40 ++++++ tests/test_workspace.py | 70 +++++++++- tests/util/test_folderparser.py | 85 +++++++++++++ tests/util/test_general.py | 83 ++++++++++++ 25 files changed, 1148 insertions(+), 298 deletions(-) create mode 100644 tests/models/test_inference.py create mode 100644 tests/test_model_upload.py create mode 100644 tests/util/test_general.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3de8e4fd..c1dcca0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,31 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Fixed + +- Dataset, weight, and search export downloads use buffered 64 KiB writes with + byte-based progress, handle missing or invalid Content-Length headers, and detect + truncated responses with both current urllib3 and supported urllib3 1.x versions. +- Local video inference uploads stream from the file instead of loading the entire video into memory. +- `Version.train()` raises on failed training instead of returning a model, including + when stale results from an earlier run are present. +- Vision Events requests follow region and environment changes made after SDK import. +- Weight downloads reject HTTP errors and replace existing weights only after the + complete response has been written successfully. +- Dataset downloads reuse a directory only after download, extraction, and YAML + processing have completed for the same version, format, and API endpoint. Interrupted + SDK downloads can be retried; other non-empty directories require `overwrite=True`. +- Dataset and model uploads propagate failures to SDK and CLI callers. Directory + uploads expose actual counts with `return_summary=True`, preserving the default + `None` result. CLI model uploads reject packaging mismatches without prompting, + and CLI upload output stays valid JSON without buffering SDK logs in memory. + Directory and model upload commands preserve HTTP status codes to distinguish + authentication, missing resources, and other failures. +- COCO imports resolve image paths relative to their annotation file, keeping + identically named images in different splits separate and rejecting ambiguous matches. +- Object detection preserves response coordinates when the input image was not + resized, and rescales boxes using the actual dimensions sent to inference. + ### Added - `Project.autolabel(..., keypoint_confidence_threshold=0.75)` and diff --git a/roboflow/__init__.py b/roboflow/__init__.py index 22aec8cb..2c2fc20b 100644 --- a/roboflow/__init__.py +++ b/roboflow/__init__.py @@ -39,7 +39,7 @@ def check_key(api_key, model, notebook, num_retries=0): response = http.post(API_URL + "/?api_key=" + api_key) if response.status_code == 401: - raise RuntimeError(response.text) + raise rfapi.AuthenticationError(response.text, status_code=response.status_code) if response.status_code != 200: # retry 5 times diff --git a/roboflow/adapters/rfapi.py b/roboflow/adapters/rfapi.py index a381f618..eeae576d 100644 --- a/roboflow/adapters/rfapi.py +++ b/roboflow/adapters/rfapi.py @@ -31,6 +31,10 @@ def __init__(self, message, status_code=None): self.status_code = status_code +class AuthenticationError(RoboflowError, RuntimeError): + """API authentication failure, compatible with existing RuntimeError handlers.""" + + class ImageUploadError(RoboflowError): def __init__(self, message, status_code=None): self.message = message @@ -49,7 +53,7 @@ def get_workspace(api_key, workspace_url): url = f"{API_URL}/{workspace_url}?api_key={api_key}" response = http.get(url) if response.status_code != 200: - raise RoboflowError(response.text) + raise RoboflowError(response.text, status_code=response.status_code) result = response.json() return result @@ -1055,7 +1059,7 @@ def init_zip_upload( body["batchName"] = batch_name response = http.post(url, params={"api_key": api_key}, json=body) if response.status_code not in (200, 201): - raise RoboflowError(response.text) + raise RoboflowError(response.text, status_code=response.status_code) return response.json() @@ -1069,7 +1073,10 @@ def upload_zip_to_signed_url(signed_url, zip_path) -> None: timeout=(60, 3600), ) if not response.ok: - raise RoboflowError(f"Zip upload to signed URL failed ({response.status_code}): {response.text}") + raise RoboflowError( + f"Zip upload to signed URL failed ({response.status_code}): {response.text}", + status_code=response.status_code, + ) def get_zip_upload_status(api_key, workspace_url, task_id) -> dict: @@ -1077,7 +1084,7 @@ def get_zip_upload_status(api_key, workspace_url, task_id) -> dict: url = f"{API_URL}/{workspace_url}/upload/zip/{task_id}" response = http.get(url, params={"api_key": api_key}) if response.status_code != 200: - raise RoboflowError(response.text) + raise RoboflowError(response.text, status_code=response.status_code) return response.json() diff --git a/roboflow/adapters/vision_events_api.py b/roboflow/adapters/vision_events_api.py index 33e6eaeb..9a2b5684 100644 --- a/roboflow/adapters/vision_events_api.py +++ b/roboflow/adapters/vision_events_api.py @@ -8,8 +8,6 @@ from roboflow.config import API_URL from roboflow.util import http -_BASE = f"{API_URL}/vision-events" - def _auth_headers(api_key: str) -> Dict[str, str]: return {"Authorization": f"Bearer {api_key}"} @@ -28,7 +26,7 @@ def write_event(api_key: str, event: Dict[str, Any]) -> dict: Raises: RoboflowError: On non-201 response status codes. """ - response = http.post(_BASE, json=event, headers=_auth_headers(api_key)) + response = http.post(f"{API_URL}/vision-events", json=event, headers=_auth_headers(api_key)) if response.status_code != 201: raise RoboflowError(response.text) return response.json() @@ -48,7 +46,7 @@ def write_batch(api_key: str, events: List[Dict[str, Any]]) -> dict: RoboflowError: On non-201 response status codes. """ response = http.post( - f"{_BASE}/batch", + f"{API_URL}/vision-events/batch", json={"events": events}, headers=_auth_headers(api_key), ) @@ -73,7 +71,7 @@ def query(api_key: str, query_params: Dict[str, Any]) -> dict: RoboflowError: On non-200 response status codes. """ response = http.post( - f"{_BASE}/query", + f"{API_URL}/vision-events/query", json=query_params, headers=_auth_headers(api_key), ) @@ -99,7 +97,7 @@ def list_use_cases(api_key: str, status: Optional[str] = None) -> dict: if status is not None: params["status"] = status response = http.get( - f"{_BASE}/use-cases", + f"{API_URL}/vision-events/use-cases", params=params, headers=_auth_headers(api_key), ) @@ -122,7 +120,7 @@ def get_custom_metadata_schema(api_key: str, use_case_id: str) -> dict: RoboflowError: On non-200 response status codes. """ response = http.get( - f"{_BASE}/custom-metadata-schema/{use_case_id}", + f"{API_URL}/vision-events/custom-metadata-schema/{use_case_id}", headers=_auth_headers(api_key), ) if response.status_code != 200: @@ -144,7 +142,7 @@ def create_use_case(api_key: str, name: str) -> dict: RoboflowError: On non-201 response status codes. """ response = http.post( - f"{_BASE}/use-cases", + f"{API_URL}/vision-events/use-cases", json={"name": name}, headers=_auth_headers(api_key), ) @@ -168,7 +166,7 @@ def rename_use_case(api_key: str, use_case_id: str, name: str) -> dict: RoboflowError: On non-200 response status codes. """ response = http.put( - f"{_BASE}/use-cases/{use_case_id}", + f"{API_URL}/vision-events/use-cases/{use_case_id}", json={"name": name}, headers=_auth_headers(api_key), ) @@ -191,7 +189,7 @@ def archive_use_case(api_key: str, use_case_id: str) -> dict: RoboflowError: On non-200 response status codes. """ response = http.post( - f"{_BASE}/use-cases/{use_case_id}/archive", + f"{API_URL}/vision-events/use-cases/{use_case_id}/archive", headers=_auth_headers(api_key), ) if response.status_code != 200: @@ -213,7 +211,7 @@ def unarchive_use_case(api_key: str, use_case_id: str) -> dict: RoboflowError: On non-200 response status codes. """ response = http.post( - f"{_BASE}/use-cases/{use_case_id}/unarchive", + f"{API_URL}/vision-events/use-cases/{use_case_id}/unarchive", headers=_auth_headers(api_key), ) if response.status_code != 200: @@ -253,7 +251,7 @@ def upload_image( m = MultipartEncoder(fields=fields) headers = _auth_headers(api_key) headers["Content-Type"] = m.content_type - response = http.post(f"{_BASE}/upload", data=m, headers=headers) + response = http.post(f"{API_URL}/vision-events/upload", data=m, headers=headers) if response.status_code != 201: raise RoboflowError(response.text) diff --git a/roboflow/cli/_output.py b/roboflow/cli/_output.py index 70ba2c6c..638a1d2d 100644 --- a/roboflow/cli/_output.py +++ b/roboflow/cli/_output.py @@ -7,8 +7,8 @@ from __future__ import annotations import contextlib -import io import json +import os import sys from typing import Any, Iterator, Optional @@ -234,8 +234,8 @@ def output_api_error( ) -> None: """Render a server-side error and exit with the correct code. - Maps an exception's HTTP status (carried as ``exc.status_code`` when the - raiser sets it — see ``_raise_for_trash_response`` in ``adapters.rfapi``) + Maps an exception's HTTP status (carried as ``exc.status_code`` or + ``exc.response.status_code`` for requests HTTP errors) to the per-CONTRIBUTING.md exit-code contract: * **401** → exit code 2 ("auth error"). ``auth_hint`` overrides ``hint`` @@ -253,6 +253,8 @@ def output_api_error( whether to retry vs. re-auth. """ status = getattr(exc, "status_code", None) + if status is None: + status = getattr(getattr(exc, "response", None), "status_code", None) if status == 401: output_error( args, @@ -279,5 +281,5 @@ def suppress_sdk_output(args: Any = None) -> Iterator[None]: are not useful CLI output in any mode. The CLI controls its own output via ``output()`` and ``output_error()``. """ - with contextlib.redirect_stdout(io.StringIO()): + with open(os.devnull, "w", encoding="utf-8") as sink, contextlib.redirect_stdout(sink): yield diff --git a/roboflow/cli/handlers/image.py b/roboflow/cli/handlers/image.py index 3f3b2bf1..26509f10 100644 --- a/roboflow/cli/handlers/image.py +++ b/roboflow/cli/handlers/image.py @@ -356,10 +356,8 @@ def _handle_upload_single(args, api_key: str, path: str) -> None: # noqa: ANN00 def _handle_upload_directory(args, api_key: str, path: str) -> None: # noqa: ANN001 - import os - import roboflow - from roboflow.cli._output import output, output_error, suppress_sdk_output + from roboflow.cli._output import output, output_api_error, output_error, suppress_sdk_output # Always suppress SDK "loading..." noise during workspace init with suppress_sdk_output(): @@ -367,7 +365,7 @@ def _handle_upload_directory(args, api_key: str, path: str) -> None: # noqa: AN rf = roboflow.Roboflow(api_key) workspace = rf.workspace(args.workspace) except Exception as exc: - output_error(args, str(exc), exit_code=3) + output_api_error(args, exc, hint="Check your workspace and API key.") return retries = getattr(args, "retries", None) or getattr(args, "num_retries", 0) or 0 @@ -376,25 +374,30 @@ def _handle_upload_directory(args, api_key: str, path: str) -> None: # noqa: AN wait = not getattr(args, "no_wait", False) try: - result = workspace.upload_dataset( - dataset_path=path, - project_name=args.project, - num_workers=args.concurrency, - batch_name=getattr(args, "batch", None), - num_retries=retries, - is_prediction=getattr(args, "is_prediction", False), - use_zip_upload=getattr(args, "zip_upload", False), - annotation_overwrite=getattr(args, "annotation_overwrite", None), - split=getattr(args, "split", None), - tags=tags, - wait=wait, - ) + with suppress_sdk_output(): + result = workspace.upload_dataset( + dataset_path=path, + project_name=args.project, + num_workers=args.concurrency, + batch_name=getattr(args, "batch", None), + num_retries=retries, + is_prediction=getattr(args, "is_prediction", False), + use_zip_upload=getattr(args, "zip_upload", False), + annotation_overwrite=getattr(args, "annotation_overwrite", None), + split=getattr(args, "split", None), + tags=tags, + wait=wait, + return_summary=True, + ) except Exception as exc: - output_error(args, str(exc)) + output_api_error(args, exc, hint="Check the failed files and retry the upload.") return - if isinstance(result, dict): + if "uploaded" not in result: status = result.get("status", "unknown") + if status not in {"completed", "pending"}: + output_error(args, f"Zip upload {status}: {result}", hint="Check the upload task and retry.") + return data = { "status": status, "task_id": result.get("task_id") or result.get("taskId"), @@ -405,15 +408,8 @@ def _handle_upload_directory(args, api_key: str, path: str) -> None: # noqa: AN output(args, data, text=f"Imported {path} to {args.project} (zip upload, status={status})") return - # Per-image fallback — count files via image extensions - count = 0 - image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"} - for root, _dirs, files in os.walk(path): - for f in files: - if os.path.splitext(f)[1].lower() in image_exts: - count += 1 - - data = {"status": "imported", "path": path, "count": count} + count = result["uploaded"] + result["duplicates"] + data = {"status": "imported", "path": path, "count": count, **result} output(args, data, text=f"Imported {count} images from {path} to {args.project}") diff --git a/roboflow/cli/handlers/model.py b/roboflow/cli/handlers/model.py index 0eb29f77..f3500542 100644 --- a/roboflow/cli/handlers/model.py +++ b/roboflow/cli/handlers/model.py @@ -311,42 +311,38 @@ def _get_model(args): # noqa: ANN001 def _upload_model(args): # noqa: ANN001 import roboflow - from roboflow.cli._output import output, output_error + from roboflow.cli._output import output, output_api_error, output_error, suppress_sdk_output + from roboflow.config import load_roboflow_api_key - api_key = args.api_key or None - rf = roboflow.Roboflow(api_key=api_key) - workspace = rf.workspace(args.workspace) - - if args.version_number is not None: - # Deploy to a specific version - project_id = args.project[0] if isinstance(args.project, list) else args.project - if not project_id: - output_error(args, "Project is required for model upload.", hint="Use -p/--project.") - return + if not args.project: + output_error(args, "At least one project is required.", hint="Use -p/--project.") + return - try: - project = workspace.project(project_id) - version = project.version(args.version_number) - version.deploy(str(args.model_type), str(args.model_path), str(args.filename)) - except Exception as exc: - output_error(args, str(exc)) - return - else: - # Deploy to multiple projects - if not args.project: - output_error(args, "At least one project is required.", hint="Use -p/--project.") - return + api_key = args.api_key or load_roboflow_api_key(args.workspace) + if not api_key: + output_error(args, "No API key found", hint="Set ROBOFLOW_API_KEY or run 'roboflow auth login'", exit_code=2) + return - try: - workspace.deploy_model( - model_type=str(args.model_type), - model_path=str(args.model_path), - project_ids=args.project, - model_name=str(args.model_name) if args.model_name else "", - filename=str(args.filename), - ) - except Exception as exc: - output_error(args, str(exc)) - return + try: + with suppress_sdk_output(): + rf = roboflow.Roboflow(api_key=api_key) + workspace = rf.workspace(args.workspace) + if args.version_number is not None: + project_id = args.project[0] if isinstance(args.project, list) else args.project + project = workspace.project(project_id) + version = project.version(args.version_number) + version.deploy(str(args.model_type), str(args.model_path), str(args.filename), interactive=False) + else: + workspace.deploy_model( + model_type=str(args.model_type), + model_path=str(args.model_path), + project_ids=args.project, + model_name=str(args.model_name) if args.model_name else "", + filename=str(args.filename), + interactive=False, + ) + except Exception as exc: + output_api_error(args, exc, hint="Check the model files and project permissions, then retry.") + return output(args, {"status": "uploaded"}, text="Model uploaded successfully.") diff --git a/roboflow/core/version.py b/roboflow/core/version.py index e78c5b2f..df76d0b9 100644 --- a/roboflow/core/version.py +++ b/roboflow/core/version.py @@ -9,14 +9,12 @@ from typing import TYPE_CHECKING, Optional, Union from dotenv import load_dotenv -from tqdm import tqdm from roboflow.adapters import rfapi from roboflow.config import ( API_URL, APP_URL, DEMO_KEYS, - TQDM_DISABLE, TYPE_ACTION_RECOGNITION, TYPE_CLASSICATION, TYPE_INSTANCE_SEGMENTATION, @@ -35,8 +33,12 @@ from roboflow.models.vlm import VLMModel from roboflow.util import http from roboflow.util.annotations import amend_data_yaml -from roboflow.util.general import extract_zip, write_line -from roboflow.util.model_processor import package_custom_weights_interactive, validate_model_type_for_project +from roboflow.util.general import extract_zip, save_response_content, write_line +from roboflow.util.model_processor import ( + package_custom_weights, + package_custom_weights_interactive, + validate_model_type_for_project, +) from roboflow.util.train_recipe import fold_epochs_into_recipe from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions @@ -373,12 +375,17 @@ def download(self, model_format=None, location=None, overwrite: bool = False): location (str): An optional path for saving the file overwrite (bool): An optional flag to overwrite an existing dataset if the dataset has already downloaded + Only completed downloads of the same dataset version, format, and API endpoint + are reused. Interrupted SDK downloads can be retried; other non-empty directories + require overwrite=True or a different location. + Returns: Dataset Object Raises: RuntimeError: If the Roboflow API returns an error with a helpful JSON body HTTPError: If the Network/Roboflow API fails and does not return JSON + FileExistsError: If a non-empty location cannot be reused without overwrite=True """ # noqa: E501 // docs self.__wait_if_generating() @@ -388,12 +395,24 @@ def download(self, model_format=None, location=None, overwrite: bool = False): if model_format not in self.exports: self.export(model_format) - # if model_format is not in - if location is None: location = self.__get_download_location() - if os.path.exists(location) and not overwrite: - return Dataset(self.name, self.version, model_format, os.path.abspath(location)) + location = os.path.abspath(location) + marker = os.path.join(location, ".roboflow-download.json") + metadata = {"id": self.id, "version": self.version, "format": model_format, "api_url": API_URL} + if not overwrite: + try: + with open(marker) as f: + cached = json.load(f) + except (OSError, ValueError): + cached = None + if cached == {**metadata, "complete": True}: + return Dataset(self.name, self.version, model_format, location) + if os.path.exists(location) and os.listdir(location) and cached != {**metadata, "complete": False}: + raise FileExistsError( + f"Cannot reuse non-empty download location '{location}'. " + "Use overwrite=True to replace dataset files, or choose a different location." + ) if self.__api_key == "coco-128-sample": link = "https://app.roboflow.com/ds/n9QwXwUK42?key=NnVCe2yMxP" @@ -415,11 +434,17 @@ def download(self, model_format=None, location=None, overwrite: bool = False): link = export_info["export"]["link"] + os.makedirs(location, exist_ok=True) + # Record an incomplete download before changing dataset files so it can be retried. + with open(marker, "w") as f: + json.dump({**metadata, "complete": False}, f) self.__download_zip(link, location, model_format) extract_zip(location, desc=f"Extracting Dataset Version Zip to {location} in {model_format}:") self.__reformat_yaml(location, model_format) # TODO: is roboflow-python a place to be munging yaml files? + with open(marker, "w") as f: + json.dump({**metadata, "complete": True}, f) - return Dataset(self.name, self.version, model_format, os.path.abspath(location)) + return Dataset(self.name, self.version, model_format, location) def export(self, model_format=None) -> bool | None: """ @@ -557,14 +582,12 @@ def live_plot(epochs, mAP, loss, title=""): else: models = {} - if "train" in version.keys(): - if "results" in version["train"].keys(): - status = "finished" - break - if "status" in version["train"].keys(): - if version["train"]["status"] == "failed": - write_line(line="Training failed") - break + training = version.get("train") or {} + if training.get("status") == "failed": + raise RuntimeError(f"Training failed for {workspace}/{project}/{self.version}") + if "results" in training: + status = "finished" + break epoch_ids: Union[np.ndarray, list] mAP: Union[np.ndarray, list] @@ -655,15 +678,19 @@ def live_plot(epochs, mAP, loss, title=""): return self._model # @warn_for_wrong_dependencies_versions([("ultralytics", "==", "8.0.196")]) - def deploy(self, model_type: str, model_path: str, filename: str = "weights/best.pt") -> None: + def deploy( + self, model_type: str, model_path: str, filename: str = "weights/best.pt", *, interactive: bool = True + ) -> None: """Uploads provided weights file to Roboflow. Args: model_type (str): The type of the model to be deployed. model_path (str): File path to the model weights to be uploaded. filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". + interactive (bool): Ask before overriding packaging mismatches. If False, raise instead. """ - bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) + package = package_custom_weights_interactive if interactive else package_custom_weights + bundle = package(model_type, model_path, filename, build_dir=model_path) self._validate_against_project_type(bundle.model_type) self._upload_zip(bundle.model_type, model_path, bundle.archive_path.name) @@ -676,41 +703,19 @@ def _upload_zip(self, model_type: str, model_path: str, model_file_name: str): f"{API_URL}/{self.workspace}/{self.project}/{self.version}" f"/uploadModel?api_key={self.__api_key}&modelType={model_type}&nocache=true" ) - try: - if res.status_code == 429: - raise RuntimeError( - "This version already has a trained model. Please generate and" - " train a new version in order to upload model to Roboflow." - ) - else: - res.raise_for_status() - except Exception as e: - print(f"An error occured when getting the model upload URL: {e}") - return + res.raise_for_status() - res = http.put( - res.json()["url"], - data=open(os.path.join(model_path, model_file_name), "rb"), - ) - try: + with open(os.path.join(model_path, model_file_name), "rb") as model_file: + res = http.put(res.json()["url"], data=model_file) res.raise_for_status() - if self.public: - print( - f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}" - ) - print( - "Share your model with the world at:" - f" {UNIVERSE_URL}/{self.workspace}/{self.project}/" - f"model/{self.version}" - ) - else: - print( - f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}" - ) - - except Exception as e: - print(f"An error occured when uploading the model: {e}") + print(f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}") + if self.public: + print( + "Share your model with the world at:" + f" {UNIVERSE_URL}/{self.workspace}/{self.project}/" + f"model/{self.version}" + ) def __download_zip(self, link, location, format): """ @@ -724,33 +729,13 @@ def __download_zip(self, link, location, format): if not os.path.exists(location): os.makedirs(location) - def bar_progress(current, total, width=80): - progress_message = ( - f"Downloading Dataset Version Zip in {location} to {format}: " - f"{current / total * 100:.0f}% [{current} / {total}] bytes" + with http.get(link, stream=True) as response: + response.raise_for_status() + save_response_content( + response, + os.path.join(location, "roboflow.zip"), + desc=f"Downloading Dataset Version Zip in {location} to {format}:", ) - sys.stdout.write("\r" + progress_message) - sys.stdout.flush() - - try: - response = http.get(link, stream=True) - - # write the zip file to the desired location - with open(location + "/roboflow.zip", "wb") as f: - total_length = int(response.headers.get("content-length")) # type: ignore[arg-type] - desc = None if TQDM_DISABLE else f"Downloading Dataset Version Zip in {location} to {format}:" - for chunk in tqdm( - response.iter_content(chunk_size=1024), - desc=desc, - total=int(total_length / 1024) + 1, - ): - if chunk: - f.write(chunk) - f.flush() - - except Exception as e: - print(f"Error when trying to download dataset @ {link}") - raise e sys.stdout.write("\n") sys.stdout.flush() diff --git a/roboflow/core/workspace.py b/roboflow/core/workspace.py index 5392b071..e7473cd7 100644 --- a/roboflow/core/workspace.py +++ b/roboflow/core/workspace.py @@ -11,12 +11,12 @@ from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Union from requests.exceptions import HTTPError -from tqdm import tqdm from roboflow.adapters import rfapi, vision_events_api from roboflow.adapters.rfapi import AnnotationSaveError, ImageUploadError, RoboflowError from roboflow.config import API_URL, APP_URL, DEMO_KEYS from roboflow.util import http +from roboflow.util.general import save_response_content if TYPE_CHECKING: from roboflow.core.device import Device @@ -521,6 +521,7 @@ def upload_dataset( wait: bool = True, poll_interval: float = 5.0, poll_timeout: float = 3600.0, + return_summary: bool = False, ) -> Optional[dict]: """ Upload a dataset to Roboflow. @@ -547,9 +548,15 @@ def upload_dataset( wait (bool, optional): zip flow only — poll for processing completion. Defaults to True. poll_interval (float, optional): zip flow only — seconds between status polls. poll_timeout (float, optional): zip flow only — total seconds to wait before timing out. + return_summary (bool, optional): return uploaded and duplicate counts for per-image uploads. Returns: - dict | None: zip flow returns the final/pending status dict; per-image flow returns None. + dict | None: zip flow returns the final/pending status dict. Per-image flow returns + None unless return_summary=True, which returns uploaded and duplicate counts. + + Raises: + RoboflowError: if any image or annotation fails, after all images have been attempted, + or if zip processing fails. """ # noqa: E501 // docs if dataset_format != "NOT_USED": print("Warning: parameter 'dataset_format' is deprecated and will be removed in a future release") @@ -709,21 +716,36 @@ def _upload(imagedesc): image_id = image["id"] annotation, annotation_time = _save_annotation(image_id, imagedesc) _log_img_upload(image_path, image, annotation, image_upload_time, image_retry_attempts, annotation_time) + return ("duplicates" if image.get("duplicate") else "uploaded"), None except ImageUploadError as e: retry_attempts = f" (with {e.retries} retries)" if e.retries > 0 else "" print(f"[ERR]{retry_attempts} {image_path} ({e.message})") + return None, f"{image_path}: {e}" except AnnotationSaveError as e: upload_time_str = f"[{image_upload_time:.1f}s]" retry_attempts = f" (with {image_retry_attempts} retries)" if image_retry_attempts > 0 else "" image_msg = f"[UPLOADED]{retry_attempts} {image_path} ({image_id}) {upload_time_str}" annotation_msg = f"annotations = ERR: {e.message}" print(f"{image_msg} / {annotation_msg}") + return None, f"{image_path}: annotations: {e}" except Exception as e: print(f"[ERR] {image_path} ({e})") + return None, f"{image_path}: {e}" + counts = {"uploaded": 0, "duplicates": 0} + failures = [] with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor: - list(executor.map(_upload, images)) - + for status, error in executor.map(_upload, images): + if error is not None: + failures.append(error) + else: + counts[status] += 1 + if failures: + raise RoboflowError( + f"Dataset upload failed for {len(failures)} of {len(images)} images:\n" + "\n".join(failures) + ) + if return_summary: + return counts return None def _get_or_create_project(self, project_id, license: str = "MIT", type: str = "object-detection"): @@ -906,6 +928,8 @@ def deploy_model( project_ids: list[str], model_name: str, filename: str = "weights/best.pt", + *, + interactive: bool = True, ): """Uploads provided weights file to Roboflow. Args: @@ -913,9 +937,11 @@ def deploy_model( model_path (str): File path to the model weights to be uploaded. project_ids (list[str]): List of project IDs to deploy the model to. filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". + interactive (bool): Ask before overriding packaging mismatches. If False, raise instead. """ from roboflow.util.model_processor import ( + package_custom_weights, package_custom_weights_interactive, validate_model_type_for_project, ) @@ -930,7 +956,8 @@ def deploy_model( if project_id not in projects_by_id: raise ValueError(f"Project {project_id} is not accessible in this workspace") - bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) + package = package_custom_weights_interactive if interactive else package_custom_weights + bundle = package(model_type, model_path, filename, build_dir=model_path) for project_id in project_ids: validate_model_type_for_project(bundle.model_type, projects_by_id[project_id].get("type", ""), project_id) @@ -949,34 +976,18 @@ def _upload_zip( res = http.post( f"{API_URL}/{self.url}/models/prepareUpload?api_key={self.__api_key}&modelType={model_type}&modelName={model_name}&projectIds={','.join(project_ids)}&nocache=true" ) - try: - res.raise_for_status() - except Exception as e: - error_message = str(e) - status_code = str(res.status_code) - - print("\n\033[91m❌ ERROR\033[0m: Failed to get model deployment URL") - print("\033[93mDetails\033[0m:", error_message) - print("\033[93mStatus\033[0m:", status_code) - print(f"\033[93mResponse\033[0m:\n{res.text}\n") - return + res.raise_for_status() # Upload the model to the signed URL - res = http.put( - res.json()["url"], - data=open(os.path.join(model_path, model_file_name), "rb"), - ) - try: + with open(os.path.join(model_path, model_file_name), "rb") as model_file: + res = http.put(res.json()["url"], data=model_file) res.raise_for_status() - for project_id in project_ids: - print( - f"View the status of your deployment for project {project_id} at:" - f" {APP_URL}/{self.url}/{project_id}/models" - ) - - except Exception as e: - print(f"An error occured when uploading the model: {e}") + for project_id in project_ids: + print( + f"View the status of your deployment for project {project_id} at:" + f" {APP_URL}/{self.url}/{project_id}/models" + ) def search( self, @@ -1204,26 +1215,12 @@ def search_export( os.makedirs(location) zip_path = os.path.join(location, "roboflow.zip") - response = session.get(download_url, stream=True) - try: - response.raise_for_status() - except HTTPError as e: - raise RoboflowError(f"Failed to download search export: {e}") - - total_length = response.headers.get("content-length") - try: - total_kib = int(total_length) // 1024 + 1 if total_length is not None else None - except (TypeError, ValueError): - total_kib = None - with open(zip_path, "wb") as f: - for chunk in tqdm( - response.iter_content(chunk_size=1024), - desc=f"Downloading search export to {location}", - total=total_kib, - ): - if chunk: - f.write(chunk) - f.flush() + with session.get(download_url, stream=True) as response: + try: + response.raise_for_status() + except HTTPError as e: + raise RoboflowError(f"Failed to download search export: {e}") + save_response_content(response, zip_path, desc=f"Downloading search export to {location}") if extract_zip: _extract_zip(location, desc=f"Extracting search export to {location}") @@ -1822,7 +1819,9 @@ def _poll_zip_status( if progress is not None and progress != last_progress: print(f" zip-upload progress: {progress}") last_progress = progress - if state in {"completed", "failed"}: + if state == "failed": + raise RoboflowError(f"Zip upload failed (task_id={task_id}): {status.get('error') or status}") + if state == "completed": return status if time.monotonic() >= deadline: raise RoboflowError( diff --git a/roboflow/models/inference.py b/roboflow/models/inference.py index 13ac5521..d16aa96f 100644 --- a/roboflow/models/inference.py +++ b/roboflow/models/inference.py @@ -1,6 +1,7 @@ import io import json import os +import tempfile import time import urllib from typing import List, Optional, Tuple @@ -8,10 +9,10 @@ from PIL import Image from requests_toolbelt.multipart.encoder import MultipartEncoder -from tqdm import tqdm from roboflow.config import API_URL from roboflow.util import http +from roboflow.util.general import save_response_content from roboflow.util.image_utils import validate_image_path from roboflow.util.prediction import PredictionGroup @@ -230,19 +231,14 @@ def predict_video( signed_url_expires = signed_url.split("&X-Goog-Expires")[1].split("&")[0].strip("=") - # make a POST request to the signed URL + # Stream the video to the signed URL. headers = {"Content-Type": "application/octet-stream"} try: with open(video_path, "rb") as f: - video_data = f.read() + result = http.put(signed_url, data=f, headers=headers) except Exception as e: - raise Exception(f"Error reading video: {e}") - - try: - result = http.put(signed_url, data=video_data, headers=headers) - except Exception as e: - raise Exception(f"There was an error uploading the video: {e}") + raise Exception(f"There was an error uploading the video: {e}") from e if not result.ok: raise Exception(f"There was an error uploading the video: {result.text}") @@ -399,18 +395,12 @@ def download(self, format="pt", location="."): pt_weights_url = r.json()["weightsUrl"] - response = http.get(pt_weights_url, stream=True) - - # write the zip file to the desired location - with open(location + "/weights.pt", "wb") as f: - total_length = int(response.headers.get("content-length")) # type: ignore[arg-type] - for chunk in tqdm( - response.iter_content(chunk_size=1024), - desc=f"Downloading weights to {location}/weights.pt", - total=int(total_length / 1024) + 1, - ): - if chunk: - f.write(chunk) - f.flush() + with http.get(pt_weights_url, stream=True) as response: + response.raise_for_status() + # Keep existing weights intact until the entire download has succeeded. + with tempfile.TemporaryDirectory(dir=location) as temp_dir: + weights_path = os.path.join(temp_dir, "weights.pt") + save_response_content(response, weights_path, desc=f"Downloading weights to {location}/weights.pt") + os.replace(weights_path, os.path.join(location, "weights.pt")) return diff --git a/roboflow/models/object_detection.py b/roboflow/models/object_detection.py index bd56f11a..9dcd97f6 100644 --- a/roboflow/models/object_detection.py +++ b/roboflow/models/object_detection.py @@ -173,7 +173,6 @@ def predict( # type: ignore[override] self.__exception_check(image_path_check=image_path) original_dimensions = None - should_resize = False # If image is local image if not hosted: import cv2 @@ -248,21 +247,14 @@ def predict( # type: ignore[override] if self.format == "json": resp_json = resp.json() - if should_resize and original_dimensions is not None: - new_preds = [] + if original_dimensions is not None and dimensions != original_dimensions: + scale_x = original_dimensions[0] / dimensions[0] + scale_y = original_dimensions[1] / dimensions[1] for p in resp_json["predictions"]: - p["x"] = int(p["x"] * (int(original_dimensions[0]) / int(self.preprocessing["resize"]["width"]))) - p["y"] = int(p["y"] * (int(original_dimensions[1]) / int(self.preprocessing["resize"]["height"]))) - p["width"] = int( - p["width"] * (int(original_dimensions[0]) / int(self.preprocessing["resize"]["width"])) - ) - p["height"] = int( - p["height"] * (int(original_dimensions[1]) / int(self.preprocessing["resize"]["height"])) - ) - - new_preds.append(p) - - resp_json["predictions"] = new_preds + p["x"] = int(p["x"] * scale_x) + p["y"] = int(p["y"] * scale_y) + p["width"] = int(p["width"] * scale_x) + p["height"] = int(p["height"] * scale_y) return PredictionGroup.create_prediction_group( resp_json, diff --git a/roboflow/util/folderparser.py b/roboflow/util/folderparser.py index 047cdaee..3b153123 100644 --- a/roboflow/util/folderparser.py +++ b/roboflow/util/folderparser.py @@ -1,5 +1,6 @@ import json import os +import posixpath import re from collections import defaultdict @@ -112,7 +113,7 @@ def _map_annotations_to_images_1to1(images, annotations): def _map_annotations_to_images_1tomany(images, annotationFiles): image_path_to_annotation_files = _build_image_to_annotationfile_index(annotationFiles) - imgRefMap, annotationMap = _build_image_and_annotation_maps(annotationFiles) + imgRefMap, annotationMap = _build_image_and_annotation_maps(annotationFiles, images) for image in tqdm(images): # Get candidate annotation files for this image @@ -123,6 +124,8 @@ def _map_annotations_to_images_1tomany(images, annotationFiles): or image_path_to_annotation_files.get(image["key"], []) or annotationFiles # Fallback to all files for non-COCO formats ) + if image["file"] in imgRefMap: + candidate_annotations = [imgRefMap[image["file"]][0]] for annotationFile in candidate_annotations: format = annotationFile["parsedType"] @@ -135,8 +138,8 @@ def _map_annotations_to_images_1tomany(images, annotationFiles): def _build_image_to_annotationfile_index(annotationFiles): """Create an index mapping possible image path keys to annotation files that reference them. - Keys include full relative path, basename, and stem to improve robustness across - different dataset layouts. Supports coco, createml, csv, multilabel_csv, jsonl. + COCO references are resolved separately against the actual image paths. + Supports createml, csv, multilabel_csv, jsonl. """ index = defaultdict(list) for annotationFile in annotationFiles: @@ -145,18 +148,7 @@ def _build_image_to_annotationfile_index(annotationFiles): if not parsedType or parsed is None: continue - if parsedType == "coco": - for imageRef in parsed.get("images", []): - file_name = _patch_sep(imageRef.get("file_name", "")).lstrip("/") - if not file_name: - continue - basename = os.path.basename(file_name) - stem = os.path.splitext(basename)[0] - index[file_name].append(annotationFile) - index[basename].append(annotationFile) - index[stem].append(annotationFile) - - elif parsedType == "createml": + if parsedType == "createml": for entry in parsed: image_name = entry.get("image") if not image_name: @@ -187,9 +179,16 @@ def _build_image_to_annotationfile_index(annotationFiles): return index -def _build_image_and_annotation_maps(annotationFiles): +def _build_image_and_annotation_maps(annotationFiles, images): imgRefMap = {} annotationMap = defaultdict(list) + paths = {image["file"]: image for image in images} + names = defaultdict(list) + stems = defaultdict(list) + for image in images: + names[image["name"]].append(image) + stems[image["key"]].append(image) + for annFile in annotationFiles: filename, parsed, parsedType = ( annFile["file"], @@ -198,19 +197,26 @@ def _build_image_and_annotation_maps(annotationFiles): ) if parsedType == "coco": for imageRef in parsed["images"]: - # Normalize and index by multiple forms to improve matching robustness - file_name = _patch_sep(imageRef["file_name"]).lstrip("/") - basename = os.path.basename(file_name) - stem = os.path.splitext(basename)[0] - - # Prefer full relative path, but also allow basename and stem - imgRefMap.update( - { - f"{filename}/{file_name}": imageRef, - f"{filename}/{basename}": imageRef, - f"{filename}/{stem}": imageRef, - } + file_name = _patch_sep(imageRef["file_name"]) + if not file_name: + continue + # Prefer paths relative to this COCO file, then dataset-root paths. + image = paths.get(posixpath.normpath(posixpath.join(annFile["dirname"], file_name))) or paths.get( + posixpath.normpath("/" + file_name.lstrip("/")) ) + if image is None: + basename = posixpath.basename(file_name) + matches = names.get(basename, []) or stems.get(posixpath.splitext(basename)[0].lower(), []) + if len(matches) > 1: + raise ValueError( + f"Ambiguous COCO image '{file_name}' in '{filename}': use a relative image path" + ) + image = matches[0] if matches else None + if image is not None: + key = image["file"] + if key in imgRefMap: + raise ValueError(f"Ambiguous COCO annotations for '{key}': multiple image references match") + imgRefMap[key] = (annFile, imageRef) for annotation in parsed["annotations"]: annotationMap[f"{filename}/{annotation['image_id']}"].append(annotation) return imgRefMap, annotationMap @@ -219,15 +225,8 @@ def _build_image_and_annotation_maps(annotationFiles): def _filterIndividualAnnotations(image, annotation, format, imgRefMap, annotationMap): parsed = annotation["parsed"] if format == "coco": - rel_path = image["file"].lstrip("/") - imgReference = ( - # Try matching by full relative path first - imgRefMap.get(f"{annotation['file']}/{rel_path}") - # Fallback: basename with extension - or imgRefMap.get(f"{annotation['file']}/{image['name']}") - # Fallback: stem (no extension) - or imgRefMap.get(f"{annotation['file']}/{image['key']}") - ) + match = imgRefMap.get(image["file"]) + imgReference = match[1] if match and match[0] is annotation else None if imgReference: # workaround to make Annotations.js correctly identify this as coco in the backend fake_annotation = { diff --git a/roboflow/util/general.py b/roboflow/util/general.py index 9368d7a2..fe9a8b8e 100644 --- a/roboflow/util/general.py +++ b/roboflow/util/general.py @@ -4,6 +4,7 @@ import zipfile from random import random +from requests import Response from tqdm import tqdm from roboflow.config import TQDM_DISABLE @@ -48,6 +49,39 @@ def __call__(self, func, *args, **kwargs): raise +def save_response_content(response: Response, path: str, desc: str | None = None) -> None: + """Write an HTTP response in bounded chunks with byte-based progress. + + Args: + response: Streaming response whose HTTP status has already been checked. + path: Destination file path. + desc: Description shown in the progress bar. + """ + total: int | None + try: + total = int(response.headers.get("content-length", "")) + if total < 0: + total = None + except (TypeError, ValueError): + total = None + # iter_content decodes compressed responses; Content-Length counts wire bytes. + if ( + response.headers.get("content-encoding", "identity").lower() != "identity" + or "transfer-encoding" in response.headers + ): + total = None + # urllib3 1.x does not enforce Content-Length by default, including for gzip. + response.raw.enforce_content_length = True + with ( + open(path, "wb") as f, + tqdm(total=total, desc=desc, unit="B", unit_scale=True, disable=bool(TQDM_DISABLE)) as progress, + ): + for chunk in response.iter_content(chunk_size=64 * 1024): + if chunk: + f.write(chunk) + progress.update(len(chunk)) + + def extract_zip(location: str, desc: str = "Extracting"): """Extract ``roboflow.zip`` inside *location* and remove the archive. diff --git a/tests/cli/test_image_handler.py b/tests/cli/test_image_handler.py index 73009ec4..90d16a7e 100644 --- a/tests/cli/test_image_handler.py +++ b/tests/cli/test_image_handler.py @@ -144,6 +144,44 @@ def test_upload_single_json_mode(self, mock_rf_cls): class TestImageUploadDirectory(unittest.TestCase): """Test the directory import path.""" + @patch("roboflow.Roboflow") + def test_api_errors_follow_exit_code_contract(self, mock_rf_cls): + from roboflow.adapters.rfapi import RoboflowError + + workspace = mock_rf_cls.return_value.workspace + with tempfile.TemporaryDirectory() as directory: + for operation in (workspace, workspace.return_value.upload_dataset): + for status, code in ((401, 2), (404, 3), (500, 1)): + with self.subTest(operation=operation, status=status): + operation.side_effect = RoboflowError("Upload rejected", status_code=status) + result = runner.invoke( + app, ["--json", "--api-key", "test-key", "image", "upload", directory, "-p", "proj"] + ) + operation.side_effect = None + self.assertEqual(result.exit_code, code) + self.assertEqual(result.stdout, "") + error = json.loads(result.stderr)["error"] + self.assertEqual(error["message"], "Upload rejected") + self.assertTrue(error["hint"]) + + @patch("roboflow.Roboflow") + def test_failed_upload_has_no_success_output(self, mock_rf_cls): + from roboflow.cli.handlers.image import _handle_upload_directory + + upload = mock_rf_cls.return_value.workspace.return_value.upload_dataset + args = _make_args(json=True, project="proj", concurrency=2) + for failure in (RuntimeError("Annotation rejected"), {"status": "failed", "error": "Invalid zip"}): + with self.subTest(failure=failure): + upload.side_effect = failure if isinstance(failure, Exception) else None + upload.return_value = failure + stdout, stderr = io.StringIO(), io.StringIO() + with patch("sys.stdout", stdout), patch("sys.stderr", stderr): + with self.assertRaises(SystemExit) as raised: + _handle_upload_directory(args, "key", "/dataset") + self.assertEqual(raised.exception.code, 1) + self.assertEqual(stdout.getvalue(), "") + self.assertIn("error", json.loads(stderr.getvalue())) + @patch("roboflow.Roboflow") def test_upload_directory(self, mock_rf_cls): from roboflow.cli.handlers.image import _handle_upload @@ -155,6 +193,7 @@ def test_upload_directory(self, mock_rf_cls): f.write("x") mock_ws = MagicMock() + mock_ws.upload_dataset.return_value = {"uploaded": 1, "duplicates": 0} mock_rf_cls.return_value.workspace.return_value = mock_ws args = _make_args( @@ -183,7 +222,7 @@ def test_upload_directory(self, mock_rf_cls): mock_ws.upload_dataset.assert_called_once() result = json.loads(buf.getvalue()) self.assertEqual(result["status"], "imported") - self.assertEqual(result["count"], 2) # .jpg and .png only + self.assertEqual(result["count"], 1) # Count the SDK results, not local files. @patch("roboflow.Roboflow") def test_upload_zip_file_routes_to_directory_handler(self, mock_rf_cls): @@ -315,8 +354,7 @@ def test_zip_upload_flag_defaults_false(self, mock_rf_cls): with tempfile.TemporaryDirectory() as tmpdir: mock_ws = MagicMock() - # MagicMock return → not a dict → per-image output branch - mock_ws.upload_dataset.return_value = None + mock_ws.upload_dataset.return_value = {"uploaded": 0, "duplicates": 0} mock_rf_cls.return_value.workspace.return_value = mock_ws args = _make_args( @@ -407,6 +445,7 @@ def test_upload_directory_omits_default_split_when_not_explicit(self, mock_rf_cl with tempfile.TemporaryDirectory() as tmpdir: mock_ws = MagicMock() + mock_ws.upload_dataset.return_value = {"uploaded": 0, "duplicates": 0} mock_rf_cls.return_value.workspace.return_value = mock_ws args = _make_args( @@ -441,6 +480,7 @@ def test_upload_directory_forwards_explicit_split(self, mock_rf_cls): with tempfile.TemporaryDirectory() as tmpdir: mock_ws = MagicMock() + mock_ws.upload_dataset.return_value = {"uploaded": 0, "duplicates": 0} mock_rf_cls.return_value.workspace.return_value = mock_ws args = _make_args( diff --git a/tests/cli/test_model_handler.py b/tests/cli/test_model_handler.py index f4ca2f2d..d6ff22a1 100644 --- a/tests/cli/test_model_handler.py +++ b/tests/cli/test_model_handler.py @@ -3,6 +3,7 @@ import io import json import sys +import tempfile import types import unittest from unittest.mock import MagicMock, patch @@ -100,6 +101,79 @@ def test_get_model_no_api_key(self, _mock_key: MagicMock) -> None: class TestModelUpload(unittest.TestCase): """Test _upload_model handler.""" + @patch("roboflow.config.load_roboflow_api_key", return_value=None) + def test_missing_api_key_exits_two(self, _mock_key): + from roboflow.cli.handlers.model import _upload_model + + with patch("sys.stderr", io.StringIO()), self.assertRaises(SystemExit) as raised: + _upload_model(self._make_args(api_key=None, json=True)) + self.assertEqual(raised.exception.code, 2) + + def test_upload_initialization_preserves_api_status(self): + import responses + + import roboflow + from roboflow.config import API_URL + + with responses.RequestsMock() as http: + http.post(API_URL + "/", body="Invalid key", status=401) + with self.assertRaisesRegex(RuntimeError, "Invalid key") as raised: + roboflow.Roboflow(api_key="test-key") + self.assertEqual(raised.exception.status_code, 401) + + with tempfile.TemporaryDirectory() as directory: + for command in ( + ["image", "upload", directory, "-p", "proj"], + ["model", "upload", "-t", "yolov8", "-m", directory, "-p", "proj"], + ): + for status, code in ((401, 2), (404, 3), (500, 1)): + with self.subTest(command=command, status=status), responses.RequestsMock() as http: + if status == 401: + http.post(API_URL + "/", body="Invalid key", status=status) + else: + http.post(API_URL + "/", json={"workspace": "test-ws"}) + http.get(API_URL + "/test-ws", body="Workspace unavailable", status=status) + result = runner.invoke(app, ["--json", "--api-key", "test-key", *command]) + self.assertEqual(result.exit_code, code) + self.assertEqual(result.stdout, "") + self.assertIn("hint", json.loads(result.stderr)["error"]) + + @patch("roboflow.Roboflow") + def test_http_errors_follow_exit_code_contract(self, mock_rf_cls): + from requests import Response + from requests.exceptions import HTTPError + + workspace = mock_rf_cls.return_value.workspace.return_value + for version in ([], ["--version", "1"]): + deploy = workspace.project.return_value.version.return_value.deploy if version else workspace.deploy_model + for status, code in ((401, 2), (404, 3), (500, 1)): + with self.subTest(version=version, status=status): + response = Response() + response.status_code = status + deploy.side_effect = HTTPError("Upload rejected", response=response) + result = runner.invoke( + app, + [ + "--json", + "--api-key", + "test-key", + "model", + "upload", + "-t", + "yolov8", + "-m", + "/model", + "-p", + "proj", + *version, + ], + ) + self.assertEqual(result.exit_code, code) + self.assertEqual(result.stdout, "") + error = json.loads(result.stderr)["error"] + self.assertEqual(error["message"], "Upload rejected") + self.assertTrue(error["hint"]) + def _make_args(self, **kwargs: object) -> types.SimpleNamespace: defaults = { "json": False, @@ -115,6 +189,56 @@ def _make_args(self, **kwargs: object) -> types.SimpleNamespace: defaults.update(kwargs) return types.SimpleNamespace(**defaults) + @patch("roboflow.Roboflow") + def test_packaging_mismatch_does_not_prompt_in_cli(self, mock_rf_cls): + from roboflow.cli.handlers.model import _upload_model + from roboflow.util.model_processor import SizeMismatchError + from tests.helpers import get_version + from tests.test_workspace import _make_workspace + + workspace = _make_workspace() + workspace.project_list = [{"id": "test-ws/proj1", "type": "object-detection"}] + mock_rf_cls.return_value.workspace.return_value = workspace + with tempfile.TemporaryDirectory() as directory, patch.object(workspace, "project") as project: + project.return_value.version.return_value = get_version() + for version_number in (1, None): + with self.subTest(version_number=version_number): + stdout, stderr = io.StringIO(), io.StringIO() + error = SizeMismatchError("Model size does not match", requested="yolov8m", detected="yolov8n") + with ( + patch("roboflow.util.model_processor._process_model", side_effect=error), + patch("builtins.input", side_effect=AssertionError("Unexpected prompt")) as prompt, + patch("sys.stdout", stdout), + patch("sys.stderr", stderr), + self.assertRaises(SystemExit) as raised, + ): + _upload_model(self._make_args(json=True, model_path=directory, version_number=version_number)) + self.assertEqual(raised.exception.code, 1) + self.assertEqual(stdout.getvalue(), "") + self.assertIn("Model size does not match", json.loads(stderr.getvalue())["error"]["message"]) + prompt.assert_not_called() + + @patch("roboflow.Roboflow") + def test_failed_upload_has_no_success_output(self, mock_rf_cls): + from roboflow.cli.handlers.model import _upload_model + + workspace = mock_rf_cls.return_value.workspace.return_value + for version_number in (1, None): + with self.subTest(version_number=version_number): + deploy = ( + workspace.project.return_value.version.return_value.deploy + if version_number + else workspace.deploy_model + ) + deploy.side_effect = RuntimeError("Upload rejected") + stdout, stderr = io.StringIO(), io.StringIO() + with patch("sys.stdout", stdout), patch("sys.stderr", stderr): + with self.assertRaises(SystemExit) as raised: + _upload_model(self._make_args(json=True, version_number=version_number)) + self.assertEqual(raised.exception.code, 1) + self.assertEqual(stdout.getvalue(), "") + self.assertIn("Upload rejected", json.loads(stderr.getvalue())["error"]["message"]) + @patch("roboflow.Roboflow") def test_upload_single_version(self, mock_rf_cls: MagicMock) -> None: from roboflow.cli.handlers.model import _upload_model @@ -122,6 +246,7 @@ def test_upload_single_version(self, mock_rf_cls: MagicMock) -> None: mock_version = MagicMock() mock_project = MagicMock() mock_project.version.return_value = mock_version + mock_version.deploy.side_effect = lambda *args, **kwargs: print("SDK deployment output") mock_workspace = MagicMock() mock_workspace.project.return_value = mock_project mock_rf = MagicMock() @@ -139,7 +264,7 @@ def test_upload_single_version(self, mock_rf_cls: MagicMock) -> None: result = json.loads(buf.getvalue()) self.assertEqual(result["status"], "uploaded") - mock_version.deploy.assert_called_once_with("yolov8", "/path/to/model", "weights/best.pt") + mock_version.deploy.assert_called_once_with("yolov8", "/path/to/model", "weights/best.pt", interactive=False) @patch("roboflow.Roboflow") def test_upload_multi_project(self, mock_rf_cls: MagicMock) -> None: diff --git a/tests/cli/test_train_handler.py b/tests/cli/test_train_handler.py index 2c35a8e1..925897d1 100644 --- a/tests/cli/test_train_handler.py +++ b/tests/cli/test_train_handler.py @@ -77,8 +77,9 @@ def test_start_success(self, mock_train: MagicMock) -> None: self.assertEqual(result["project"], "my-project") self.assertEqual(result["version"], 1) + @patch("roboflow.adapters.rfapi.get_version", return_value={"version": {"exports": ["yolov5pytorch"]}}) @patch("roboflow.adapters.rfapi.start_version_training") - def test_start_with_all_options(self, mock_train: MagicMock) -> None: + def test_start_with_all_options(self, mock_train: MagicMock, mock_version: MagicMock) -> None: from roboflow.cli.handlers.train import _start mock_train.return_value = True @@ -98,6 +99,7 @@ def test_start_with_all_options(self, mock_train: MagicMock) -> None: finally: sys.stdout = old_stdout + mock_version.assert_called_once_with("test-key", "test-ws", "my-project", "1") mock_train.assert_called_once_with( "test-key", "test-ws", diff --git a/tests/models/test_inference.py b/tests/models/test_inference.py new file mode 100644 index 00000000..eccc97f9 --- /dev/null +++ b/tests/models/test_inference.py @@ -0,0 +1,131 @@ +import gzip +import io +import os +import tempfile +import unittest +from unittest.mock import MagicMock, patch + +import requests +import responses +from urllib3.response import HTTPResponse + +from roboflow.config import API_URL +from roboflow.models.inference import InferenceModel +from roboflow.models.object_detection import ObjectDetectionModel + + +class TestVideoUpload(unittest.TestCase): + @responses.activate + def test_stream_is_open_during_upload_and_closed_afterwards(self): + model = ObjectDetectionModel("test-key", "workspace/project/1") + signed_url = "https://upload.test/video?signature=test&X-Goog-Expires=3600" + responses.post(f"{API_URL}/video_upload_signed_url", json={"signed_url": signed_url}) + responses.post(f"{API_URL}/videoinfer/", json={"job_id": "job-1"}) + with tempfile.TemporaryDirectory() as directory: + video_path = os.path.join(directory, "video.mp4") + with open(video_path, "wb") as video: + video.write(b"video bytes") + for outcome in (200, 403, requests.ConnectionError("Connection lost")): + with self.subTest(outcome=outcome): + + def upload(url, data, headers): + self.assertEqual(url, signed_url) + self.assertFalse(data.closed) + self.assertEqual(data.read(5), b"video") + self.assertEqual(data.read(6), b" bytes") + if isinstance(outcome, Exception): + raise outcome + response = requests.Response() + response.status_code = outcome + response._content = b"Upload rejected" if outcome == 403 else b"" + return response + + before = len(responses.calls) + with patch("roboflow.models.inference.http.put", side_effect=upload) as put: + if outcome == 200: + self.assertEqual(model.predict_video(video_path), ("job-1", signed_url, "3600")) + self.assertEqual(len(responses.calls) - before, 2) + else: + with self.assertRaisesRegex(Exception, "error uploading the video"): + model.predict_video(video_path) + self.assertEqual(len(responses.calls) - before, 1) + self.assertTrue(put.call_args.kwargs["data"].closed) + + @responses.activate + def test_hosted_video_does_not_upload_a_file(self): + model = ObjectDetectionModel("test-key", "workspace/project/1") + responses.post(f"{API_URL}/videoinfer/", json={"job_id": "job-1"}) + with patch("roboflow.models.inference.http.put") as put: + result = model.predict_video("https://video.test/video.mp4") + self.assertEqual(result, ("job-1", "https://video.test/video.mp4", None)) + put.assert_not_called() + + +class TestDownloadWeights(unittest.TestCase): + def setUp(self): + self.model = InferenceModel("test-key", "workspace/project/1") + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.weights_path = os.path.join(self.directory.name, "weights.pt") + with open(self.weights_path, "wb") as f: + f.write(b"existing weights") + self.metadata_url = f"{API_URL}/workspace/project/1/ptFile" + self.weights_url = "https://download.test/weights.pt" + + @responses.activate + def test_http_error_does_not_replace_existing_weights(self): + responses.get(self.metadata_url, json={"weightsUrl": self.weights_url}) + for status in (403, 404, 500): + with self.subTest(status=status): + responses.get(self.weights_url, status=status, body="error") + with self.assertRaises(requests.HTTPError): + self.model.download(location=self.directory.name) + with open(self.weights_path, "rb") as f: + self.assertEqual(f.read(), b"existing weights") + self.assertEqual(os.listdir(self.directory.name), ["weights.pt"]) + + def test_interrupted_stream_preserves_weights_and_cleans_temporary_files(self): + def broken_stream(**kwargs): + yield b"partial weights" + raise requests.exceptions.ChunkedEncodingError("Download interrupted") + + metadata = MagicMock() + metadata.json.return_value = {"weightsUrl": self.weights_url} + response = MagicMock(headers={"content-length": "1024"}) + response.__enter__.return_value = response + response.iter_content.side_effect = broken_stream + with patch("roboflow.models.inference.http.get", side_effect=[metadata, response]): + with self.assertRaisesRegex(requests.exceptions.ChunkedEncodingError, "interrupted"): + self.model.download(location=self.directory.name) + with open(self.weights_path, "rb") as f: + self.assertEqual(f.read(), b"existing weights") + self.assertEqual(os.listdir(self.directory.name), ["weights.pt"]) + response.__exit__.assert_called_once() + + def test_short_stream_without_transport_error_does_not_replace_weights(self): + metadata = MagicMock() + metadata.json.return_value = {"weightsUrl": self.weights_url} + for encoding, body in (("identity", b"truncated"), ("gzip", gzip.compress(b"partial weights"))): + with self.subTest(encoding=encoding): + response = requests.Response() + response.status_code = 200 + response.headers.update({"content-length": str(len(body) + 10), "content-encoding": encoding}) + response.raw = HTTPResponse( + body=io.BytesIO(body), headers=response.headers, preload_content=False, enforce_content_length=False + ) + with patch("roboflow.models.inference.http.get", side_effect=[metadata, response]): + with self.assertRaises(requests.exceptions.ChunkedEncodingError): + self.model.download(location=self.directory.name) + with open(self.weights_path, "rb") as f: + self.assertEqual(f.read(), b"existing weights") + self.assertEqual(os.listdir(self.directory.name), ["weights.pt"]) + + @responses.activate + def test_success_replaces_weights_with_exact_response_bytes(self): + weights = b"new weights" * 1000 + responses.get(self.metadata_url, json={"weightsUrl": self.weights_url}) + responses.get(self.weights_url, body=weights, headers={"Content-Length": str(len(weights))}) + self.model.download(location=self.directory.name) + with open(self.weights_path, "rb") as f: + self.assertEqual(f.read(), weights) + self.assertEqual(os.listdir(self.directory.name), ["weights.pt"]) diff --git a/tests/models/test_object_detection.py b/tests/models/test_object_detection.py index 772d8b00..a5c2bd74 100644 --- a/tests/models/test_object_detection.py +++ b/tests/models/test_object_detection.py @@ -1,7 +1,11 @@ +import base64 +import io +import os +import tempfile import unittest import responses -from PIL import UnidentifiedImageError +from PIL import Image, UnidentifiedImageError from requests.exceptions import HTTPError from roboflow.config import OBJECT_DETECTION_URL @@ -98,6 +102,38 @@ def test_predict_with_a_numpy_array_request(self): self.assertDictEqual(request.params, self._default_params) self.assertIsNotNone(request.body) + @responses.activate + def test_predict_scales_boxes_only_when_the_sent_image_is_resized(self): + prediction = {"x": 100.5, "y": 50.25, "width": 40.5, "height": 20.25, "class": "helmet", "confidence": 0.9} + unchanged = (100.5, 50.25, 40.5, 20.25) + cases = [ + ((200, 100), True, (200, 100), unchanged), + ((640, 640), True, (640, 640), unchanged), + ((1280, 640), True, (640, 640), (201, 50, 81, 20)), + ((640, 1280), True, (640, 640), (100, 100, 40, 40)), + ((1280, 320), True, (640, 640), (201, 25, 81, 10)), + ((1280, 1280), True, (640, 640), (201, 100, 81, 40)), + ((1280, 1280), False, (1280, 1280), unchanged), + ] + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "input.png") + for size, resize, sent_size, expected in cases: + with self.subTest(size=size, resize=resize): + Image.new("RGB", size).save(path) + preprocessing = ( + {"resize": {"format": "Stretch to", "width": "640", "height": "640"}} if resize else {} + ) + model = ObjectDetectionModel(self.api_key, self.version_id, preprocessing=preprocessing) + responses.add(responses.POST, self.api_url, json={"predictions": [prediction]}) + + result = model.predict(path).json() + + sent = Image.open(io.BytesIO(base64.b64decode(responses.calls[-1].request.body))) + self.assertEqual(sent.size, sent_size) + box = result["predictions"][0] + self.assertEqual(tuple(box[key] for key in ("x", "y", "width", "height")), expected) + self.assertEqual(result["image"], {"width": str(size[0]), "height": str(size[1])}) + def test_predict_with_local_wrong_image_request(self): image_path = "tests/images/not_an_image.txt" instance = ObjectDetectionModel(self.api_key, self.version_id, version=self.version) diff --git a/tests/test_model_upload.py b/tests/test_model_upload.py new file mode 100644 index 00000000..ba25f862 --- /dev/null +++ b/tests/test_model_upload.py @@ -0,0 +1,59 @@ +"""Model upload failures must reach SDK callers and close the upload file.""" + +import os +import tempfile +import unittest +from functools import partial +from unittest.mock import patch + +import requests +import responses + +from roboflow.config import API_URL +from roboflow.util import http +from tests.helpers import get_version +from tests.test_workspace import _make_workspace + + +class TestModelUpload(unittest.TestCase): + @responses.activate + def test_upload_status_and_file_lifetime(self): + version = get_version(workspace_name="test-ws", project_name="project") + workspace = _make_workspace() + with tempfile.TemporaryDirectory() as directory: + with open(os.path.join(directory, "model.zip"), "wb") as model_file: + model_file.write(b"model archive") + for scope in ("version", "workspace"): + for prepare_status, upload_status in ((403, 200), (429, 200), (200, 500), (200, 200)): + with self.subTest(scope=scope, prepare=prepare_status, upload=upload_status): + responses.reset() + if scope == "version": + responses.get( + f"{API_URL}/test-ws/project/1/uploadModel", + status=prepare_status, + json={"url": "https://upload.test/model"}, + ) + upload = partial(version._upload_zip, "yolov8", directory, "model.zip") + else: + responses.post( + f"{API_URL}/test-ws/models/prepareUpload", + status=prepare_status, + json={"url": "https://upload.test/model"}, + ) + upload = partial( + workspace._upload_zip, "yolov8", directory, ["project"], "model", "model.zip" + ) + + if prepare_status == 200: + responses.put("https://upload.test/model", status=upload_status) + with patch("roboflow.util.http.put", wraps=http.put) as put: + if prepare_status != 200 or upload_status != 200: + with self.assertRaises(requests.HTTPError): + upload() + else: + upload() + if prepare_status == 200: + self.assertTrue(put.call_args.kwargs["data"].closed) + self.assertEqual(responses.calls[-1].request.body, b"model archive") + else: + self.assertEqual(len(responses.calls), 1) diff --git a/tests/test_project.py b/tests/test_project.py index 59bb14bb..8292bbc9 100644 --- a/tests/test_project.py +++ b/tests/test_project.py @@ -588,7 +588,8 @@ def test_project_upload_dataset(self): try: params = {"dataset_path": "/test/dataset", "project_name": PROJECT_NAME} params.update(case.get("params", {})) - self.workspace.upload_dataset(**params) + with self.assertRaisesRegex(RoboflowError, "failed for 1 of 1 images"): + self.workspace.upload_dataset(**params) finally: for mock in mocks.values(): mock.stop() diff --git a/tests/test_version.py b/tests/test_version.py index d00ae4cc..e149cf8b 100644 --- a/tests/test_version.py +++ b/tests/test_version.py @@ -1,6 +1,9 @@ +import io import json import os +import tempfile import unittest +import zipfile from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -65,11 +68,167 @@ def test_download_raises_exception_on_api_failure(self): def test_download_returns_dataset(self, *_): responses.add(responses.GET, self.api_url, json={"export": {"link": None}}) mock_generating_url_response(self.generating_url) - dataset = self.version.download("coco", location="/my-spot") - self.assertEqual(dataset.name, self.version.name) - self.assertEqual(dataset.version, self.version.version) - self.assertEqual(dataset.model_format, "coco") - self.assertEqual(dataset.location, os.path.abspath("/my-spot")) + with tempfile.TemporaryDirectory() as location: + dataset = self.version.download("coco", location=location) + self.assertEqual(dataset.name, self.version.name) + self.assertEqual(dataset.version, self.version.version) + self.assertEqual(dataset.model_format, "coco") + self.assertEqual(dataset.location, os.path.abspath(location)) + + +class TestDownloadCache(unittest.TestCase): + def setUp(self): + self.version = get_version(exports=["coco", "voc"]) + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.location = self.directory.name + self.marker = os.path.join(self.location, ".roboflow-download.json") + self.zip_url = "https://download.test/dataset.zip" + for target, kwargs in ( + ("roboflow.core.version.Version._Version__wait_if_generating", {}), + ("roboflow.adapters.rfapi.get_version_export", {"return_value": {"export": {"link": self.zip_url}}}), + ): + mock = patch(target, **kwargs) + mock.start() + self.addCleanup(mock.stop) + archive = io.BytesIO() + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("train/image.jpg", b"image bytes") + self.archive = archive.getvalue() + + def _serve_zip(self, body=None, status=200): + """Serve a dataset archive without contacting Roboflow.""" + body = self.archive if body is None else body + responses.get(self.zip_url, body=body, status=status, headers={"Content-Length": str(len(body))}) + + @responses.activate + def test_unrecognized_directory_requires_explicit_overwrite(self): + image_path = os.path.join(self.location, "train", "image.jpg") + os.makedirs(os.path.dirname(image_path)) + with open(image_path, "wb") as f: + f.write(b"user data") + self._serve_zip() + with self.assertRaisesRegex(FileExistsError, "overwrite=True"): + self.version.download("coco", self.location) + with open(image_path, "rb") as f: + self.assertEqual(f.read(), b"user data") + self.assertEqual(len(responses.calls), 0) + self.version.download("coco", self.location, overwrite=True) + with open(image_path, "rb") as f: + self.assertEqual(f.read(), b"image bytes") + + @responses.activate + def test_existing_directory_is_downloaded_then_reused(self): + self._serve_zip() + for _ in range(2): + self.version.download("coco", self.location) + self.assertEqual(len(responses.calls), 1) + with open(os.path.join(self.location, "train", "image.jpg"), "rb") as f: + self.assertEqual(f.read(), b"image bytes") + + @responses.activate + def test_failed_download_and_extraction_are_retried(self): + for status, body in ((403, b"Forbidden"), (200, b"broken zip")): + with self.subTest(status=status): + responses.reset() + self._serve_zip(body, status) + with self.assertRaises((requests.HTTPError, zipfile.BadZipFile)): + self.version.download("coco", self.location, overwrite=True) + with open(self.marker) as f: + self.assertFalse(json.load(f)["complete"]) + self._serve_zip() + self.version.download("coco", self.location) + self.assertEqual(len(responses.calls), 2) + + @responses.activate + def test_failed_postprocessing_invalidates_previous_cache(self): + self._serve_zip() + self.version.download("coco", self.location) + with patch.object(self.version, "_Version__reformat_yaml", side_effect=ValueError("Invalid YAML")): + with self.assertRaisesRegex(ValueError, "Invalid YAML"): + self.version.download("coco", self.location, overwrite=True) + with open(self.marker) as f: + self.assertFalse(json.load(f)["complete"]) + self.version.download("coco", self.location) + self.assertEqual(len(responses.calls), 3) + + @responses.activate + def test_other_formats_versions_and_invalid_metadata_are_not_reused(self): + self._serve_zip() + self.version.download("coco", self.location) + with self.assertRaises(FileExistsError): + self.version.download("voc", self.location) + self.version.download("voc", self.location, overwrite=True) + self.version.version = "2" + with self.assertRaises(FileExistsError): + self.version.download("voc", self.location) + self.version.download("voc", self.location, overwrite=True) + self.version.id = "other-workspace/other-project/2" + with self.assertRaises(FileExistsError): + self.version.download("voc", self.location) + self.version.download("voc", self.location, overwrite=True) + with open(self.marker, "w") as f: + f.write("{") + with self.assertRaises(FileExistsError): + self.version.download("voc", self.location) + self.version.download("voc", self.location, overwrite=True) + self.assertEqual(len(responses.calls), 5) + + @responses.activate + def test_region_switch_does_not_reuse_dataset(self): + self._serve_zip() + self.version.download("coco", self.location) + with patch("roboflow.core.version.API_URL", "https://api.other-region.test"): + with self.assertRaises(FileExistsError): + self.version.download("coco", self.location) + self.assertEqual(len(responses.calls), 1) + self.version.download("coco", self.location, overwrite=True) + self.assertEqual(len(responses.calls), 2) + + +class TestTrainResult(unittest.TestCase): + def setUp(self): + self.version = get_version(exports=["yolov5pytorch"]) + for target in ( + "roboflow.core.version.Version._Version__wait_if_generating", + "roboflow.adapters.rfapi.start_version_training", + "roboflow.core.version.time.sleep", + ): + mock = patch(target) + mock.start() + self.addCleanup(mock.stop) + + @patch("roboflow.adapters.rfapi.get_version") + def test_failure_never_returns_a_model(self, get_version): + for existing_model in (None, MagicMock()): + for training in ({"status": "failed"}, {"status": "failed", "results": {"old": "result"}}): + with self.subTest(existing_model=existing_model, training=training): + self.version.model = existing_model + get_version.return_value = {"version": {"train": training}} + with self.assertRaisesRegex(RuntimeError, "Training failed for test-workspace/test-project/1"): + self.version.train(model_type="yolov8n") + self.assertIs(self.version._model, existing_model) + + @patch("roboflow.adapters.rfapi.get_version") + def test_failure_after_polling_raises(self, get_version): + get_version.side_effect = [ + {"version": {"train": {"status": "running"}}}, + {"version": {"train": {"status": "failed"}}}, + ] + with self.assertRaisesRegex(RuntimeError, "Training failed"): + self.version.train(model_type="yolov8n") + self.assertIsNone(self.version._model) + self.assertEqual(get_version.call_count, 2) + + @patch("roboflow.adapters.rfapi.get_version") + def test_completed_training_returns_model(self, get_version): + get_version.side_effect = [ + {"version": {"train": {"status": "running"}}}, + {"version": {"train": {"results": {"mAP": 0.9}}}}, + ] + model = self.version.train(model_type="yolov8n") + self.assertIsInstance(model, ObjectDetectionModel) + self.assertIs(model, self.version._model) class TestExport(unittest.TestCase): diff --git a/tests/test_vision_events.py b/tests/test_vision_events.py index 0ab2bb7c..3c31e89a 100644 --- a/tests/test_vision_events.py +++ b/tests/test_vision_events.py @@ -2,6 +2,7 @@ import os import tempfile import unittest +from unittest.mock import patch import responses @@ -17,6 +18,45 @@ class TestVisionEvents(unittest.TestCase): API_KEY = "test_key" WORKSPACE = "test-ws" + @responses.activate + def test_existing_workspace_uses_refreshed_region_and_environment(self): + from roboflow import config + + ws = self._make_workspace() + cases = ( + ("us", "prod", None, "https://api.roboflow.com"), + ("eu", "prod", None, "https://api.roboflow.eu"), + ("us", "staging", None, "https://api.roboflow.one"), + ("eu", "staging", None, "https://api.roboflow-eu.one"), + ("eu", "prod", "https://api.test", "https://api.test"), + ("us", "prod", None, "https://api.roboflow.com"), + ) + try: + with tempfile.TemporaryDirectory() as directory: + for region, environment, override, expected in cases: + with self.subTest(region=region, environment=environment, override=override): + responses.reset() + with patch.dict( + os.environ, + { + "ROBOFLOW_REGION": region, + "ROBOFLOW_ENVIRONMENT": environment, + "ROBOFLOW_CONFIG_DIR": os.path.join(directory, "config.json"), + }, + ): + os.environ.pop("API_URL", None) + if override: + os.environ["API_URL"] = override + config.refresh_region_urls() + responses.post(f"{expected}/vision-events", json={"eventId": "event"}, status=201) + responses.get(f"{expected}/vision-events/use-cases", json={"useCases": []}) + ws.write_vision_event({"eventId": "event"}) + ws.list_vision_event_use_cases() + for call in responses.calls: + self.assertEqual(call.request.headers["Authorization"], f"Bearer {self.API_KEY}") + finally: + config.refresh_region_urls() + def _make_workspace(self): from roboflow.core.workspace import Workspace diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 5c33b5b9..0433b789 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -4,9 +4,10 @@ import tempfile import unittest import zipfile -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from roboflow.core.workspace import Workspace, _zip_directory +from roboflow.adapters.rfapi import AnnotationSaveError, ImageUploadError, RoboflowError +from roboflow.core.workspace import Workspace, _poll_zip_status, _zip_directory def _make_workspace(): @@ -49,6 +50,71 @@ def test_filters_hidden_and_junk_entries(self): os.unlink(zip_path) +class TestDatasetUploadResults(unittest.TestCase): + @patch("roboflow.util.folderparser.parsefolder", return_value={"location": "/dataset", "images": []}) + def test_default_success_return_is_backward_compatible(self, parsefolder): + workspace = _make_workspace() + project = MagicMock(id="test-ws/project", type="object-detection") + with patch.object(workspace, "_get_or_create_project", return_value=(project, False)): + self.assertIsNone(workspace.upload_dataset("/dataset", "project")) + + @patch("roboflow.util.folderparser.parsefolder") + def test_counts_uploaded_and_duplicate_images(self, parsefolder): + project = MagicMock(id="test-ws/project", type="object-detection") + project.upload_image.side_effect = [ + ({"id": "new", "success": True}, 0.1, 0), + ({"id": "existing", "duplicate": True}, 0.1, 0), + ] + parsefolder.return_value = { + "location": "/dataset", + "images": [{"file": f"/{name}.jpg", "split": "train"} for name in ("new", "existing")], + } + workspace = _make_workspace() + with patch.object(workspace, "_get_or_create_project", return_value=(project, False)): + result = workspace.upload_dataset("/dataset", "project", num_workers=1, return_summary=True) + self.assertEqual(result, {"uploaded": 1, "duplicates": 1}) + + @patch("roboflow.util.folderparser.parsefolder") + def test_reports_all_failures_after_attempting_every_image(self, parsefolder): + project = MagicMock(id="test-ws/project", type="object-detection") + + def upload_image(image_path, **kwargs): + if image_path.endswith("image.jpg"): + raise ImageUploadError("Upload rejected") + if image_path.endswith("unexpected.jpg"): + raise ValueError("Invalid image") + return {"id": image_path, "success": True}, 0.1, 0 + + project.upload_image.side_effect = upload_image + project.save_annotation.side_effect = AnnotationSaveError("Annotation rejected") + parsefolder.return_value = { + "location": "/dataset", + "images": [ + {"file": "/image.jpg", "split": "train"}, + {"file": "/annotation.jpg", "split": "train", "annotationfile": {"file": "/annotation.xml"}}, + {"file": "/unexpected.jpg", "split": "train"}, + {"file": "/ok.jpg", "split": "train"}, + ], + } + workspace = _make_workspace() + with patch.object(workspace, "_get_or_create_project", return_value=(project, False)): + with self.assertRaisesRegex(RoboflowError, "failed for 3 of 4 images") as raised: + workspace.upload_dataset("/dataset", "project", num_workers=2) + self.assertEqual(project.upload_image.call_count, 4) + for message in ( + "image.jpg: Upload rejected", + "annotation.jpg: annotations: Annotation rejected", + "unexpected.jpg: Invalid image", + ): + self.assertIn(message, str(raised.exception)) + + @patch("roboflow.adapters.rfapi.get_zip_upload_status") + def test_failed_zip_processing_raises(self, get_status): + get_status.return_value = {"status": "failed", "error": "Invalid annotations"} + with self.assertRaisesRegex(RoboflowError, "task_id=task-1.*Invalid annotations"): + _poll_zip_status("key", "workspace", "task-1", 0, 10) + + class TestWorkspaceAsyncTasks(unittest.TestCase): @patch("roboflow.adapters.rfapi.fork_project") def test_fork_project_uses_workspace_destination(self, mock_fork): diff --git a/tests/util/test_folderparser.py b/tests/util/test_folderparser.py index 4f9ddb5b..c1893ac8 100644 --- a/tests/util/test_folderparser.py +++ b/tests/util/test_folderparser.py @@ -3,6 +3,7 @@ import tempfile import unittest from os.path import abspath, dirname +from pathlib import Path from roboflow.util import folderparser @@ -206,6 +207,90 @@ def test_coco_root_annotation_matches_images_in_subdirs(self): self.assertEqual(ann_data["annotations"][0]["bbox"], [10, 20, 100, 200]) +class TestCocoImageMatching(unittest.TestCase): + def _parse(self, image_paths, documents): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for path in image_paths: + image = root / path + image.parent.mkdir(parents=True, exist_ok=True) + image.touch() + for path, references in documents.items(): + annotation = root / path + annotation.parent.mkdir(parents=True, exist_ok=True) + data = { + "info": {}, + "licenses": [], + "categories": [{"id": category, "name": str(category)} for _, category in references], + "images": [ + {"id": i, "file_name": name, "width": 4, "height": 3} for i, (name, _) in enumerate(references) + ], + "annotations": [ + { + "id": i, + "image_id": i, + "category_id": category, + "segmentation": {"size": [3, 4], "counts": [0, 3, 3, 3, 3]}, + } + for i, (_, category) in enumerate(references) + ], + } + annotation.write_text(json.dumps(data), encoding="utf-8") + parsed = folderparser.parsefolder(directory) + return {image["file"]: json.loads(image["annotationfile"]["rawText"]) for image in parsed["images"]} + + def test_split_annotations_match_their_own_images(self): + result = self._parse( + ["train/same.jpg", "valid/same.jpg"], + { + "train/_annotations.coco.json": [("same.jpg", 11)], + "valid/_annotations.coco.json": [("same.jpg", 22)], + }, + ) + self.assertEqual(result["/train/same.jpg"]["annotations"][0]["category_id"], 11) + self.assertEqual(result["/valid/same.jpg"]["annotations"][0]["category_id"], 22) + self.assertEqual( + result["/valid/same.jpg"]["annotations"][0]["segmentation"], + {"size": [3, 4], "counts": [0, 3, 3, 3, 3]}, + ) + + def test_one_coco_file_can_reference_duplicate_basenames(self): + result = self._parse( + ["train/same.jpg", "valid/same.jpg"], + {"_annotations.coco.json": [("train/same.jpg", 11), ("valid/same.jpg", 22)]}, + ) + self.assertEqual(result["/train/same.jpg"]["annotations"][0]["category_id"], 11) + self.assertEqual(result["/valid/same.jpg"]["annotations"][0]["category_id"], 22) + + def test_nested_coco_paths_are_normalized(self): + for name in ("images/same.jpg", "images\\same.jpg", "./images/same.jpg", "train/images/same.jpg"): + with self.subTest(name=name): + result = self._parse(["train/images/same.jpg"], {"train/_annotations.coco.json": [(name, 11)]}) + self.assertEqual(result["/train/images/same.jpg"]["images"][0]["file_name"], name) + + def test_unique_basename_and_stem_fallbacks_still_work(self): + for name in ("old/path/same.jpg", "C:\\old\\path\\same.jpg", "old/path/same.png"): + with self.subTest(name=name): + result = self._parse(["same.jpg"], {"_annotations.coco.json": [(name, 11)]}) + self.assertEqual(result["/same.jpg"]["annotations"][0]["category_id"], 11) + + def test_ambiguous_image_fallback_raises(self): + with self.assertRaisesRegex(ValueError, "[Aa]mbiguous.*same"): + self._parse( + ["train/same.jpg", "valid/same.jpg"], + {"_annotations.coco.json": [("old/same.jpg", 11)]}, + ) + + def test_competing_coco_references_raise(self): + cases = [ + {"first.coco.json": [("same.jpg", 11)], "second.coco.json": [("same.jpg", 22)]}, + {"_annotations.coco.json": [("old/same.jpg", 11), ("other/same.jpg", 22)]}, + ] + for documents in cases: + with self.subTest(documents=documents), self.assertRaisesRegex(ValueError, "[Aa]mbiguous.*same"): + self._parse(["same.jpg"], documents) + + def _assertJsonMatchesFile(actual, filename): with open(filename) as file: expected = json.load(file) diff --git a/tests/util/test_general.py b/tests/util/test_general.py new file mode 100644 index 00000000..06086463 --- /dev/null +++ b/tests/util/test_general.py @@ -0,0 +1,83 @@ +import gzip +import io +import os +import tempfile +import unittest +from unittest.mock import MagicMock, patch + +import requests +import responses + +from roboflow.util import http +from roboflow.util.general import save_response_content + + +class TestSaveResponseContent(unittest.TestCase): + def test_preserves_bytes_with_known_unknown_or_invalid_content_length(self): + content = bytes(range(256)) * 1024 + b"tail" + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "download") + for length in (str(len(content)), None, "invalid", "-1"): + with self.subTest(length=length): + response = requests.Response() + response.raw = io.BytesIO(content) + if length is not None: + response.headers["content-length"] = length + save_response_content(response, path) + with open(path, "rb") as f: + self.assertEqual(f.read(), content) + + @patch("roboflow.util.general.tqdm") + def test_progress_counts_bytes_and_skips_empty_chunks(self, tqdm): + response = MagicMock(headers={}) + response.iter_content.return_value = iter([b"ab", b"", b"cde"]) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "download") + save_response_content(response, path) + with open(path, "rb") as f: + self.assertEqual(f.read(), b"abcde") + self.assertIsNone(tqdm.call_args.kwargs["total"]) + self.assertEqual( + [call.args[0] for call in tqdm.return_value.__enter__.return_value.update.call_args_list], [2, 3] + ) + + def test_empty_response_creates_empty_file(self): + response = requests.Response() + response.raw = io.BytesIO() + response.headers["content-length"] = "0" + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "download") + save_response_content(response, path) + self.assertEqual(os.path.getsize(path), 0) + + @responses.activate + @patch("roboflow.util.general.tqdm") + def test_compressed_response_uses_decoded_bytes_without_a_wire_length_total(self, tqdm): + content = b"model weights" * 10000 + compressed = gzip.compress(content) + responses.get( + "https://download.test/weights", + body=compressed, + headers={"Content-Encoding": "gzip", "Content-Length": str(len(compressed))}, + ) + with ( + tempfile.TemporaryDirectory() as directory, + http.get("https://download.test/weights", stream=True) as response, + ): + path = os.path.join(directory, "download") + save_response_content(response, path) + with open(path, "rb") as f: + self.assertEqual(f.read(), content) + self.assertIsNone(tqdm.call_args.kwargs["total"]) + + @patch("roboflow.util.general.tqdm") + def test_transfer_encoding_takes_precedence_over_content_length(self, tqdm): + response = requests.Response() + response.raw = io.BytesIO(b"decoded chunks") + response.headers.update({"Transfer-Encoding": "chunked", "Content-Length": "999"}) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "download") + save_response_content(response, path) + with open(path, "rb") as f: + self.assertEqual(f.read(), b"decoded chunks") + self.assertIsNone(tqdm.call_args.kwargs["total"])