diff --git a/roboflow/__init__.py b/roboflow/__init__.py index 056ae873..260c1261 100644 --- a/roboflow/__init__.py +++ b/roboflow/__init__.py @@ -21,7 +21,7 @@ CLIPModel = None # type: ignore[assignment,misc] GazeModel = None # type: ignore[assignment,misc] -__version__ = "1.3.12" +__version__ = "1.3.13" def check_key(api_key, model, notebook, num_retries=0): 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 10ae4bd1..46389b35 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, PurePosixPath, PureWindowsPath +from typing import Any import yaml @@ -19,14 +51,46 @@ 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", +) -# Minimum rf-detr release shipping `RFDETR.export_for_roboflow`. +# Minimum rf-detr release shipping ``RFDETR.export_for_roboflow`` (used to rebuild +# an upload bundle from a raw PyTorch-Lightning checkpoint). RFDETR_MIN_VERSION = "1.8.0" # rf-detr model_type -> RFDETR subclass name. Single source of truth for both the -# supported-type check and the `from_checkpoint` fallback (used when a raw checkpoint -# lacks the metadata rf-detr needs to infer its own class). +# supported-type check and the ``from_checkpoint`` fallback (used when a raw +# checkpoint lacks the metadata rf-detr needs to infer its own class). _RFDETR_MODEL_TYPE_TO_CLASS = { # Detection "rfdetr-base": "RFDETRBase", @@ -45,6 +109,131 @@ "rfdetr-seg-2xlarge": "RFDETRSeg2XLarge", } +SUPPORTED_RFDETR_TYPES = tuple(_RFDETR_MODEL_TYPE_TO_CLASS) + +DEFAULT_WEIGHTS_FILENAME = "weights/best.pt" + +# 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: """Canonical task for a deploy model_type string. @@ -59,8 +248,53 @@ def task_of_model_type(model_type: str) -> str: return TASK_DET +def _checkpoint_args_as_dict(raw_args: Any) -> dict[str, Any]: + """Coerce a checkpoint's ``args`` (dict, argparse.Namespace, or None) to a dict. + + A corrupt checkpoint may store ``args`` as a scalar or list; those have no + meaningful attributes, so treat them as no-args rather than letting + ``vars()`` raise a bare ``TypeError`` outside the ModelPackagingError contract. + """ + if isinstance(raw_args, dict): + return raw_args + if hasattr(raw_args, "__dict__"): + return dict(vars(raw_args)) + return {} + + +def _resolve_within_source(source_dir: Path, filename: str) -> Path: + """Resolve ``filename`` against ``source_dir``, refusing to escape it. + + ``filename`` is documented as relative to ``model_path`` and is forwarded + verbatim by hosted callers (the MCP server). An absolute path or ``..`` + segments would let a caller read weights from outside the model directory, + so reject both instead of silently packaging a file the caller never + intended. ``source_dir`` is already resolved; resolving the join collapses + ``..`` and follows symlinks before the containment check. + + Absoluteness is tested under both OS conventions, not just the host's: the + hosted MCP forwards caller-supplied paths, so a POSIX absolute like + ``/etc/passwd`` must be rejected even when packaging happens to run on + Windows (where ``Path.is_absolute()`` alone would miss it, and vice versa). + """ + if PurePosixPath(filename).is_absolute() or PureWindowsPath(filename).is_absolute(): + raise ModelPackagingError(f"filename '{filename}' must be a path relative to model_path, not an absolute path.") + resolved = (source_dir / filename).resolve() + if resolved != source_dir and source_dir not in resolved.parents: + raise ModelPackagingError( + f"filename '{filename}' resolves outside model_path '{source_dir}'. " + "It must point to a checkpoint inside the model directory." + ) + # '' / '.' resolve to model_path itself and a subdirectory passes the + # containment check; all would reach torch.load() and leak a raw + # IsADirectoryError outside the ModelPackagingError contract. + if resolved.is_dir(): + raise ModelPackagingError(f"filename '{filename}' must point to a checkpoint file, not a directory.") + return resolved + + 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, @@ -72,70 +306,208 @@ 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 = DEFAULT_WEIGHTS_FILENAME, + *, + 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.") + _resolve_within_source(source_dir, filename) + + 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).expanduser().resolve() + 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 = DEFAULT_WEIGHTS_FILENAME, + *, + 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. - 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 + 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 "yolonas" in model_type: - return _process_yolonas + if model_type.startswith(("paligemma", "paligemma2", "florence-2")): + 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, + ) - if "rfdetr" in model_type: - return _process_rfdetr - return _process_yolo +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) -> Optional[str]: +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 @@ -149,87 +521,221 @@ 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]: + # A corrupt checkpoint may store args as a scalar/None; coerce via the shared + # helper so it becomes {} rather than raising a raw TypeError from vars(). + return {k: v for k, v in _checkpoint_args_as_dict(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. + + 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 - elif "yolov12" in model_type: - try: - import torch - import ultralytics - 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: + if allow_mismatch: + warnings.append( + f"Could not infer a size for '{model_type}' from the checkpoint; " + f"uploading the bare family name as requested. Roboflow may reject it " + f"if it requires an explicit size." + ) + return model_type + 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 + + 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( - "\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" + +def _require_model_attr(model_instance: Any, attr: str, model_type: str) -> Any: + """Return ``model_instance.`` or raise a ModelPackagingError. + + Roboflow's server-side conversion needs these fields; a stripped checkpoint + missing one would otherwise raise a raw ``AttributeError`` outside the + ModelPackagingError contract (an opaque 500 for hosted callers). + """ + value = getattr(model_instance, attr, None) + if value is None: + raise ModelPackagingError( + f"The {model_type} checkpoint's model is missing '{attr}'; it does not look " + "like a complete Ultralytics training checkpoint. Re-export it from your training run." + ) + return value + + +def _require_checkpoint_field(checkpoint: Any, key: str, model_type: str) -> Any: + """Return ``checkpoint[key]`` or raise a ModelPackagingError (see _require_model_attr).""" + if not isinstance(checkpoint, dict) or key not in checkpoint: + raise ModelPackagingError( + f"The {model_type} checkpoint is missing '{key}'; it does not look like a " + "complete Ultralytics training checkpoint. Re-export it from your training run." ) + return checkpoint[key] - 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: @@ -237,28 +743,16 @@ 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") + model_yaml = _require_model_attr(model_instance, "yaml", model_type) if ( "-cls" in model_type or model_type.startswith("yolov10") @@ -266,79 +760,57 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, or model_type.startswith("yolov12") or model_type.startswith("yolo26") ): - nc = model_instance.yaml["nc"] - args = model["train_args"] + if not isinstance(model_yaml, dict) or "nc" not in model_yaml: + raise ModelPackagingError( + f"The {model_type} checkpoint's model config (model.yaml) is missing 'nc'; " + "it does not look like a complete Ultralytics training checkpoint." + ) + nc = model_yaml["nc"] + args = _require_checkpoint_field(checkpoint, "train_args", model_type) 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) - + nc = _require_model_attr(model_instance, "nc", model_type) + args = _require_model_attr(model_instance, "args", model_type) + model_artifacts: dict[str, Any] = { + "names": class_names, + "yaml": model_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"], - }, + "nc": _require_model_attr(model_instance, "nc", model_type), + "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 @@ -352,120 +824,418 @@ def _detect_rfdetr_task(checkpoint) -> Optional[str]: model_name = checkpoint.get("model_name") if isinstance(model_name, str): name = model_name.lower() - # Keypoint rf-detr checkpoints (e.g. 'RFDETRKeypointPreview') are not a supported - # upload type; classify them as pose so the model_type task check below rejects them - # instead of silently uploading a keypoint model as object detection. + # Keypoint rf-detr checkpoints (e.g. 'RFDETRKeypointPreview') are not a + # supported upload type; classify them as pose so the task check rejects + # them instead of silently uploading a keypoint model as detection. if "keypoint" in name: return TASK_POSE return TASK_SEG if TASK_SEG in name 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 = _checkpoint_args_as_dict(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 _is_ptl_checkpoint(checkpoint) -> bool: - """True if `checkpoint` is a raw PyTorch-Lightning rf-detr checkpoint dict.""" - return isinstance(checkpoint, dict) and "pytorch-lightning_version" in checkpoint +def _rfdetr_checkpoint_pe_size(checkpoint: Any) -> int | None: + """Return an RF-DETR checkpoint's position-encoding grid size (tokens per side). + 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 + args = _checkpoint_args_as_dict(checkpoint.get("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 _require_rfdetr(): - """Lazily import `rfdetr` and verify it ships the upload-bundle helpers. - Raises a RuntimeError with an actionable hint if rfdetr is missing or too old. +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. """ - try: - import rfdetr - except ImportError: - raise RuntimeError( - "rfdetr is required to upload PyTorch-Lightning rf-detr checkpoints. " - f"Please install it with `pip install 'rfdetr>={RFDETR_MIN_VERSION}'`." + 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 not hasattr(rfdetr.RFDETR, "export_for_roboflow"): - raise RuntimeError( - "The installed rfdetr is too old to upload PyTorch-Lightning rf-detr checkpoints. " - f"Please upgrade it with `pip install --upgrade 'rfdetr>={RFDETR_MIN_VERSION}'`." + 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 - return rfdetr +def _find_rfdetr_checkpoint(model_path: Path, filename: str, warnings: list[str]) -> Path: + """Locate the rf-detr checkpoint. -def _process_rfdetr(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - _supported_types = list(_RFDETR_MODEL_TYPE_TO_CLASS) - if model_type not in _supported_types: - raise ValueError(f"Model type {model_type} not supported. Supported types are {_supported_types}") + An explicitly-requested ``filename`` (anything other than the default) must + exist: falling back to a different checkpoint on a typo would silently + package the wrong weights. Only the default path falls back to discovering + the first top-level .pt/.pth file (sorted for determinism), preserving how + rf-detr uploads located the checkpoint before ``filename`` was honored. + """ + requested_file = model_path / filename + if requested_file.exists(): + return requested_file + + if filename != DEFAULT_WEIGHTS_FILENAME: + raise MissingFileError( + f"RF-DETR weights file '{requested_file}' was not found. Set filename to the " + f"checkpoint's exact .pt or .pth path relative to model_path." + ) - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model path {model_path} does not exist.") + 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)." + ) + if len(discovered) > 1: + others = ", ".join(path.name for path in discovered) + warnings.append( + f"Weights file '{requested_file}' was not found and '{model_path}' holds multiple " + f"checkpoints ({others}); packaging '{discovered[0].name}'. Set filename to pick a " + "specific checkpoint if that is not the one you want." + ) + else: + warnings.append( + f"Weights file '{requested_file}' was not found; using discovered checkpoint " + f"'{discovered[0].name}' instead." + ) + return discovered[0] - model_files = os.listdir(model_path) - pt_file = next((f for f in model_files if f.endswith(".pt") or f.endswith(".pth")), None) - # Honor the caller-selected `filename` when it points at a valid checkpoint — rf-detr - # training dirs routinely hold several (`checkpoint.pth`, `checkpoint_best_ema.pth`, - # `weights.pt`, ...), so first-match discovery could upload the wrong one. Discovery - # above remains the fallback when `filename` is absent/invalid. - if filename and filename.endswith((".pt", ".pth")) and os.path.exists(os.path.join(model_path, filename)): - pt_file = filename - if pt_file is None: - raise RuntimeError("No .pt or .pth model file found in the provided path") +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 + class_names = _checkpoint_args_as_dict(raw_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." + ) + + 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 _is_ptl_checkpoint(checkpoint: Any) -> bool: + """True if the checkpoint is a raw PyTorch-Lightning rf-detr checkpoint dict.""" + return isinstance(checkpoint, dict) and "pytorch-lightning_version" in checkpoint + + +def _require_rfdetr() -> Any: + """Import ``rfdetr`` and verify it ships the upload-bundle helpers. - import torch + Raises :class:`MissingDependencyError` (a ModelPackagingError, so callers see + an actionable 400 rather than an opaque server error) when ``rfdetr`` is + missing or too old to export a Roboflow upload bundle. + """ + try: + import rfdetr + except ImportError as exc: + raise MissingDependencyError( + "The 'rfdetr' package is required to package raw PyTorch-Lightning rf-detr " + f"checkpoints. Install it with `pip install 'rfdetr>={RFDETR_MIN_VERSION}'`." + ) from exc + + if not hasattr(rfdetr.RFDETR, "export_for_roboflow"): + raise MissingDependencyError( + "The installed 'rfdetr' is too old to package raw PyTorch-Lightning rf-detr " + f"checkpoints. Upgrade it with `pip install --upgrade 'rfdetr>={RFDETR_MIN_VERSION}'`." + ) + return rfdetr - checkpoint = torch.load(os.path.join(model_path, pt_file), map_location="cpu", weights_only=False) +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] = [] + + checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings) + checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu") + + # Task detection + mismatch runs for every checkpoint shape (it also rejects + # keypoint rf-detr, which is not a supported upload type). 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." - ) + 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." + ) if _is_ptl_checkpoint(checkpoint): # Raw PyTorch-Lightning checkpoint: let rf-detr rebuild a proper upload - # bundle (weights.pt with `args.resolution` + class_names.txt). + # bundle (weights.pt with args.resolution + class_names.txt) into build_dir, + # so the caller's model_path is never mutated. rfdetr = _require_rfdetr() - pth = os.path.join(model_path, pt_file) try: - model = rfdetr.RFDETR.from_checkpoint(pth) + model = rfdetr.RFDETR.from_checkpoint(str(checkpoint_path)) except ValueError: - # Checkpoint lacks model_name/pretrain_weights signals; fall back to - # the already-validated user-provided model_type to pick the subclass. + # Checkpoint lacks model_name/pretrain_weights signals; fall back to the + # already-validated model_type to pick the RFDETR subclass. model_cls = getattr(rfdetr, _RFDETR_MODEL_TYPE_TO_CLASS[model_type]) - model = model_cls(pretrain_weights=pth) - model.export_for_roboflow(model_path) # writes weights.pt + class_names.txt + model = model_cls(pretrain_weights=str(checkpoint_path)) + model.export_for_roboflow(str(build_dir)) # writes weights.pt + class_names.txt else: - get_classnames_txt_for_rfdetr(model_path, pt_file, checkpoint=checkpoint) + # Roboflow's server-side RF-DETR conversion reads checkpoint["args"] (the + # class names, class count, and model config). A bare inference state_dict — + # e.g. {"model": } with nothing else — would otherwise package and + # upload fine, then fail conversion with an opaque KeyError: 'args'. Catch it + # here so the caller gets an actionable error before uploading. + if not isinstance(checkpoint, dict) or checkpoint.get("args") is None: + raise ModelPackagingError( + f"The RF-DETR checkpoint '{checkpoint_path.name}' is missing its 'args' " + "metadata; it looks like a bare inference state_dict. Roboflow's weight " + "conversion needs the full training checkpoint (args carries the class " + "names, class count, and model config). Re-export the checkpoint from your " + "training run, or download the deploy checkpoint from Roboflow." + ) - # 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")) + model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) + + weights_dest = build_dir / "weights.pt" + # In the legacy deploy flow build_dir is model_path, so a checkpoint already + # named weights.pt is its own destination; copying would raise SameFileError. + if checkpoint_path.resolve() != weights_dest.resolve(): + shutil.copy(checkpoint_path, weights_dest) + _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 - required_files = ["weights.pt"] - optional_files = ["results.csv", "results.png", "model_artifacts.json", "class_names.txt"] +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)}." + ) - 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) + 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]]: + if model_type != "yolonas": + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for YOLO-NAS upload. The only " + "supported YOLO-NAS type is 'yolonas'; the architecture size goes in opt.yaml " + "as 'architecture: yolo_nas_s' (or _m / _l)." + ) + 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") + # A SuperGradients YOLO-NAS checkpoint carries processing_params.class_names. + # A bare state_dict (e.g. torch.save(net.state_dict())) lacks it and would + # otherwise raise a raw KeyError/TypeError instead of an actionable error. + processing_params = checkpoint.get("processing_params") if isinstance(checkpoint, dict) else None + class_names = processing_params.get("class_names") if isinstance(processing_params, dict) else None + if not class_names: + raise ModelPackagingError( + f"The YOLO-NAS checkpoint '{weights_path.name}' is missing " + "processing_params.class_names; it looks like a bare state_dict. Provide the " + "full training checkpoint saved by SuperGradients." + ) + 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)}.") - 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) + 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, [] - 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) @@ -485,7 +1255,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" @@ -503,122 +1273,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/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index bee9a67f..c28cfdd0 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -1,30 +1,42 @@ +import json import os import sys +import tarfile import tempfile +import types import unittest import zipfile +from pathlib import Path from types import SimpleNamespace from unittest import mock -try: - # torch is an optional, lazily-imported SDK dependency; absent in CI. Tests that - # round-trip a real checkpoint through `_process_rfdetr` are skipped without it. - import torch - - _HAS_TORCH = True -except ImportError: - _HAS_TORCH = False - 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 ( _RFDETR_MODEL_TYPE_TO_CLASS, + MissingFileError, + ModelPackagingError, + SizeMismatchError, + TaskMismatchError, + UnsupportedModelError, + _checkpoint_args_as_dict, _detect_rfdetr_task, _detect_yolo_task, + _filtered_args, + _infer_yolo_size, _is_ptl_checkpoint, - _process_rfdetr, + _legacy_yolo_args, _require_rfdetr, + _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, ) @@ -113,6 +125,34 @@ def test_unrecognized_returns_none(self): self.assertIsNone(_detect_rfdetr_task({"model_name": None})) self.assertIsNone(_detect_rfdetr_task({"args": SimpleNamespace(other=1)})) + def test_scalar_args_do_not_raise_typeerror(self): + # A corrupt checkpoint storing args as a bare scalar must not escape the + # ModelPackagingError contract with a raw vars() TypeError. + self.assertIsNone(_detect_rfdetr_task({"args": 640})) + + +class CheckpointArgsAsDictTest(unittest.TestCase): + def test_coerces_dict_namespace_none_and_scalar(self): + self.assertEqual(_checkpoint_args_as_dict({"a": 1}), {"a": 1}) + self.assertEqual(_checkpoint_args_as_dict(SimpleNamespace(a=1)), {"a": 1}) + self.assertEqual(_checkpoint_args_as_dict(None), {}) + self.assertEqual(_checkpoint_args_as_dict(640), {}) + self.assertEqual(_checkpoint_args_as_dict(["not", "a", "dict"]), {}) + + +class FilteredArgsTest(unittest.TestCase): + def test_keeps_only_upload_keys_from_dict_or_namespace(self): + self.assertEqual( + _filtered_args({"model": "m", "imgsz": 640, "batch": 8, "lr0": 0.01}), + {"model": "m", "imgsz": 640, "batch": 8}, + ) + self.assertEqual(_filtered_args(SimpleNamespace(imgsz=320, batch=4, extra=1)), {"imgsz": 320, "batch": 4}) + + def test_scalar_or_none_args_do_not_raise(self): + # A corrupt .args must coerce to {} instead of raising a raw TypeError. + self.assertEqual(_filtered_args(None), {}) + self.assertEqual(_filtered_args(640), {}) + class GetClassnamesTxtForRfdetrTest(unittest.TestCase): def _classnames(self, args): @@ -131,6 +171,655 @@ 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, []) + + def test_bare_family_uninferable_raises_by_default(self): + with self.assertRaises(SizeMismatchError): + _resolve_yolo_size("yolov10", _FakeYoloModel({}), []) + + def test_bare_family_uninferable_proceeds_under_allow_mismatch(self): + # The interactive retry sets allow_mismatch=True after the user confirms; + # this branch must then converge (return) rather than raise forever. + warnings: list = [] + self.assertEqual( + _resolve_yolo_size("yolov10", _FakeYoloModel({}), warnings, allow_mismatch=True), + "yolov10", + ) + self.assertTrue(warnings and "bare family name" in warnings[0]) + + +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).resolve()) + 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_yolonas_requires_exact_model_type(self): + # Only 'yolonas' is valid; a suffixed typo like 'yolonas-foo' passes the + # family prefix gate but must be rejected before upload, not by the backend. + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("yolonas-foo", 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_rfdetr_bare_state_dict_without_args_fails_before_upload(self): + # A stripped inference checkpoint ({"model": state_dict} with no "args") + # would package fine but fail Roboflow's server-side conversion with an + # opaque KeyError: 'args'. Packaging must reject it up front. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"model": {"backbone.weight": object()}}) + with _import_patch({"torch": torch}): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("rfdetr-base", str(model_dir), filename="weights.pt") + self.assertIn("args", str(ctx.exception)) + self.assertIn("state_dict", str(ctx.exception)) + + def test_yolonas_bare_state_dict_without_class_names_fails_before_upload(self): + # A bare YOLO-NAS state_dict lacks processing_params.class_names; it must + # raise an actionable error instead of a raw KeyError/TypeError. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights").mkdir() + (model_dir / "weights" / "best.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"backbone.weight": object()}) # bare state_dict, no processing_params + with _import_patch({"torch": torch}): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", str(model_dir)) + self.assertIn("class_names", str(ctx.exception)) + self.assertIn("state_dict", str(ctx.exception)) + + def test_rfdetr_explicit_missing_filename_does_not_fall_back(self): + # A typo'd explicit filename must fail loudly, not silently upload a + # different checkpoint that happens to be in the directory. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "some_other_checkpoint.pth").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + with self.assertRaises(MissingFileError): + package_custom_weights("rfdetr-base", str(model_dir), filename="checkpoint_epoch50.pth") + + def test_rfdetr_legacy_deploy_layout_does_not_self_copy(self): + # Legacy deploy passes build_dir=model_path with a top-level weights.pt; + # copying weights.pt onto itself would raise shutil.SameFileError. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights.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), filename="weights.pt", build_dir=model_dir + ) + with zipfile.ZipFile(bundle.archive_path) as archive: + self.assertIn("weights.pt", archive.namelist()) + + 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() + + def test_rejects_absolute_filename(self): + # A hosted caller (the MCP server) forwards filename verbatim; an absolute + # path — POSIX or Windows-style, regardless of the packaging host's OS — + # must be rejected rather than reading weights from outside model_path. + with tempfile.TemporaryDirectory() as tmp: + for bad in ("/etc/passwd", "C:\\Windows\\System32\\config"): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", tmp, filename=bad) + self.assertIn("absolute", str(ctx.exception)) + + def test_rejects_filename_escaping_model_path(self): + # '..' segments must not let the caller escape the model directory. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) / "model" + model_dir.mkdir() + (Path(tmp) / "secret.pt").write_bytes(b"outside") + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", str(model_dir), filename="../secret.pt") + self.assertIn("outside model_path", str(ctx.exception)) + + def test_rejects_filename_pointing_at_a_directory(self): + # '' / '.' resolve to model_path itself and a subdirectory stays inside + # it; all three would otherwise reach torch.load() and leak a raw + # IsADirectoryError outside the ModelPackagingError contract. + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "weights").mkdir() + for bad in ("", ".", "weights"): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", tmp, filename=bad) + self.assertIn("not a directory", str(ctx.exception)) + + def _attempt_ultralytics_yolo(self, model_type, checkpoint): + fake_ultralytics = types.ModuleType("ultralytics") + fake_ultralytics.__version__ = "8.3.0" + fake_torch = _fake_torch(checkpoint) + 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}), + ): + # allow_dependency_mismatch keeps a version mismatch from masking the + # completeness error under test. + return package_custom_weights(model_type, str(model_dir), allow_dependency_mismatch=True) + + def test_yolov8_missing_nc_attr_raises_packaging_error(self): + # A stripped Ultralytics checkpoint that lacks model.nc must raise a + # ModelPackagingError, not a raw AttributeError (an opaque 500 in the MCP). + class DetectionModel: + names = {0: "cat"} + yaml = {"nc": 1, "depth_multiple": 0.33, "width_multiple": 0.25} + args = {"imgsz": 640} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("nc", str(ctx.exception)) + + def test_yolov8_missing_args_attr_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + nc = 1 + yaml = {"nc": 1, "depth_multiple": 0.33, "width_multiple": 0.25} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("args", str(ctx.exception)) + + def test_yolov8_missing_yaml_attr_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + nc = 1 + args = {"imgsz": 640} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("yaml", str(ctx.exception)) + + def test_yolov11_missing_train_args_raises_packaging_error(self): + # yolov10/11/12/26 read args from checkpoint["train_args"]; a checkpoint + # without it must fail with the ModelPackagingError contract. + class DetectionModel: + names = {0: "cat"} + yaml = {"nc": 1, "scale": "n"} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov11n", {"model": DetectionModel()}) + self.assertIn("train_args", str(ctx.exception)) + + def test_yolov11_missing_nc_in_model_yaml_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + yaml = {"scale": "n"} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov11n", {"model": DetectionModel(), "train_args": {"imgsz": 640}}) + self.assertIn("nc", str(ctx.exception)) + + def test_rfdetr_multiple_discovered_checkpoints_warns_which_used(self): + # With no explicit filename and several checkpoints present, discovery is + # ambiguous; the warning must name the candidates and which one was used. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "alpha.pt").write_bytes(b"checkpoint") + (model_dir / "beta.pth").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: + warning = "\n".join(bundle.warnings) + self.assertIn("multiple", warning) + self.assertIn("alpha.pt", warning) + self.assertIn("beta.pth", warning) + finally: + bundle.cleanup() + + +class ProcessHuggingfaceTest(unittest.TestCase): + """Packaging for HuggingFace-backed models (paligemma / florence-2).""" + + def test_missing_companion_files_raise_missing_file(self): + # A safetensors checkpoint without its tokenizer/preprocessor companions + # now raises MissingFileError instead of prompting for the files. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "model.safetensors").write_bytes(b"weights") + with mock.patch("builtins.input", side_effect=AssertionError("must not prompt")): + with self.assertRaises(MissingFileError) as ctx: + package_custom_weights("florence-2-base", str(model_dir)) + message = str(ctx.exception) + self.assertIn("tokenizer.json", message) + self.assertIn("preprocessor_config.json", message) + + def test_no_model_file_raises_missing_file(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "readme.txt").write_text("no weights here") + with self.assertRaises(MissingFileError) as ctx: + package_custom_weights("paligemma-3b-pt-224", str(model_dir)) + self.assertIn("safetensors", str(ctx.exception)) + + def test_npz_checkpoint_packages_into_tar(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "model.npz").write_bytes(b"npz-bytes") + bundle = package_custom_weights("paligemma-3b-pt-224", str(model_dir)) + try: + self.assertTrue(bundle.archive_path.name.endswith(".tar")) + with tarfile.open(bundle.archive_path) as tar: + self.assertIn("model.npz", tar.getnames()) + finally: + bundle.cleanup() + + def test_unsupported_huggingface_type_raises(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("florence-2-tiny", str(tmp)) + + +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") + + class RfdetrModelTypeToClassTest(unittest.TestCase): def test_representative_mappings(self): self.assertEqual(_RFDETR_MODEL_TYPE_TO_CLASS["rfdetr-seg-medium"], "RFDETRSegMedium") @@ -140,8 +829,7 @@ def test_keys_are_rfdetr_types_and_values_are_class_names(self): for model_type, class_name in _RFDETR_MODEL_TYPE_TO_CLASS.items(): self.assertTrue(model_type.startswith("rfdetr-"), model_type) self.assertTrue(class_name.startswith("RFDETR"), class_name) - # Segmentation types must map to Seg classes (and detection types must not). - for model_type, class_name in _RFDETR_MODEL_TYPE_TO_CLASS.items(): + # Segmentation types map to Seg classes (and detection types must not). self.assertEqual("seg" in model_type, "Seg" in class_name, model_type) @@ -164,15 +852,12 @@ def __init__(self, class_names=("cat", "dog")): self.class_names = list(class_names) def export_for_roboflow(self, output_dir): - torch.save({"dummy": True}, os.path.join(output_dir, "weights.pt")) - with open(os.path.join(output_dir, "class_names.txt"), "w") as f: - for name in self.class_names: - f.write(name + "\n") + (Path(output_dir) / "weights.pt").write_bytes(b"rebuilt-weights") + (Path(output_dir) / "class_names.txt").write_text("\n".join(self.class_names) + "\n") def _make_fake_rfdetr(*, from_checkpoint_raises=False, capabilities=True): - """Build a fake `rfdetr` module for injection via sys.modules.""" - + """Build a fake ``rfdetr`` module for injection via sys.modules.""" stub_model = _StubBundleModel() calls = {"from_checkpoint": 0, "fallback_constructed": 0, "constructor_kwargs": None} @@ -192,13 +877,11 @@ def __init__(self, *, pretrain_weights): module = SimpleNamespace() module.RFDETR = _RFDETR - # The SDK fallback resolves the subclass by name via _RFDETR_MODEL_TYPE_TO_CLASS, + # The fallback resolves the subclass by name via _RFDETR_MODEL_TYPE_TO_CLASS, # e.g. "rfdetr-seg-medium" -> getattr(rfdetr, "RFDETRSegMedium"). module.RFDETRSegMedium = _SizedModel - if capabilities: - _RFDETR.export_for_roboflow = _StubBundleModel.export_for_roboflow # capability marker on class - + _RFDETR.export_for_roboflow = _StubBundleModel.export_for_roboflow # capability marker module._calls = calls return module @@ -206,16 +889,15 @@ def __init__(self, *, pretrain_weights): class RequireRfdetrTest(unittest.TestCase): def test_raises_when_not_installed(self): with mock.patch.dict(sys.modules, {"rfdetr": None}): - with self.assertRaises(RuntimeError) as ctx: + with self.assertRaises(ModelPackagingError) as ctx: _require_rfdetr() self.assertIn("pip install", str(ctx.exception).lower()) self.assertIn("rfdetr", str(ctx.exception).lower()) def test_raises_when_capability_missing(self): - # rfdetr present but RFDETR lacks export_for_roboflow (too old) fake = SimpleNamespace(RFDETR=type("RFDETR", (), {})) with mock.patch.dict(sys.modules, {"rfdetr": fake}): - with self.assertRaises(RuntimeError) as ctx: + with self.assertRaises(ModelPackagingError) as ctx: _require_rfdetr() self.assertIn("upgrade", str(ctx.exception).lower()) @@ -225,107 +907,61 @@ def test_returns_module_when_capable(self): self.assertIs(_require_rfdetr(), fake) -@unittest.skipUnless(_HAS_TORCH, "requires torch") -class ProcessRfdetrPtlTest(unittest.TestCase): - def _write_ptl_checkpoint(self, model_path, *, segmentation_head=False, class_names=("cat", "dog")): - checkpoint = { - "pytorch-lightning_version": "2.1.0", - "args": {"segmentation_head": segmentation_head, "class_names": list(class_names)}, - } - torch.save(checkpoint, os.path.join(model_path, "checkpoint_best_ema.pth")) +class PackageRfdetrPtlTest(unittest.TestCase): + """PyTorch-Lightning rf-detr checkpoints are rebuilt via rfdetr into build_dir.""" + + def _package(self, model_type, fake_rfdetr, *, segmentation_head=False): + with tempfile.TemporaryDirectory() as model_dir: + (Path(model_dir) / "checkpoint_best_ema.pth").write_bytes(b"raw-ptl") + ckpt = { + "pytorch-lightning_version": "2.1.0", + "args": {"segmentation_head": segmentation_head, "class_names": ["cat", "dog"]}, + } + torch = _fake_torch(ckpt) + with _import_patch({"torch": torch}), mock.patch.dict(sys.modules, {"rfdetr": fake_rfdetr}): + bundle = package_custom_weights(model_type, model_dir, filename="checkpoint_best_ema.pth") + try: + with zipfile.ZipFile(bundle.archive_path) as archive: + names = archive.namelist() + finally: + bundle.cleanup() + return bundle, names def test_from_checkpoint_success_produces_bundle(self): fake = _make_fake_rfdetr() - with tempfile.TemporaryDirectory() as model_path: - self._write_ptl_checkpoint(model_path) - with mock.patch.dict(sys.modules, {"rfdetr": fake}): - zip_name, model_type = _process_rfdetr("rfdetr-base", model_path, "checkpoint_best_ema.pth") - self.assertEqual(model_type, "rfdetr-base") - self.assertEqual(fake._calls["from_checkpoint"], 1) - self.assertEqual(fake._calls["fallback_constructed"], 0) - with zipfile.ZipFile(os.path.join(model_path, zip_name)) as z: - self.assertIn("weights.pt", z.namelist()) + bundle, names = self._package("rfdetr-base", fake) + self.assertEqual(bundle.model_type, "rfdetr-base") + self.assertEqual(fake._calls["from_checkpoint"], 1) + self.assertEqual(fake._calls["fallback_constructed"], 0) + self.assertIn("weights.pt", names) + self.assertIn("class_names.txt", names) def test_from_checkpoint_valueerror_falls_back_to_model_type(self): fake = _make_fake_rfdetr(from_checkpoint_raises=True) - with tempfile.TemporaryDirectory() as model_path: - self._write_ptl_checkpoint(model_path, segmentation_head=True) - pth = os.path.join(model_path, "checkpoint_best_ema.pth") - with mock.patch.dict(sys.modules, {"rfdetr": fake}): - zip_name, model_type = _process_rfdetr("rfdetr-seg-medium", model_path, "checkpoint_best_ema.pth") - self.assertEqual(model_type, "rfdetr-seg-medium") - self.assertEqual(fake._calls["from_checkpoint"], 1) - self.assertEqual(fake._calls["fallback_constructed"], 1) - self.assertEqual(fake._calls["constructor_kwargs"], {"pretrain_weights": pth}) - with zipfile.ZipFile(os.path.join(model_path, zip_name)) as z: - self.assertIn("weights.pt", z.namelist()) + bundle, names = self._package("rfdetr-seg-medium", fake, segmentation_head=True) + self.assertEqual(bundle.model_type, "rfdetr-seg-medium") + self.assertEqual(fake._calls["from_checkpoint"], 1) + self.assertEqual(fake._calls["fallback_constructed"], 1) + self.assertIn("weights.pt", names) def test_ptl_path_raises_when_rfdetr_absent(self): - with tempfile.TemporaryDirectory() as model_path: - self._write_ptl_checkpoint(model_path) - with mock.patch.dict(sys.modules, {"rfdetr": None}): - with self.assertRaises(RuntimeError): - _process_rfdetr("rfdetr-base", model_path, "checkpoint_best_ema.pth") + with tempfile.TemporaryDirectory() as model_dir: + (Path(model_dir) / "checkpoint_best_ema.pth").write_bytes(b"raw-ptl") + ckpt = {"pytorch-lightning_version": "2.1.0", "args": {"segmentation_head": False}} + torch = _fake_torch(ckpt) + with _import_patch({"torch": torch}), mock.patch.dict(sys.modules, {"rfdetr": None}): + with self.assertRaises(ModelPackagingError): + package_custom_weights("rfdetr-base", model_dir, filename="checkpoint_best_ema.pth") def test_keypoint_checkpoint_rejected_before_export(self): - # A keypoint checkpoint must be rejected by the task check (before any rfdetr - # import/export), not uploaded as an object-detection rfdetr-base model. - with tempfile.TemporaryDirectory() as model_path: - checkpoint = { - "pytorch-lightning_version": "2.1.0", - "model_name": "RFDETRKeypointPreview", - "args": {"segmentation_head": False}, - } - torch.save(checkpoint, os.path.join(model_path, "checkpoint_best_ema.pth")) - with mock.patch.dict(sys.modules, {"rfdetr": None}): - with self.assertRaises(ValueError) as ctx: - _process_rfdetr("rfdetr-base", model_path, "checkpoint_best_ema.pth") - self.assertIn("pose", str(ctx.exception)) - - -@unittest.skipUnless(_HAS_TORCH, "requires torch") -class ProcessRfdetrLegacyTest(unittest.TestCase): - def _write_legacy_checkpoint(self, model_path): - checkpoint = {"args": {"segmentation_head": False, "class_names": ["cat", "dog"]}} - torch.save(checkpoint, os.path.join(model_path, "weights.pt")) - - def test_legacy_path_produces_bundle_without_importing_rfdetr(self): - with tempfile.TemporaryDirectory() as model_path: - self._write_legacy_checkpoint(model_path) - # Make any attempt to import rfdetr fail loudly. - with mock.patch.dict(sys.modules, {"rfdetr": None}): - zip_name, model_type = _process_rfdetr("rfdetr-base", model_path, "weights.pt") - self.assertEqual(model_type, "rfdetr-base") - with zipfile.ZipFile(os.path.join(model_path, zip_name)) as z: - names = z.namelist() - self.assertIn("weights.pt", names) - self.assertIn("class_names.txt", names) - - -@unittest.skipUnless(_HAS_TORCH, "requires torch") -class ProcessRfdetrFilenameSelectionTest(unittest.TestCase): - """The caller-provided `filename` must win over first-match directory discovery.""" - - @staticmethod - def _write_det_ckpt(path, class_names): - torch.save({"args": {"segmentation_head": False, "class_names": class_names}}, path) - - def test_honors_filename_over_directory_order(self): - with tempfile.TemporaryDirectory() as model_path: - self._write_det_ckpt(os.path.join(model_path, "distractor.pth"), ["distractor_class"]) - self._write_det_ckpt(os.path.join(model_path, "checkpoint_best_ema.pth"), ["cat", "dog"]) - # Force directory order so the distractor would be picked under first-match discovery; - # the fix must select `filename` regardless. - listing = ["distractor.pth", "checkpoint_best_ema.pth"] - with ( - mock.patch.dict(sys.modules, {"rfdetr": None}), - mock.patch("roboflow.util.model_processor.os.listdir", return_value=listing), - ): - zip_name, _ = _process_rfdetr("rfdetr-base", model_path, "checkpoint_best_ema.pth") - with zipfile.ZipFile(os.path.join(model_path, zip_name)) as z: - classes = z.read("class_names.txt").decode().split() - self.assertIn("cat", classes) - self.assertNotIn("distractor_class", classes) + # A keypoint checkpoint is rejected by the task check before any rfdetr use. + with tempfile.TemporaryDirectory() as model_dir: + (Path(model_dir) / "checkpoint_best_ema.pth").write_bytes(b"raw-ptl") + ckpt = {"pytorch-lightning_version": "2.1.0", "model_name": "RFDETRKeypointPreview"} + torch = _fake_torch(ckpt) + with _import_patch({"torch": torch}): + with self.assertRaises(TaskMismatchError): + package_custom_weights("rfdetr-base", model_dir, filename="checkpoint_best_ema.pth") if __name__ == "__main__":