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"])