diff --git a/roboflow/core/version.py b/roboflow/core/version.py index fd4a3f5f..4d8a9af0 100644 --- a/roboflow/core/version.py +++ b/roboflow/core/version.py @@ -34,8 +34,8 @@ from roboflow.models.vlm import VLMModel from roboflow.util.annotations import amend_data_yaml from roboflow.util.general import extract_zip, write_line -from roboflow.util.model_processor import process, validate_model_type_for_project -from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions, normalize_yolo_model_type +from roboflow.util.model_processor import package_custom_weights_interactive, validate_model_type_for_project +from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions if TYPE_CHECKING: import numpy as np @@ -497,14 +497,10 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best 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". """ - model_type = normalize_yolo_model_type(model_type) - zip_file_name, model_type = process(model_type, model_path, filename) + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) - if zip_file_name is None: - raise RuntimeError("Failed to process model") - - self._validate_against_project_type(model_type) - self._upload_zip(model_type, model_path, zip_file_name) + self._validate_against_project_type(bundle.model_type) + self._upload_zip(bundle.model_type, model_path, bundle.archive_path.name) def _validate_against_project_type(self, model_type: str) -> None: validate_model_type_for_project(model_type, self.type, self.project) diff --git a/roboflow/core/workspace.py b/roboflow/core/workspace.py index c1dca153..64a5c430 100644 --- a/roboflow/core/workspace.py +++ b/roboflow/core/workspace.py @@ -787,8 +787,10 @@ def deploy_model( filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". """ - from roboflow.util.model_processor import process, validate_model_type_for_project - from roboflow.util.versions import normalize_yolo_model_type + from roboflow.util.model_processor import ( + package_custom_weights_interactive, + validate_model_type_for_project, + ) if not project_ids: raise ValueError("At least one project ID must be provided") @@ -800,16 +802,12 @@ def deploy_model( if project_id not in projects_by_id: raise ValueError(f"Project {project_id} is not accessible in this workspace") - model_type = normalize_yolo_model_type(model_type) - zip_file_name, model_type = process(model_type, model_path, filename) - - if zip_file_name is None: - raise RuntimeError("Failed to process model") + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) for project_id in project_ids: - validate_model_type_for_project(model_type, projects_by_id[project_id].get("type", ""), project_id) + validate_model_type_for_project(bundle.model_type, projects_by_id[project_id].get("type", ""), project_id) - self._upload_zip(model_type, model_path, project_ids, model_name, zip_file_name) + self._upload_zip(bundle.model_type, model_path, project_ids, model_name, bundle.archive_path.name) def _upload_zip( self, diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 0022c11a..03cfe8bf 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -1,8 +1,40 @@ +"""Packaging of custom model weights for Roboflow upload. + +The public, non-interactive entry point is :func:`package_custom_weights`. It +only builds the upload archive; it never prompts, prints, or uploads, so it is +safe to call from servers and other headless environments (for example the +Roboflow MCP server):: + + from roboflow.util.model_processor import package_custom_weights + + bundle = package_custom_weights("yolov8n", "runs/detect/train") + try: + ... # upload bundle.archive_path + finally: + bundle.cleanup() + +Expected, user-correctable failures raise :class:`ModelPackagingError` +subclasses; anything else escaping these helpers is a bug. + +The legacy :func:`process` entry point and the ``Version.deploy`` / +``Workspace.deploy_model`` flows wrap the packaging step with +:func:`package_custom_weights_interactive`, which preserves the historical +print-and-confirm CLI behavior. +""" + +from __future__ import annotations + import json +import math import os import shutil +import tarfile +import tempfile import zipfile -from typing import Callable, Optional +from dataclasses import dataclass +from importlib import import_module +from pathlib import Path +from typing import Any import yaml @@ -19,7 +51,177 @@ TYPE_OBJECT_DETECTION, TYPE_SEMANTIC_SEGMENTATION, ) -from roboflow.util.versions import print_warn_for_wrong_dependencies_versions +from roboflow.util.versions import get_wrong_dependencies_versions, normalize_yolo_model_type + +SUPPORTED_MODELS = ( + "yolov5", + "yolov7", + "yolov7-seg", + "yolov8", + "yolov9", + "yolov10", + "yolov11", + "yolov12", + "yolo26", + "yolonas", + "paligemma", + "paligemma2", + "florence-2", + "rfdetr", +) + +SUPPORTED_HUGGINGFACE_TYPES = ( + "florence-2-base", + "florence-2-large", + "paligemma-3b-pt-224", + "paligemma-3b-pt-448", + "paligemma-3b-pt-896", + "paligemma2-3b-pt-224", + "paligemma2-3b-pt-448", + "paligemma2-3b-pt-896", + "paligemma2-3b-pt-224-peft", + "paligemma2-3b-pt-448-peft", + "paligemma2-3b-pt-896-peft", +) + +SUPPORTED_RFDETR_TYPES = ( + # Detection models + "rfdetr-base", + "rfdetr-nano", + "rfdetr-small", + "rfdetr-medium", + "rfdetr-large", + "rfdetr-xlarge", + "rfdetr-2xlarge", + # Segmentation models + "rfdetr-seg-nano", + "rfdetr-seg-small", + "rfdetr-seg-medium", + "rfdetr-seg-large", + "rfdetr-seg-xlarge", + "rfdetr-seg-2xlarge", +) + +# YOLO families Roboflow rejects without a size suffix (e.g. `yolov8` must be +# `yolov8n`/`yolov8s`/...). Legacy yolov5/7/9 go through the opt.yaml path and are +# intentionally excluded. +ULTRALYTICS_YOLO_FAMILIES = ("yolov8", "yolov10", "yolov11", "yolov12", "yolo26") + +# Canonical (depth_multiple, width_multiple) -> size letter for the classic YOLO +# scaling (v5/v8/v9/v10). Newer families (v11+) instead store an explicit ``scale`` +# letter in the model yaml, which is read first. +YOLO_DEPTH_WIDTH_TO_SIZE = { + (0.33, 0.25): "n", + (0.33, 0.50): "s", + (0.67, 0.75): "m", + (0.67, 1.00): "b", + (1.00, 1.00): "l", + (1.00, 1.25): "x", +} + +# Position-encoding grid size (DINOv2 tokens per side) each *known* RF-DETR variant +# is built with, mirroring rfdetr/config.py. Roboflow reconstructs the architecture +# from the model_type at the variant's default resolution, so a checkpoint trained at +# that default must match or state_dict loading fails on the backbone +# position_embeddings. Variants absent here (e.g. detection xlarge/2xlarge, which have +# no standard config) are not grid-checked. A checkpoint trained at a custom +# resolution may not match any entry; that case warns rather than blocks. +RFDETR_POSITIONAL_ENCODING_SIZE = { + "rfdetr-nano": 24, + "rfdetr-small": 32, + "rfdetr-medium": 36, + "rfdetr-base": 37, + "rfdetr-large": 44, + "rfdetr-seg-nano": 26, + "rfdetr-seg-small": 32, + "rfdetr-seg-medium": 36, + "rfdetr-seg-large": 42, + "rfdetr-seg-xlarge": 52, + "rfdetr-seg-2xlarge": 64, +} + + +class ModelPackagingError(Exception): + """Custom weights could not be packaged for a user-correctable reason. + + Consumers can treat any instance of this class as an expected input problem + (bad model_type, missing files, mismatched metadata, ...) and surface the + message to the user. Exceptions that are not ModelPackagingError indicate + bugs and are deliberately not wrapped. + """ + + +class UnsupportedModelError(ModelPackagingError, ValueError): + """The model_type is not supported for custom weights upload.""" + + +class TaskMismatchError(ModelPackagingError, ValueError): + """The model_type's task conflicts with the checkpoint or the project type.""" + + +class MissingFileError(ModelPackagingError, FileNotFoundError): + """A file required for packaging was not found.""" + + +class MissingDependencyError(ModelPackagingError, RuntimeError): + """A Python package required to package these weights is not installed.""" + + +class DependencyMismatchError(ModelPackagingError, RuntimeError): + """An installed dependency version differs from the recommended one. + + Retry with ``allow_dependency_mismatch=True`` to package anyway. + """ + + retry_flag = "allow_dependency_mismatch" + + def __init__(self, message: str, *, dependency: str, required: str, installed: str): + super().__init__(message) + self.dependency = dependency + self.required = required + self.installed = installed + + +class SizeMismatchError(ModelPackagingError, ValueError): + """The declared model size/variant conflicts with the checkpoint architecture. + + Retry with ``allow_size_mismatch=True`` to package the requested model_type + as-is. + """ + + retry_flag = "allow_size_mismatch" + + def __init__(self, message: str, *, requested: str, detected: str | None = None): + super().__init__(message) + self.requested = requested + self.detected = detected + + +@dataclass(frozen=True) +class ModelUploadBundle: + """Packaged archive ready to upload through the Roboflow API. + + ``model_type`` is the resolved type (it may differ from the requested one, + e.g. ``yolov8`` filled in as ``yolov8n`` from the checkpoint architecture). + ``owns_build_dir`` is True when :func:`package_custom_weights` created a + temporary build directory; call :meth:`cleanup` once the archive has been + consumed. + """ + + archive_path: Path + build_dir: Path + model_type: str + warnings: tuple[str, ...] = () + owns_build_dir: bool = False + + @property + def size_bytes(self) -> int: + return self.archive_path.stat().st_size + + def cleanup(self) -> None: + """Remove the build directory if this bundle created it (no-op otherwise).""" + if self.owns_build_dir: + shutil.rmtree(self.build_dir, ignore_errors=True) def task_of_model_type(model_type: str) -> str: @@ -36,7 +238,7 @@ def task_of_model_type(model_type: str) -> str: def validate_model_type_for_project(model_type: str, project_type: str, project_id: str) -> None: - """Raise ValueError if model_type's task doesn't match the Roboflow project type.""" + """Raise TaskMismatchError if model_type's task doesn't match the Roboflow project type.""" expected = { TYPE_OBJECT_DETECTION: TASK_DET, TYPE_INSTANCE_SEGMENTATION: TASK_SEG, @@ -48,70 +250,207 @@ def validate_model_type_for_project(model_type: str, project_type: str, project_ return actual = task_of_model_type(model_type) if actual != expected: - raise ValueError( + raise TaskMismatchError( f"Project '{project_id}' is type '{project_type}' (task '{expected}') " f"but model_type '{model_type}' implies task '{actual}'." ) +def package_custom_weights( + model_type: str, + model_path: str, + filename: str = "weights/best.pt", + *, + build_dir: str | Path | None = None, + allow_dependency_mismatch: bool = False, + allow_size_mismatch: bool = False, +) -> ModelUploadBundle: + """Package locally trained custom weights into a Roboflow upload archive. + + This is the public packaging entry point. It is non-interactive and free of + side effects on ``model_path``: it never prompts, prints, or writes into the + source directory. Heavy dependencies (torch, ultralytics) are imported + lazily, only for the model families that need them. + + Args: + model_type: Roboflow model type (e.g. "yolov8n", "rfdetr-base"). + model_path: Directory containing the trained model artifacts. + filename: Weights file path, relative to ``model_path``. + build_dir: Directory to write intermediate artifacts and the final + archive into. Defaults to a fresh temporary directory owned by the + returned bundle; call ``bundle.cleanup()`` when done. + allow_dependency_mismatch: Record a warning instead of raising + DependencyMismatchError when an installed dependency version is not + the recommended one. + allow_size_mismatch: Record a warning instead of raising + SizeMismatchError when the declared model size/variant conflicts + with the checkpoint architecture. + + Returns: + ModelUploadBundle with the archive path, the resolved model_type, and + any warnings collected while packaging. + + Raises: + ModelPackagingError: (or a subclass) for user-correctable problems. + """ + normalized_model_type = normalize_yolo_model_type(model_type.strip()) + source_dir = Path(model_path).expanduser().resolve() + if not source_dir.is_dir(): + raise MissingFileError(f"Model path '{model_path}' does not exist or is not a directory.") + + owns_build_dir = build_dir is None + if build_dir is None: + build_path = Path(tempfile.mkdtemp(prefix="roboflow-package-")) + else: + build_path = Path(build_dir) + build_path.mkdir(parents=True, exist_ok=True) + + try: + archive_path, resolved_model_type, warnings = _process_model( + model_type=normalized_model_type, + model_path=source_dir, + filename=filename, + build_dir=build_path, + allow_dependency_mismatch=allow_dependency_mismatch, + allow_size_mismatch=allow_size_mismatch, + ) + except BaseException: + if owns_build_dir: + shutil.rmtree(build_path, ignore_errors=True) + raise + + return ModelUploadBundle( + archive_path=archive_path, + build_dir=build_path, + model_type=resolved_model_type, + warnings=tuple(warnings), + owns_build_dir=owns_build_dir, + ) + + +def package_custom_weights_interactive( + model_type: str, + model_path: str, + filename: str = "weights/best.pt", + *, + build_dir: str | Path | None = None, +) -> ModelUploadBundle: + """Package weights with the historical interactive SDK behavior. + + Used by ``Version.deploy`` and ``Workspace.deploy_model``: warnings are + printed, and dependency/size mismatches ask for confirmation before + retrying with the corresponding override. Declining re-raises the error. + """ + allow_dependency_mismatch = False + allow_size_mismatch = False + while True: + try: + bundle = package_custom_weights( + model_type, + model_path, + filename, + build_dir=build_dir, + allow_dependency_mismatch=allow_dependency_mismatch, + allow_size_mismatch=allow_size_mismatch, + ) + except (DependencyMismatchError, SizeMismatchError) as error: + print(error) + answer = input("Would you like to continue anyway? y/n: ") + if answer.lower() != "y": + raise + if isinstance(error, DependencyMismatchError): + allow_dependency_mismatch = True + else: + allow_size_mismatch = True + continue + for warning in bundle.warnings: + print(warning) + return bundle + + def process(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - processor = _get_processor_function(model_type) - return processor(model_type, model_path, filename) - - -def _get_processor_function(model_type: str) -> Callable: - supported_models = [ - "yolov5", - "yolov7", - "yolov7-seg", - "yolov8", - "yolov9", - "yolov10", - "yolov11", - "yolov12", - "yolo26", - "yolonas", - "paligemma", - "paligemma2", - "florence-2", - "rfdetr", - ] - - if not any(supported_model in model_type for supported_model in supported_models): - raise (ValueError(f"Model type {model_type} not supported. Supported models are {supported_models}")) + """Legacy packaging entry point, kept for backwards compatibility. + + Preserves the historical contract end to end: packages into ``model_path`` + (intermediate artifacts and the final archive land there), prints packaging + warnings, asks for confirmation on dependency/size mismatches, and returns + ``(archive_file_name, resolved_model_type)``. Headless code should call + :func:`package_custom_weights` instead. + """ + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) + return bundle.archive_path.name, bundle.model_type + + +def _process_model( + *, + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_dependency_mismatch: bool, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + if not model_type.startswith(SUPPORTED_MODELS): + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for custom weights upload. " + f"It must start with a supported family: {', '.join(SUPPORTED_MODELS)}." + ) if model_type.startswith(("paligemma", "paligemma2", "florence-2")): - if any(model in model_type for model in ["paligemma", "paligemma2", "florence-2"]): - supported_hf_types = [ - "florence-2-base", - "florence-2-large", - "paligemma-3b-pt-224", - "paligemma-3b-pt-448", - "paligemma-3b-pt-896", - "paligemma2-3b-pt-224", - "paligemma2-3b-pt-448", - "paligemma2-3b-pt-896", - "paligemma2-3b-pt-224-peft", - "paligemma2-3b-pt-448-peft", - "paligemma2-3b-pt-896-peft", - ] - if model_type not in supported_hf_types: - raise RuntimeError( - f"{model_type} not supported for this type of upload." - f"Supported upload types are {supported_hf_types}" - ) - return _process_huggingface - - if "yolonas" in model_type: - return _process_yolonas - - if "rfdetr" in model_type: - return _process_rfdetr - - return _process_yolo - - -def _detect_yolo_task(model_instance) -> Optional[str]: + return _process_huggingface(model_type, model_path, build_dir) + if model_type.startswith("yolonas"): + return _process_yolonas(model_type, model_path, filename, build_dir) + if model_type.startswith("rfdetr"): + return _process_rfdetr(model_type, model_path, filename, build_dir, allow_size_mismatch) + return _process_yolo( + model_type, + model_path, + filename, + build_dir, + allow_dependency_mismatch, + allow_size_mismatch, + ) + + +def _import_required_module(module_name: str, install_hint: str) -> Any: + try: + return import_module(module_name) + except ImportError as exc: + raise MissingDependencyError( + f"The '{module_name}' Python package is required to package these " + f"custom weights. Install it with `{install_hint}`." + ) from exc + + +def _check_dependency_version( + *, + dependency: str, + operator: str, + required_version: str, + allow_mismatch: bool, + warnings: list[str], +) -> None: + mismatches = get_wrong_dependencies_versions([(dependency, operator, required_version)]) + if not mismatches: + return + _, _, _, installed = mismatches[0] + message = ( + f"{dependency}{operator}{required_version} is recommended for this " + f"upload, but {dependency} {installed} is installed." + ) + if allow_mismatch: + warnings.append(message) + return + raise DependencyMismatchError( + f"{message} Retry with allow_dependency_mismatch=True to package with the " + f"installed version, or `pip install {dependency}{operator}{required_version}`.", + dependency=dependency, + required=f"{dependency}{operator}{required_version}", + installed=installed, + ) + + +def _detect_yolo_task(model_instance: Any) -> str | None: """Detect the training task of an Ultralytics model instance via its class name.""" if model_instance is None: return None @@ -125,87 +464,188 @@ def _detect_yolo_task(model_instance) -> Optional[str]: }.get(type(model_instance).__name__) -def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - if "yolov8" in model_type: - try: - import torch - import ultralytics +def _class_names_from_model_instance(model_instance: Any) -> list[str]: + names = getattr(model_instance, "names", None) + if isinstance(names, list): + return names + if isinstance(names, dict): + return [name for _, name in sorted(names.items(), key=lambda item: item[0])] + raise ModelPackagingError("Could not extract class names from the model checkpoint.") - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov8" - " models. Please install it with `pip install ultralytics`" - ) - print_warn_for_wrong_dependencies_versions([("ultralytics", "==", "8.0.196")], ask_to_continue=True) +def _filtered_args(args: Any) -> dict[str, Any]: + if not isinstance(args, dict): + args = vars(args) + return {k: v for k, v in args.items() if k in {"model", "imgsz", "batch"}} - elif "yolov10" in model_type: - try: - import torch - import ultralytics - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov10" - " models. Please install it with `pip install ultralytics`" - ) +def _load_checkpoint(torch_module: Any, checkpoint_path: Path, *, map_location: str | None = None) -> Any: + kwargs: dict[str, Any] = {"weights_only": False} + if map_location is not None: + kwargs["map_location"] = map_location + return torch_module.load(checkpoint_path, **kwargs) - elif "yolov5" in model_type or "yolov7" in model_type or "yolov9" in model_type: - try: - import torch - except ImportError: - raise RuntimeError( - f"The torch python package is required to deploy {model_type} models." - " Please install it with `pip install torch`" - ) - elif "yolov11" in model_type: - try: - import torch - import ultralytics +def _legacy_yolo_args(opts: dict[str, Any], opt_path: Path) -> dict[str, Any]: + """Return required legacy YOLO upload args from opt.yaml.""" + if "imgsz" in opts: + image_size = opts["imgsz"] + elif "img_size" in opts: + image_size = opts["img_size"] + else: + raise ModelPackagingError(f"{opt_path} is missing required key 'imgsz' or 'img_size'.") + if "batch_size" not in opts: + raise ModelPackagingError(f"{opt_path} is missing required key 'batch_size'.") + return {"imgsz": image_size, "batch": opts["batch_size"]} - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov11" - " models. Please install it with `pip install ultralytics`" - ) - print_warn_for_wrong_dependencies_versions([("ultralytics", ">=", "8.3.0")], ask_to_continue=True) +def _infer_yolo_size(model_instance: Any) -> str | None: + """Infer a YOLO size letter (n/s/m/l/x/...) from a loaded checkpoint. - elif "yolov12" in model_type: - try: - import torch - import ultralytics + Prefers an explicit ``scale`` letter in the model yaml (set by newer + Ultralytics), then maps the ``(depth_multiple, width_multiple)`` pair used by + the classic scaling. Returns None when neither signal is present. + """ + yaml_cfg = getattr(model_instance, "yaml", None) or {} + scale = yaml_cfg.get("scale") + if isinstance(scale, str) and len(scale) == 1 and scale.isalpha(): + return scale.lower() + + depth = yaml_cfg.get("depth_multiple") + width = yaml_cfg.get("width_multiple") + if isinstance(depth, (int, float)) and isinstance(width, (int, float)): + for (ref_depth, ref_width), letter in YOLO_DEPTH_WIDTH_TO_SIZE.items(): + if abs(depth - ref_depth) < 1e-6 and abs(width - ref_width) < 1e-6: + return letter + return None - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov12" - " models. Please install it from `https://github.com/sunsmarterjie/yolov12`" + +def _resolve_yolo_size( + model_type: str, + model_instance: Any, + warnings: list[str], + allow_mismatch: bool = False, +) -> str: + """Fill in or check a YOLO model_type's size suffix against the checkpoint. + + Roboflow rejects bare family names (e.g. ``yolov8``) with an + ``InvalidModelTypeException`` because it needs the model size, and a size that + disagrees with the weights fails conversion. A *missing* size is inferred and + filled in. A *supplied* size that conflicts with the inferred one raises so the + caller can confirm — unless ``allow_mismatch`` is set, in which case the + supplied size is packaged as-is with a warning. A user size is also kept when + the size cannot be inferred. Returns the resolved model_type. + """ + core = model_type.lower().split("-", 1)[0] + family = next((f for f in ULTRALYTICS_YOLO_FAMILIES if core.startswith(f)), None) + if family is None: + return model_type + + inferred = _infer_yolo_size(model_instance) + provided = core[len(family) :] + task_suffix = model_type[len(core) :] + + if inferred is None: + if not provided: + raise SizeMismatchError( + f"model_type '{model_type}' is missing a size suffix and the size " + f"could not be inferred from the checkpoint. Specify it explicitly, " + f"e.g. '{family}n', '{family}s', '{family}m', '{family}l', '{family}x'.", + requested=model_type, ) + return model_type + + if not provided: + warnings.append( + f"Inferred model size '{family}{inferred}' from the checkpoint " + f"architecture (model_type was '{model_type}')." + ) + return f"{family}{inferred}{task_suffix}" + + if provided == inferred: + return model_type - print( - "\n!!! ATTENTION !!!\n" - "Model must be trained and uploaded using ultralytics from https://github.com/sunsmarterjie/yolov12\n" - "or through the Roboflow platform\n" - "!!! ATTENTION !!!\n" + if allow_mismatch: + warnings.append( + f"model_type '{model_type}' declares size '{provided}', but the checkpoint " + f"architecture is '{family}{inferred}'. Packaging as '{model_type}' as requested." ) + return model_type + + raise SizeMismatchError( + f"You specified model_type '{model_type}' (size '{provided}'), but the " + f"checkpoint architecture is '{family}{inferred}'. They don't match, so " + f"Roboflow's weight conversion would fail. Upload as '{family}{inferred}" + f"{task_suffix}', or set allow_size_mismatch=True to upload " + f"'{model_type}' exactly as specified.", + requested=model_type, + detected=f"{family}{inferred}{task_suffix}", + ) + - print_warn_for_wrong_dependencies_versions([("ultralytics", "==", "8.3.63")], ask_to_continue=True) +def _process_yolo( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_dependency_mismatch: bool, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + warnings: list[str] = [] + torch = _import_required_module("torch", "pip install torch") + ultralytics = None + if "yolov8" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics==8.0.196") + _check_dependency_version( + dependency="ultralytics", + operator="==", + required_version="8.0.196", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) + elif "yolov10" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") + elif "yolov11" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install 'ultralytics>=8.3.0'") + _check_dependency_version( + dependency="ultralytics", + operator=">=", + required_version="8.3.0", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) + elif "yolov12" in model_type: + ultralytics = _import_required_module( + "ultralytics", + "pip install git+https://github.com/sunsmarterjie/yolov12.git", + ) + warnings.append( + "YOLOv12 uploads must use the Ultralytics fork from " + "https://github.com/sunsmarterjie/yolov12 or a Roboflow-trained model." + ) + _check_dependency_version( + dependency="ultralytics", + operator="==", + required_version="8.3.63", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) elif "yolo26" in model_type: - try: - import torch - import ultralytics + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolo26" - " models. Please install it with `pip install ultralytics`" - ) + checkpoint_path = model_path / filename + if not checkpoint_path.exists(): + raise MissingFileError(f"Model weights file '{checkpoint_path}' was not found.") - model = torch.load(os.path.join(model_path, filename), weights_only=False) + checkpoint = _load_checkpoint(torch, checkpoint_path) + if not isinstance(checkpoint, dict): + raise ModelPackagingError(f"Model weights file '{checkpoint_path}' is not a supported checkpoint dictionary.") + model_instance = checkpoint.get("model") or checkpoint.get("ema") + if model_instance is None: + raise ModelPackagingError("Could not find a 'model' or 'ema' entry in the checkpoint.") - model_instance = model["model"] if "model" in model and model["model"] is not None else model["ema"] + model_type = _resolve_yolo_size(model_type, model_instance, warnings, allow_size_mismatch) detected_task = _detect_yolo_task(model_instance) if detected_task: @@ -213,28 +653,15 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, if existing_task == TASK_DET and detected_task != TASK_DET: model_type = f"{model_type}-{detected_task}" elif existing_task != detected_task: - raise ValueError( + raise TaskMismatchError( f"model_type '{model_type}' implies task '{existing_task}' but the " f".pt file is a '{detected_task}' checkpoint. Use a matching model_type." ) - if isinstance(model_instance.names, list): - class_names = model_instance.names - else: - class_names = [] - for i, val in enumerate(model_instance.names): - class_names.append((val, model_instance.names[val])) - class_names.sort(key=lambda x: x[0]) - class_names = [x[1] for x in class_names] - - if ( - "yolov8" in model_type - or "yolov10" in model_type - or "yolov11" in model_type - or "yolov12" in model_type - or "yolo26" in model_type - ): - # try except for backwards compatibility with older versions of ultralytics + class_names = _class_names_from_model_instance(model_instance) + if any(name in model_type for name in ULTRALYTICS_YOLO_FAMILIES): + if ultralytics is None: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") if ( "-cls" in model_type or model_type.startswith("yolov10") @@ -243,78 +670,51 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, or model_type.startswith("yolo26") ): nc = model_instance.yaml["nc"] - args = model["train_args"] + args = checkpoint["train_args"] else: nc = model_instance.nc args = model_instance.args - try: - model_artifacts = { - "names": class_names, - "yaml": model_instance.yaml, - "nc": nc, - "args": {k: val for k, val in args.items() if ((k == "model") or (k == "imgsz") or (k == "batch"))}, - "ultralytics_version": ultralytics.__version__, - "model_type": model_type, - } - except Exception: - model_artifacts = { - "names": class_names, - "yaml": model_instance.yaml, - "nc": nc, - "args": { - k: val for k, val in args.__dict__.items() if ((k == "model") or (k == "imgsz") or (k == "batch")) - }, - "ultralytics_version": ultralytics.__version__, - "model_type": model_type, - } - elif "yolov5" in model_type or "yolov7" in model_type or "yolov9" in model_type: - # parse from yaml for yolov5 - - with open(os.path.join(model_path, "opt.yaml")) as stream: - opts = yaml.safe_load(stream) - + model_artifacts: dict[str, Any] = { + "names": class_names, + "yaml": model_instance.yaml, + "nc": nc, + "args": _filtered_args(args), + "ultralytics_version": ultralytics.__version__, + "model_type": model_type, + } + else: + # yolov5 / yolov7 / yolov9 read their upload args from opt.yaml + opt_path = model_path / "opt.yaml" + if not opt_path.exists(): + raise MissingFileError(f"You must provide an opt.yaml file at '{opt_path}' for {model_type} uploads.") + with opt_path.open() as stream: + opts = yaml.safe_load(stream) or {} model_artifacts = { "names": class_names, "nc": model_instance.nc, - "args": { - "imgsz": opts["imgsz"] if "imgsz" in opts else opts["img_size"], - "batch": opts["batch_size"], - }, + "args": _legacy_yolo_args(opts, opt_path), "model_type": model_type, } if hasattr(model_instance, "yaml"): model_artifacts["yaml"] = model_instance.yaml - with open(os.path.join(model_path, "model_artifacts.json"), "w") as fp: - json.dump(model_artifacts, fp) - - torch.save(model_instance.state_dict(), os.path.join(model_path, "state_dict.pt")) - - list_files = [ - "results.csv", - "results.png", - "model_artifacts.json", - "state_dict.pt", - ] - - zip_file_name = "roboflow_deploy.zip" - - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in list_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write( - os.path.join(model_path, file), - arcname=file, - compress_type=zipfile.ZIP_DEFLATED, - ) - else: - if file in ["model_artifacts.json", "state_dict.pt"]: - raise (ValueError(f"File {file} not found. Please make sure to provide a valid model path.")) - - return zip_file_name, model_type + (build_dir / "model_artifacts.json").write_text(json.dumps(model_artifacts)) + torch.save(model_instance.state_dict(), build_dir / "state_dict.pt") + + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (model_path / "results.csv", "results.csv", False), + (model_path / "results.png", "results.png", False), + (build_dir / "model_artifacts.json", "model_artifacts.json", True), + (build_dir / "state_dict.pt", "state_dict.pt", True), + ], + ) + return archive_path, model_type, warnings -def _detect_rfdetr_task(checkpoint) -> Optional[str]: +def _detect_rfdetr_task(checkpoint: Any) -> str | None: """Detect the training task of an rf-detr checkpoint. rf-detr currently only supports weight upload for detection and instance @@ -328,83 +728,328 @@ def _detect_rfdetr_task(checkpoint) -> Optional[str]: model_name = checkpoint.get("model_name") if isinstance(model_name, str): return TASK_SEG if TASK_SEG in model_name.lower() else TASK_DET - args = checkpoint.get("args") - if args is None: + raw_args = checkpoint.get("args") + if raw_args is None: return None - seg_head = args.get("segmentation_head") if isinstance(args, dict) else getattr(args, "segmentation_head", None) - if seg_head is True: + args = raw_args if isinstance(raw_args, dict) else vars(raw_args) + segmentation_head = args.get("segmentation_head") + if segmentation_head is True: return TASK_SEG - if seg_head is False: + if segmentation_head is False: return TASK_DET return None -def _process_rfdetr(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - _supported_types = [ - # Detection models - "rfdetr-base", - "rfdetr-nano", - "rfdetr-small", - "rfdetr-medium", - "rfdetr-large", - "rfdetr-xlarge", - "rfdetr-2xlarge", - # Segmentation models - "rfdetr-seg-nano", - "rfdetr-seg-small", - "rfdetr-seg-medium", - "rfdetr-seg-large", - "rfdetr-seg-xlarge", - "rfdetr-seg-2xlarge", - ] - if model_type not in _supported_types: - raise ValueError(f"Model type {model_type} not supported. Supported types are {_supported_types}") - - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model path {model_path} does not exist.") - - model_files = os.listdir(model_path) - pt_file = next((f for f in model_files if f.endswith(".pt") or f.endswith(".pth")), None) - - if pt_file is None: - raise RuntimeError("No .pt or .pth model file found in the provided path") - - import torch - - checkpoint = torch.load(os.path.join(model_path, pt_file), map_location="cpu", weights_only=False) +def _rfdetr_checkpoint_pe_size(checkpoint: Any) -> int | None: + """Return an RF-DETR checkpoint's position-encoding grid size (tokens per side). - detected_task = _detect_rfdetr_task(checkpoint) - if detected_task: - implied_task = task_of_model_type(model_type) - if detected_task != implied_task: - raise ValueError( - f"model_type '{model_type}' implies task '{implied_task}' but the " - f".pt is a '{detected_task}' rfdetr checkpoint. Use a matching model_type." + Prefers the explicit ``positional_encoding_size`` arg, then ``resolution // + patch_size``, then derives it from the backbone ``position_embeddings`` tensor + (``grid² + 1`` tokens). Returns None when the geometry cannot be determined. + """ + if not isinstance(checkpoint, dict): + return None + raw_args = checkpoint.get("args") + if isinstance(raw_args, dict): + args = raw_args + elif raw_args is not None: + args = dict(vars(raw_args)) + else: + args = {} + + pe = args.get("positional_encoding_size") + if isinstance(pe, int) and pe > 0: + return pe + resolution = args.get("resolution") + patch_size = args.get("patch_size") + if isinstance(resolution, int) and isinstance(patch_size, int) and patch_size > 0: + return resolution // patch_size + + state_dict = checkpoint.get("model") + if isinstance(state_dict, dict): + for key, tensor in state_dict.items(): + if not key.endswith("position_embeddings"): + continue + shape = getattr(tensor, "shape", None) + if shape is not None and len(shape) == 3: + grid = math.isqrt(int(shape[1]) - 1) + if grid > 0 and grid * grid == int(shape[1]) - 1: + return grid + return None + + +def _resolve_rfdetr_variant( + model_type: str, + checkpoint: Any, + warnings: list[str], + allow_mismatch: bool = False, +) -> str: + """Check an RF-DETR model_type's size variant against the checkpoint geometry. + + Roboflow rebuilds the architecture from ``model_type`` at the variant's default + resolution before loading the weights, so a variant whose position-encoding grid + differs from the checkpoint fails conversion with a ``position_embeddings`` size + mismatch. Two cases: + + * The checkpoint's grid matches a *different* known variant (e.g. uploaded as + ``rfdetr-seg-nano`` but the grid is ``rfdetr-seg-small``) — a high-confidence + mislabel. Raise, naming the variant that fits, so the caller can confirm. + * The grid matches *no* known variant — likely a custom training resolution. We + cannot tell whether the backend supports it, so warn and proceed rather than + block a possibly-valid upload. + + ``allow_mismatch`` always proceeds with the requested variant (with a warning). + The detection-vs-segmentation task is handled separately and is not changed + here. Returns the resolved model_type. + """ + expected = RFDETR_POSITIONAL_ENCODING_SIZE.get(model_type) + actual = _rfdetr_checkpoint_pe_size(checkpoint) + if actual is None or expected is None or actual == expected: + return model_type + + task = task_of_model_type(model_type) + match = next( + ( + name + for name, grid in RFDETR_POSITIONAL_ENCODING_SIZE.items() + if grid == actual and task_of_model_type(name) == task + ), + None, + ) + + if match is not None and not allow_mismatch: + raise SizeMismatchError( + f"You specified model_type '{model_type}' (a {expected}x{expected} " + f"position-encoding grid), but the checkpoint was trained with " + f"{actual}x{actual}, which matches '{match}'. They don't match, so " + f"Roboflow's weight conversion would fail to load the backbone position " + f"embeddings. Upload as '{match}', or set allow_size_mismatch=True to " + f"upload '{model_type}' exactly as specified.", + requested=model_type, + detected=match, + ) + + if match is None: + warnings.append( + f"model_type '{model_type}' expects a {expected}x{expected} position-encoding " + f"grid, but the checkpoint is {actual}x{actual} and matches no known RF-DETR " + f"variant (it may use a custom training resolution). Packaging as " + f"'{model_type}'; Roboflow's conversion may reject it if it rebuilds at the " + f"variant's default resolution." + ) + else: + warnings.append( + f"model_type '{model_type}' expects a {expected}x{expected} grid, but the " + f"checkpoint is {actual}x{actual} (matches '{match}'). Packaging as " + f"'{model_type}' as requested." + ) + return model_type + + +def _find_rfdetr_checkpoint(model_path: Path, filename: str, warnings: list[str]) -> Path: + """Locate the rf-detr checkpoint, preserving the historical discovery behavior. + + The requested ``filename`` wins when it exists. Otherwise fall back to the + first top-level .pt/.pth file (sorted for determinism), which is how rf-detr + uploads located the checkpoint before ``filename`` was honored for them. + """ + requested_file = model_path / filename + if requested_file.exists(): + return requested_file + + discovered = sorted(path for path in model_path.iterdir() if path.is_file() and path.suffix in {".pt", ".pth"}) + if not discovered: + raise MissingFileError( + f"No .pt or .pth checkpoint found in '{model_path}' (and '{requested_file}' does not exist)." + ) + warnings.append( + f"Weights file '{requested_file}' was not found; using discovered checkpoint '{discovered[0].name}' instead." + ) + return discovered[0] + + +def _write_rfdetr_class_names(model_path: Path, build_dir: Path, checkpoint: Any) -> Path: + class_names_path = model_path / "class_names.txt" + if class_names_path.exists(): + class_names = class_names_path.read_text().splitlines() + else: + raw_args = checkpoint.get("args") if isinstance(checkpoint, dict) else None + if raw_args is None: + args: dict[str, Any] = {} + elif isinstance(raw_args, dict): + args = raw_args + else: + args = dict(vars(raw_args)) + class_names = args.get("class_names") or [] + if not class_names: + raise MissingFileError( + f"No class_names.txt file found in '{model_path}', and the RF-DETR " + "checkpoint does not include args with class_names. This should only " + "happen on rfdetr models trained before version 1.1.0. Create " + "class_names.txt with one class per line or retrain with a newer " + "rfdetr library." ) - get_classnames_txt_for_rfdetr(model_path, pt_file, checkpoint=checkpoint) + if "background_class83422" not in class_names: + class_names = ["background_class83422", *class_names] + output_path = build_dir / "class_names.txt" + output_path.write_text("\n".join(class_names) + "\n") + return output_path + + +def _process_rfdetr( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + if model_type not in SUPPORTED_RFDETR_TYPES: + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for RF-DETR upload. " + f"Supported types are: {', '.join(SUPPORTED_RFDETR_TYPES)}." + ) + torch = _import_required_module("torch", "pip install torch") + warnings: list[str] = [] - # Copy the .pt file to weights.pt if not already named weights.pt - if pt_file != "weights.pt": - shutil.copy(os.path.join(model_path, pt_file), os.path.join(model_path, "weights.pt")) + checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings) + checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu") - required_files = ["weights.pt"] + detected_task = _detect_rfdetr_task(checkpoint) + if detected_task and detected_task != task_of_model_type(model_type): + raise TaskMismatchError( + f"model_type '{model_type}' implies task '{task_of_model_type(model_type)}', " + f"but the checkpoint is a '{detected_task}' RF-DETR model. Use a matching model_type." + ) + + model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) + + shutil.copy(checkpoint_path, build_dir / "weights.pt") + _write_rfdetr_class_names(model_path, build_dir, checkpoint) + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (build_dir / "weights.pt", "weights.pt", True), + (model_path / "results.csv", "results.csv", False), + (model_path / "results.png", "results.png", False), + (model_path / "model_artifacts.json", "model_artifacts.json", False), + (build_dir / "class_names.txt", "class_names.txt", False), + ], + ) + return archive_path, model_type, warnings - optional_files = ["results.csv", "results.png", "model_artifacts.json", "class_names.txt"] - zip_file_name = "roboflow_deploy.zip" - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in required_files: - zipMe.write(os.path.join(model_path, file), arcname=file, compress_type=zipfile.ZIP_DEFLATED) +def _process_huggingface( + model_type: str, + model_path: Path, + build_dir: Path, +) -> tuple[Path, str, list[str]]: + if model_type not in SUPPORTED_HUGGINGFACE_TYPES: + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for this type of upload. " + f"Supported types are: {', '.join(SUPPORTED_HUGGINGFACE_TYPES)}." + ) + + model_files = [path for path in model_path.iterdir() if path.is_file()] + safetensors_files = [path for path in model_files if path.suffix == ".safetensors"] + npz_file = next((path for path in model_files if path.suffix == ".npz"), None) + if safetensors_files: + required = { + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer_config.json", + "tokenizer.json", + } + missing = sorted(required - {path.name for path in model_files}) + if missing: + raise MissingFileError(f"Missing files required for a PyTorch {model_type} upload: {', '.join(missing)}.") + files_to_deploy = model_files + elif npz_file is not None: + files_to_deploy = [npz_file] + else: + raise MissingFileError(f"No .npz or .safetensors model file found in '{model_path}'.") + + archive_path = build_dir / "roboflow_deploy.tar" + with tarfile.open(archive_path, "w") as tar: + for path in files_to_deploy: + tar.add(path, arcname=path.name) + return archive_path, model_type, [] + + +def _process_yolonas( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, +) -> tuple[Path, str, list[str]]: + torch = _import_required_module("torch", "pip install torch") + weights_path = model_path / filename + if not weights_path.exists(): + raise MissingFileError(f"Model weights file '{weights_path}' was not found.") + + checkpoint = _load_checkpoint(torch, weights_path, map_location="cpu") + class_names = checkpoint["processing_params"]["class_names"] + opt_path = model_path / "opt.yaml" + if not opt_path.exists(): + raise MissingFileError( + f"You must create an opt.yaml file at '{opt_path}' of the format:\n" + f"imgsz: \n" + f"batch_size: \n" + f"architecture: \n" + ) + with opt_path.open() as stream: + opts = yaml.safe_load(stream) or {} + missing = [key for key in ("imgsz", "batch_size", "architecture") if key not in opts] + if missing: + raise ModelPackagingError(f"{opt_path} lacks required keys: {', '.join(missing)}.") + + model_artifacts = { + "names": class_names, + "nc": len(class_names), + "args": { + "imgsz": opts["imgsz"], + "batch": opts["batch_size"], + "architecture": opts["architecture"], + }, + "model_type": model_type, + } + (build_dir / "model_artifacts.json").write_text(json.dumps(model_artifacts)) + shutil.copy(weights_path, build_dir / "state_dict.pt") + + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (model_path / "results.json", "results.json", False), + (model_path / "results.png", "results.png", False), + (build_dir / "model_artifacts.json", "model_artifacts.json", True), + (build_dir / "state_dict.pt", "state_dict.pt", True), + ], + ) + return archive_path, model_type, [] - for file in optional_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write(os.path.join(model_path, file), arcname=file, compress_type=zipfile.ZIP_DEFLATED) - return zip_file_name, model_type +def _write_zip( + archive_path: Path, + files: list[tuple[Path, str, bool]], +) -> None: + with zipfile.ZipFile(archive_path, "w") as zip_file: + for path, arcname, required in files: + if path.exists(): + zip_file.write(path, arcname=arcname, compress_type=zipfile.ZIP_DEFLATED) + elif required: + raise MissingFileError(f"Required upload artifact '{path}' was not found.") def get_classnames_txt_for_rfdetr(model_path: str, pt_file: str, checkpoint=None): + """Legacy rf-detr class-names helper, kept for backwards compatibility. + + Writes (and mutates) ``class_names.txt`` inside ``model_path``. The packaging + flow uses :func:`_write_rfdetr_class_names` instead, which leaves the source + directory untouched. + """ class_names_path = os.path.join(model_path, "class_names.txt") if os.path.exists(class_names_path): maybe_prepend_dummy_class(class_names_path) @@ -424,7 +1069,7 @@ def get_classnames_txt_for_rfdetr(model_path: str, pt_file: str, checkpoint=None maybe_prepend_dummy_class(class_names_path) return class_names_path - raise FileNotFoundError( + raise MissingFileError( f"No class_names.txt file found in model path {model_path}.\n" f"This should only happen on rfdetr models trained before version 1.1.0.\n" f"Please re-train your model with the latest version of the rfdetr library, or\n" @@ -442,122 +1087,3 @@ def maybe_prepend_dummy_class(class_name_file: str): class_names.insert(0, dummy_class) with open(class_name_file, "w") as f: f.writelines(class_names) - - -def _process_huggingface( - model_type: str, model_path: str, filename: str = "fine-tuned-paligemma-3b-pt-224.f16.npz" -) -> tuple[str, str]: - # Check if model_path exists - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model path {model_path} does not exist.") - model_files = os.listdir(model_path) - print(f"Model files found in {model_path}: {model_files}") - - files_to_deploy = [] - - # Find first .npz file in model_path - npz_filename = next((file for file in model_files if file.endswith(".npz")), None) - if any([file.endswith(".safetensors") for file in model_files]): - print(f"Found .safetensors file in model path. Deploying PyTorch {model_type} model.") - necessary_files = [ - "preprocessor_config.json", - "special_tokens_map.json", - "tokenizer_config.json", - "tokenizer.json", - ] - for file in necessary_files: - if file not in model_files: - print("Missing necessary file", file) - res = input("Do you want to continue? (y/n)") - if res.lower() != "y": - exit(1) - for file in model_files: - files_to_deploy.append(file) - elif npz_filename is not None: - print(f"Found .npz file {npz_filename} in model path. Deploying JAX PaliGemma model.") - files_to_deploy.append(npz_filename) - else: - raise FileNotFoundError(f"No .npz or .safetensors file found in model path {model_path}.") - - if len(files_to_deploy) == 0: - raise FileNotFoundError(f"No valid files found in model path {model_path}.") - print(f"Zipping files for deploy: {files_to_deploy}") - - import tarfile - - tar_file_name = "roboflow_deploy.tar" - - with tarfile.open(os.path.join(model_path, tar_file_name), "w") as tar: - for file in files_to_deploy: - tar.add(os.path.join(model_path, file), arcname=file) - - print("Uploading to Roboflow... May take several minutes.") - - return tar_file_name, model_type - - -def _process_yolonas(model_type: str, model_path: str, filename: str = "weights/best.pt") -> tuple[str, str]: - try: - import torch - except ImportError: - raise RuntimeError( - "The torch python package is required to deploy yolonas models. Please install it with `pip install torch`" - ) - - model = torch.load(os.path.join(model_path, filename), map_location="cpu") - class_names = model["processing_params"]["class_names"] - - opt_path = os.path.join(model_path, "opt.yaml") - if not os.path.exists(opt_path): - raise RuntimeError( - f"You must create an opt.yaml file at {os.path.join(model_path, '')} of the format:\n" - f"imgsz: \n" - f"batch_size: \n" - f"architecture: \n" - ) - with open(os.path.join(model_path, "opt.yaml")) as stream: - opts = yaml.safe_load(stream) - required_keys = ["imgsz", "batch_size", "architecture"] - for key in required_keys: - if key not in opts: - raise RuntimeError(f"{opt_path} lacks required key {key}. Required keys: {required_keys}") - - model_artifacts = { - "names": class_names, - "nc": len(class_names), - "args": { - "imgsz": opts["imgsz"] if "imgsz" in opts else opts["img_size"], - "batch": opts["batch_size"], - "architecture": opts["architecture"], - }, - "model_type": model_type, - } - - with open(os.path.join(model_path, "model_artifacts.json"), "w") as fp: - json.dump(model_artifacts, fp) - - shutil.copy(os.path.join(model_path, filename), os.path.join(model_path, "state_dict.pt")) - - list_files = [ - "results.json", - "results.png", - "model_artifacts.json", - "state_dict.pt", - ] - - zip_file_name = "roboflow_deploy.zip" - - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in list_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write( - os.path.join(model_path, file), - arcname=file, - compress_type=zipfile.ZIP_DEFLATED, - ) - else: - if file in ["model_artifacts.json", filename]: - raise (ValueError(f"File {file} not found. Please make sure to provide a valid model path.")) - - return zip_file_name, model_type diff --git a/setup.py b/setup.py index 85f671ab..49278e77 100644 --- a/setup.py +++ b/setup.py @@ -31,6 +31,11 @@ extras_require={ "desktop": ["opencv-python==4.8.0.74"], "dev": [ + # numpy 2.5.0 (2026-06-21) ships PEP 695 `type` statement stubs that + # mypy cannot parse under python_version = "3.10". Cap in the dev + # extra only so typecheck CI keeps passing; runtime installs + # (requirements.txt) stay unconstrained for end users. See #498. + "numpy<2.5", "mypy", "responses", "ruff", diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 80408602..50f14a86 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -1,14 +1,36 @@ +import json import os +import sys import tempfile +import types import unittest +import zipfile +from pathlib import Path from types import SimpleNamespace +from unittest import mock from roboflow.config import TASK_CLS, TASK_DET, TASK_OBB, TASK_POSE, TASK_SEG, TASK_SEM +from roboflow.util import model_processor from roboflow.util.model_processor import ( + MissingFileError, + ModelPackagingError, + SizeMismatchError, + TaskMismatchError, + UnsupportedModelError, _detect_rfdetr_task, _detect_yolo_task, + _infer_yolo_size, + _legacy_yolo_args, + _resolve_rfdetr_variant, + _resolve_yolo_size, + _rfdetr_checkpoint_pe_size, + _write_rfdetr_class_names, get_classnames_txt_for_rfdetr, + package_custom_weights, + package_custom_weights_interactive, + process, task_of_model_type, + validate_model_type_for_project, ) @@ -110,5 +132,405 @@ def test_namespace_args(self): ) +class ValidateModelTypeForProjectTest(unittest.TestCase): + def test_rejects_detection_for_classification(self): + with self.assertRaises(TaskMismatchError) as ctx: + validate_model_type_for_project("yolov8", "classification", "widgets") + self.assertIn("classification", str(ctx.exception)) + self.assertIn("task 'cls'", str(ctx.exception)) + + def test_task_mismatch_is_a_value_error(self): + # Callers that caught the historical ValueError keep working. + with self.assertRaises(ValueError): + validate_model_type_for_project("yolov8", "classification", "widgets") + + def test_unknown_project_type_is_ignored(self): + validate_model_type_for_project("yolov8", "some-new-type", "widgets") + + +class LegacyYoloArgsTest(unittest.TestCase): + def test_reports_missing_batch_size(self): + with self.assertRaises(ModelPackagingError) as ctx: + _legacy_yolo_args({"imgsz": 640}, Path("opt.yaml")) + self.assertIn("batch_size", str(ctx.exception)) + + def test_reports_missing_image_size(self): + with self.assertRaises(ModelPackagingError) as ctx: + _legacy_yolo_args({"batch_size": 8}, Path("opt.yaml")) + self.assertIn("imgsz", str(ctx.exception)) + + def test_accepts_either_image_size_key(self): + self.assertEqual(_legacy_yolo_args({"imgsz": 640, "batch_size": 8}, Path("x")), {"imgsz": 640, "batch": 8}) + self.assertEqual(_legacy_yolo_args({"img_size": 416, "batch_size": 4}, Path("x")), {"imgsz": 416, "batch": 4}) + + +class _FakeYoloModel: + def __init__(self, yaml): + self.yaml = yaml + + +class InferYoloSizeTest(unittest.TestCase): + def test_from_depth_width_multiples(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_infer_yolo_size(model), "n") + + def test_explicit_scale_letter_wins(self): + model = _FakeYoloModel({"scale": "m", "depth_multiple": 0.67, "width_multiple": 0.75}) + self.assertEqual(_infer_yolo_size(model), "m") + + def test_unknown_returns_none(self): + self.assertIsNone(_infer_yolo_size(_FakeYoloModel({}))) + + +class ResolveYoloSizeTest(unittest.TestCase): + def test_fills_bare_family_from_architecture(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8", model, warnings), "yolov8n") + self.assertTrue(warnings and "Inferred model size 'yolov8n'" in warnings[0]) + + def test_preserves_task_suffix_when_filling(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.50}) + self.assertEqual(_resolve_yolo_size("yolov8-seg", model, []), "yolov8s-seg") + + def test_raises_when_size_cannot_be_inferred(self): + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_yolo_size("yolov8", _FakeYoloModel({}), []) + self.assertIn("could not be inferred", str(ctx.exception)) + self.assertIn("yolov8n", str(ctx.exception)) + + def test_raises_on_declared_size_conflict(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_yolo_size("yolov8m", model, []) + self.assertIn("yolov8n", str(ctx.exception)) + self.assertIn("allow_size_mismatch=True", str(ctx.exception)) + self.assertEqual(ctx.exception.requested, "yolov8m") + self.assertEqual(ctx.exception.detected, "yolov8n") + + def test_allow_mismatch_keeps_declared_size(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8m", model, warnings, allow_mismatch=True), "yolov8m") + self.assertTrue(warnings and "as requested" in warnings[0]) + + def test_keeps_user_size_when_not_inferable(self): + warnings: list = [] + self.assertEqual(_resolve_yolo_size("yolov8m", _FakeYoloModel({}), warnings), "yolov8m") + self.assertEqual(warnings, []) + + def test_keeps_matching_sized_type_without_warning(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8n", model, warnings), "yolov8n") + self.assertEqual(warnings, []) + + +class ResolveRfdetrVariantTest(unittest.TestCase): + def test_raises_on_size_conflict_naming_the_fit(self): + checkpoint = {"args": {"resolution": 384, "patch_size": 12}} + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_rfdetr_variant("rfdetr-seg-nano", checkpoint, []) + self.assertIn("rfdetr-seg-small", str(ctx.exception)) + self.assertIn("32x32", str(ctx.exception)) + self.assertIn("allow_size_mismatch=True", str(ctx.exception)) + + def test_allow_mismatch_keeps_requested_variant(self): + warnings: list = [] + checkpoint = {"args": {"resolution": 384, "patch_size": 12}} + resolved = _resolve_rfdetr_variant("rfdetr-seg-nano", checkpoint, warnings, allow_mismatch=True) + self.assertEqual(resolved, "rfdetr-seg-nano") + self.assertTrue(warnings and "as requested" in warnings[0]) + + def test_keeps_matching_grid_without_warning(self): + warnings: list = [] + checkpoint = {"args": {"positional_encoding_size": 32}} + self.assertEqual(_resolve_rfdetr_variant("rfdetr-seg-small", checkpoint, warnings), "rfdetr-seg-small") + self.assertEqual(warnings, []) + + def test_warns_but_allows_custom_resolution(self): + warnings: list = [] + resolved = _resolve_rfdetr_variant("rfdetr-seg-nano", {"args": {"positional_encoding_size": 99}}, warnings) + self.assertEqual(resolved, "rfdetr-seg-nano") + self.assertTrue(warnings and "matches no known RF-DETR variant" in warnings[0]) + + def test_pe_size_derived_from_position_embeddings_tensor(self): + class FakeTensor: + shape = (1, 1025, 384) + + checkpoint = {"model": {"backbone.0.encoder.encoder.embeddings.position_embeddings": FakeTensor()}} + self.assertEqual(_rfdetr_checkpoint_pe_size(checkpoint), 32) + + +class WriteRfdetrClassNamesTest(unittest.TestCase): + def test_fails_cleanly_without_checkpoint_args(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(MissingFileError) as ctx: + _write_rfdetr_class_names(Path(tmp), Path(tmp), checkpoint={}) + self.assertIn("does not include args with class_names", str(ctx.exception)) + + def test_does_not_mutate_existing_class_names_file(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as build: + source = Path(tmp) / "class_names.txt" + source.write_text("cat\ndog\n") + output = _write_rfdetr_class_names(Path(tmp), Path(build), checkpoint={}) + self.assertEqual(source.read_text(), "cat\ndog\n") + self.assertEqual( + output.read_text().splitlines(), + ["background_class83422", "cat", "dog"], + ) + + +def _fake_torch(load_result, calls=None): + module = types.ModuleType("torch") + + def load(path, **kwargs): + if calls is not None: + calls.append((Path(path), kwargs)) + return load_result + + def save(obj, path): + Path(path).write_bytes(b"fake-state-dict") + + module.load = load + module.save = save + return module + + +def _import_patch(modules): + def _import(module_name, install_hint): + return modules[module_name] + + return mock.patch.object(model_processor, "_import_required_module", side_effect=_import) + + +def _write_yolonas_inputs(model_dir: Path): + weights = model_dir / "weights" / "best.pt" + weights.parent.mkdir() + weights.write_bytes(b"checkpoint") + (model_dir / "opt.yaml").write_text("imgsz: 640\nbatch_size: 8\narchitecture: yolo_nas_s\n") + + +class PackageCustomWeightsTest(unittest.TestCase): + """Contract tests for the public non-interactive helper.""" + + def _package_yolonas(self, model_dir: Path, **kwargs): + calls: list = [] + torch = _fake_torch({"processing_params": {"class_names": ["widget"]}}, calls) + with _import_patch({"torch": torch}): + bundle = package_custom_weights("yolonas", str(model_dir), **kwargs) + return bundle, calls + + def test_never_prompts_or_exits(self): + prompt_guard = mock.patch( + "builtins.input", side_effect=AssertionError("package_custom_weights must not prompt") + ) + exit_guard = mock.patch.object(sys, "exit", side_effect=AssertionError("package_custom_weights must not exit")) + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + with prompt_guard, exit_guard: + bundle, calls = self._package_yolonas(model_dir) + try: + self.assertTrue(bundle.archive_path.exists()) + self.assertEqual(calls[0][1], {"weights_only": False, "map_location": "cpu"}) + finally: + bundle.cleanup() + + def test_does_not_write_into_model_path(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + before = sorted(path for path in model_dir.rglob("*")) + bundle, _ = self._package_yolonas(model_dir) + try: + self.assertEqual(sorted(path for path in model_dir.rglob("*")), before) + self.assertNotEqual(bundle.build_dir, model_dir) + self.assertTrue(bundle.owns_build_dir) + finally: + bundle.cleanup() + self.assertFalse(bundle.build_dir.exists()) + + def test_explicit_build_dir_is_used_and_not_cleaned_up(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as build: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + bundle, _ = self._package_yolonas(model_dir, build_dir=build) + self.assertEqual(bundle.build_dir, Path(build)) + self.assertFalse(bundle.owns_build_dir) + bundle.cleanup() + self.assertTrue(bundle.archive_path.exists()) + + def test_owned_build_dir_is_removed_on_failure(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as fake_build: + with mock.patch.object(model_processor.tempfile, "mkdtemp", return_value=fake_build): + with self.assertRaises(UnsupportedModelError): + package_custom_weights("not-a-model", tmp) + self.assertFalse(Path(fake_build).exists()) + + def test_missing_model_path_raises_missing_file(self): + with self.assertRaises(MissingFileError): + package_custom_weights("yolonas", "/nonexistent/path/for/test") + + def test_family_must_be_a_prefix_not_a_substring(self): + # 'foo-yolov8n' merely contains a family token; the backend would + # reject it after upload, so the gate must reject it up front. + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("foo-yolov8n", tmp) + + def test_rfdetr_falls_back_to_discovered_checkpoint(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "other.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + bundle = package_custom_weights("rfdetr-base", str(model_dir)) + try: + self.assertTrue(any("other.pt" in warning for warning in bundle.warnings)) + with zipfile.ZipFile(bundle.archive_path) as archive: + self.assertIn("weights.pt", archive.namelist()) + self.assertIn("class_names.txt", archive.namelist()) + finally: + bundle.cleanup() + + def test_rfdetr_without_any_checkpoint_raises(self): + with tempfile.TemporaryDirectory() as tmp: + torch = _fake_torch({}) + with _import_patch({"torch": torch}): + with self.assertRaises(MissingFileError): + package_custom_weights("rfdetr-base", tmp) + + def test_yolov8_full_flow_builds_artifacts(self): + checkpoint_names = {1: "dog", 0: "cat"} + + class DetectionModel: + names = checkpoint_names + nc = 2 + yaml = {"nc": 2, "depth_multiple": 0.33, "width_multiple": 0.25} + args = {"model": "yolov8n.yaml", "imgsz": 640, "batch": 16, "lr0": 0.01} + + def state_dict(self): + return {"weight": b"w"} + + fake_ultralytics = types.ModuleType("ultralytics") + fake_ultralytics.__version__ = "8.0.196" + fake_torch = _fake_torch({"model": DetectionModel()}) + + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + weights = model_dir / "weights" / "best.pt" + weights.parent.mkdir() + weights.write_bytes(b"checkpoint") + + with ( + _import_patch({"torch": fake_torch, "ultralytics": fake_ultralytics}), + mock.patch.dict(sys.modules, {"ultralytics": fake_ultralytics}), + ): + bundle = package_custom_weights("yolov8", str(model_dir)) + try: + self.assertEqual(bundle.model_type, "yolov8n") + self.assertTrue(any("Inferred model size 'yolov8n'" in warning for warning in bundle.warnings)) + with zipfile.ZipFile(bundle.archive_path) as archive: + artifacts = json.loads(archive.read("model_artifacts.json")) + self.assertIn("state_dict.pt", archive.namelist()) + self.assertEqual(artifacts["names"], ["cat", "dog"]) + self.assertEqual(artifacts["model_type"], "yolov8n") + self.assertEqual(artifacts["ultralytics_version"], "8.0.196") + self.assertEqual(artifacts["args"], {"model": "yolov8n.yaml", "imgsz": 640, "batch": 16}) + finally: + bundle.cleanup() + + +class ProcessCompatTest(unittest.TestCase): + """The legacy process() entry point keeps its historical contract.""" + + def test_packages_into_model_path_and_returns_tuple(self): + calls: list = [] + torch = _fake_torch({"processing_params": {"class_names": ["widget"]}}, calls) + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + with _import_patch({"torch": torch}): + zip_file_name, model_type = process("yolonas", str(model_dir), "weights/best.pt") + + self.assertEqual(zip_file_name, "roboflow_deploy.zip") + self.assertEqual(model_type, "yolonas") + # Historical side effects: artifacts and archive land in model_path. + self.assertTrue((model_dir / "roboflow_deploy.zip").exists()) + self.assertTrue((model_dir / "model_artifacts.json").exists()) + self.assertTrue((model_dir / "state_dict.pt").exists()) + + def test_prompts_and_retries_on_mismatch_like_before(self): + error = model_processor.DependencyMismatchError( + "wrong ultralytics", + dependency="ultralytics", + required="ultralytics==8.0.196", + installed="8.3.0", + ) + bundle = model_processor.ModelUploadBundle( + archive_path=Path("roboflow_deploy.zip"), + build_dir=Path("."), + model_type="yolov8n", + ) + outcomes = [error, bundle] + + def fake_package(*args, **kwargs): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + self.assertTrue(kwargs["allow_dependency_mismatch"]) + return outcome + + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=fake_package), + mock.patch("builtins.input", return_value="y"), + mock.patch("builtins.print"), + ): + zip_file_name, model_type = process("yolov8m", "/models", "weights/best.pt") + + self.assertEqual(zip_file_name, "roboflow_deploy.zip") + self.assertEqual(model_type, "yolov8n") + + +class PackageCustomWeightsInteractiveTest(unittest.TestCase): + def _bundle(self): + return model_processor.ModelUploadBundle( + archive_path=Path("roboflow_deploy.zip"), + build_dir=Path("."), + model_type="yolov8n", + warnings=("some warning",), + ) + + def test_retries_with_size_override_on_confirmation(self): + error = SizeMismatchError("size conflict", requested="yolov8m", detected="yolov8n") + outcomes = [error, self._bundle()] + + def fake_package(*args, **kwargs): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + self.assertTrue(kwargs["allow_size_mismatch"]) + return outcome + + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=fake_package), + mock.patch("builtins.input", return_value="y"), + mock.patch("builtins.print"), + ): + bundle = package_custom_weights_interactive("yolov8m", "/models") + self.assertEqual(bundle.model_type, "yolov8n") + + def test_reraises_when_user_declines(self): + error = SizeMismatchError("size conflict", requested="yolov8m") + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=error), + mock.patch("builtins.input", return_value="n"), + mock.patch("builtins.print"), + ): + with self.assertRaises(SizeMismatchError): + package_custom_weights_interactive("yolov8m", "/models") + + if __name__ == "__main__": unittest.main()