Sync from arm/Optimized-Runtimes via hub-sync

9422383 verified 6 days ago

6.57 kB

"""Load and validate one converted YOLO26n model directory."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import cast

from .errors import Yolo26nError

# These values define the converted-model/runtime contract. They describe the
# only model family and native target this purpose-built runtime accepts.
_MODEL_KIND = "yolo26n-exported-model"
_MODEL_VERSION = 1
_TARGET = "linux-aarch64-raspberry-pi-5"

# The Raspberry Pi 5 runtime owns a fixed four-worker pool, matching its four
# CPU cores, so the thread count is part of model/runtime compatibility.
_TARGET_THREAD_COUNT = 4
_BOX_COORDINATE_CHANNELS = 4

# Graph sizing and memory planning are implemented only for these profiles.
_SUPPORTED_RESOLUTIONS = (320, 480, 640, 800, 960, 1280)

# YOLO26n predicts at its P3, P4, and P5 feature levels.
_DETECTION_STRIDES = (8, 16, 32)


@dataclass(frozen=True, slots=True)
class _LoadedModel:
    """Validated metadata and serialized weights required by a session."""

    input_shape: tuple[int, int, int, int]
    output_shape: tuple[int, int, int]
    class_names: tuple[str, ...]
    model_bytes: bytes


@dataclass(frozen=True, slots=True)
class _ArtifactRecord:
    relative_path: Path
    size: int
    sha256: str

    @classmethod
    def from_manifest(cls, value: object) -> _ArtifactRecord:
        """Parse the model artifact without accepting unknown fields."""
        record = _require_mapping(value, "artifact")
        if set(record) != {"bytes", "path", "sha256"}:
            raise Yolo26nError("invalid artifact in YOLO26n model manifest")

        path_value = record["path"]
        size = record["bytes"]
        digest = record["sha256"]
        if (
            not isinstance(path_value, str)
            or type(size) is not int
            or size < 0
            or not isinstance(digest, str)
            or len(digest) != 64
            or any(character not in "0123456789abcdef" for character in digest)
        ):
            raise Yolo26nError("invalid artifact in YOLO26n model manifest")

        relative_path = Path(path_value)
        if relative_path.is_absolute() or ".." in relative_path.parts:
            raise Yolo26nError("invalid artifact path in YOLO26n model manifest")
        return cls(relative_path=relative_path, size=size, sha256=digest)

    def load(self, model_directory: Path) -> bytes:
        """Read the model bytes and verify their manifest record."""
        path = model_directory / self.relative_path
        try:
            model_bytes = path.read_bytes()
        except OSError as error:
            raise Yolo26nError(
                f"YOLO26n model artifact {self.relative_path.as_posix()!r} cannot be read"
            ) from error
        if len(model_bytes) != self.size or hashlib.sha256(model_bytes).hexdigest() != self.sha256:
            raise Yolo26nError(
                f"YOLO26n model artifact {self.relative_path.as_posix()!r} "
                "does not match its manifest"
            )
        return model_bytes


def _load_model(model_directory: str | Path) -> _LoadedModel:
    """Validate one converted model before native code receives its weights."""
    directory = Path(model_directory).expanduser().resolve()
    if not directory.is_dir():
        raise FileNotFoundError(f"YOLO26n model directory does not exist: {directory}")

    manifest = _read_manifest(directory / "manifest.json")
    if (manifest.get("kind"), manifest.get("version")) != (_MODEL_KIND, _MODEL_VERSION):
        raise Yolo26nError("unsupported YOLO26n model manifest kind or version")

    model = _require_mapping(manifest.get("model"), "model record")
    input_shape = _require_shape(model.get("input_shape"), "input shape")
    output_shape = _require_shape(model.get("output_shape"), "output shape")
    resolution = input_shape[2] if len(input_shape) == 4 else None
    if (
        len(input_shape) != 4
        or input_shape[:2] != (1, 3)
        or input_shape[2] != input_shape[3]
        or resolution not in _SUPPORTED_RESOLUTIONS
    ):
        raise Yolo26nError("unsupported YOLO26n input shape")
    if (
        len(output_shape) != 3
        or output_shape[0] != 1
        or output_shape[1] <= _BOX_COORDINATE_CHANNELS
        or output_shape[2] != _detection_positions(resolution)
    ):
        raise Yolo26nError("unsupported YOLO26n output shape")

    expected_target = {
        "version": 1,
        "target": _TARGET,
        "input_shape": list(input_shape),
        "input_dtype": "fp32",
        "output_dtype": "fp32",
        "thread_count": _TARGET_THREAD_COUNT,
    }
    if model.get("target") != expected_target:
        raise Yolo26nError("unsupported YOLO26n target contract")

    names = model.get("class_names")
    if (
        not isinstance(names, list)
        or len(names) != output_shape[1] - _BOX_COORDINATE_CHANNELS
        or any(not isinstance(name, str) or not name for name in names)
    ):
        raise Yolo26nError("invalid class names in YOLO26n model manifest")

    artifact = _ArtifactRecord.from_manifest(manifest.get("artifact"))
    return _LoadedModel(
        input_shape=input_shape,
        output_shape=output_shape,
        class_names=tuple(cast(list[str], names)),
        model_bytes=artifact.load(directory),
    )


def _read_manifest(path: Path) -> Mapping[str, object]:
    try:
        value: object = json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        raise FileNotFoundError(f"YOLO26n model manifest does not exist: {path}") from None
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise Yolo26nError(f"invalid YOLO26n model manifest: {path}") from error
    return _require_mapping(value, "manifest")


def _require_mapping(value: object, name: str) -> Mapping[str, object]:
    if not isinstance(value, dict) or any(not isinstance(key, str) for key in value):
        raise Yolo26nError(f"invalid {name} in YOLO26n model manifest")
    return cast(Mapping[str, object], value)


def _require_shape(value: object, name: str) -> tuple[int, ...]:
    if (
        not isinstance(value, list)
        or not value
        or any(type(dimension) is not int or dimension <= 0 for dimension in value)
    ):
        raise Yolo26nError(f"invalid {name} in YOLO26n model manifest")
    return tuple(value)


def _detection_positions(resolution: int) -> int:
    return sum((resolution // stride) ** 2 for stride in _DETECTION_STRIDES)

	"""Load and validate one converted YOLO26n model directory."""

	from __future__ import annotations

	import hashlib
	import json
	from collections.abc import Mapping
	from dataclasses import dataclass
	from pathlib import Path
	from typing import cast

	from .errors import Yolo26nError

	# These values define the converted-model/runtime contract. They describe the
	# only model family and native target this purpose-built runtime accepts.
	_MODEL_KIND = "yolo26n-exported-model"
	_MODEL_VERSION = 1
	_TARGET = "linux-aarch64-raspberry-pi-5"

	# The Raspberry Pi 5 runtime owns a fixed four-worker pool, matching its four
	# CPU cores, so the thread count is part of model/runtime compatibility.
	_TARGET_THREAD_COUNT = 4
	_BOX_COORDINATE_CHANNELS = 4

	# Graph sizing and memory planning are implemented only for these profiles.
	_SUPPORTED_RESOLUTIONS = (320, 480, 640, 800, 960, 1280)

	# YOLO26n predicts at its P3, P4, and P5 feature levels.
	_DETECTION_STRIDES = (8, 16, 32)


	@dataclass(frozen=True, slots=True)
	class _LoadedModel:
	"""Validated metadata and serialized weights required by a session."""

	input_shape: tuple[int, int, int, int]
	output_shape: tuple[int, int, int]
	class_names: tuple[str, ...]
	model_bytes: bytes


	@dataclass(frozen=True, slots=True)
	class _ArtifactRecord:
	relative_path: Path
	size: int
	sha256: str

	@classmethod
	def from_manifest(cls, value: object) -> _ArtifactRecord:
	"""Parse the model artifact without accepting unknown fields."""
	record = _require_mapping(value, "artifact")
	if set(record) != {"bytes", "path", "sha256"}:
	raise Yolo26nError("invalid artifact in YOLO26n model manifest")

	path_value = record["path"]
	size = record["bytes"]
	digest = record["sha256"]
	if (
	not isinstance(path_value, str)
	or type(size) is not int
	or size < 0
	or not isinstance(digest, str)
	or len(digest) != 64
	or any(character not in "0123456789abcdef" for character in digest)
	):
	raise Yolo26nError("invalid artifact in YOLO26n model manifest")

	relative_path = Path(path_value)
	if relative_path.is_absolute() or ".." in relative_path.parts:
	raise Yolo26nError("invalid artifact path in YOLO26n model manifest")
	return cls(relative_path=relative_path, size=size, sha256=digest)

	def load(self, model_directory: Path) -> bytes:
	"""Read the model bytes and verify their manifest record."""
	path = model_directory / self.relative_path
	try:
	model_bytes = path.read_bytes()
	except OSError as error:
	raise Yolo26nError(
	f"YOLO26n model artifact {self.relative_path.as_posix()!r} cannot be read"
	) from error
	if len(model_bytes) != self.size or hashlib.sha256(model_bytes).hexdigest() != self.sha256:
	raise Yolo26nError(
	f"YOLO26n model artifact {self.relative_path.as_posix()!r} "
	"does not match its manifest"
	)
	return model_bytes


	def _load_model(model_directory: str \| Path) -> _LoadedModel:
	"""Validate one converted model before native code receives its weights."""
	directory = Path(model_directory).expanduser().resolve()
	if not directory.is_dir():
	raise FileNotFoundError(f"YOLO26n model directory does not exist: {directory}")

	manifest = _read_manifest(directory / "manifest.json")
	if (manifest.get("kind"), manifest.get("version")) != (_MODEL_KIND, _MODEL_VERSION):
	raise Yolo26nError("unsupported YOLO26n model manifest kind or version")

	model = _require_mapping(manifest.get("model"), "model record")
	input_shape = _require_shape(model.get("input_shape"), "input shape")
	output_shape = _require_shape(model.get("output_shape"), "output shape")
	resolution = input_shape[2] if len(input_shape) == 4 else None
	if (
	len(input_shape) != 4
	or input_shape[:2] != (1, 3)
	or input_shape[2] != input_shape[3]
	or resolution not in _SUPPORTED_RESOLUTIONS
	):
	raise Yolo26nError("unsupported YOLO26n input shape")
	if (
	len(output_shape) != 3
	or output_shape[0] != 1
	or output_shape[1] <= _BOX_COORDINATE_CHANNELS
	or output_shape[2] != _detection_positions(resolution)
	):
	raise Yolo26nError("unsupported YOLO26n output shape")

	expected_target = {
	"version": 1,
	"target": _TARGET,
	"input_shape": list(input_shape),
	"input_dtype": "fp32",
	"output_dtype": "fp32",
	"thread_count": _TARGET_THREAD_COUNT,
	}
	if model.get("target") != expected_target:
	raise Yolo26nError("unsupported YOLO26n target contract")

	names = model.get("class_names")
	if (
	not isinstance(names, list)
	or len(names) != output_shape[1] - _BOX_COORDINATE_CHANNELS
	or any(not isinstance(name, str) or not name for name in names)
	):
	raise Yolo26nError("invalid class names in YOLO26n model manifest")

	artifact = _ArtifactRecord.from_manifest(manifest.get("artifact"))
	return _LoadedModel(
	input_shape=input_shape,
	output_shape=output_shape,
	class_names=tuple(cast(list[str], names)),
	model_bytes=artifact.load(directory),
	)


	def _read_manifest(path: Path) -> Mapping[str, object]:
	try:
	value: object = json.loads(path.read_text(encoding="utf-8"))
	except FileNotFoundError:
	raise FileNotFoundError(f"YOLO26n model manifest does not exist: {path}") from None
	except (UnicodeDecodeError, json.JSONDecodeError) as error:
	raise Yolo26nError(f"invalid YOLO26n model manifest: {path}") from error
	return _require_mapping(value, "manifest")


	def _require_mapping(value: object, name: str) -> Mapping[str, object]:
	if not isinstance(value, dict) or any(not isinstance(key, str) for key in value):
	raise Yolo26nError(f"invalid {name} in YOLO26n model manifest")
	return cast(Mapping[str, object], value)


	def _require_shape(value: object, name: str) -> tuple[int, ...]:
	if (
	not isinstance(value, list)
	or not value
	or any(type(dimension) is not int or dimension <= 0 for dimension in value)
	):
	raise Yolo26nError(f"invalid {name} in YOLO26n model manifest")
	return tuple(value)


	def _detection_positions(resolution: int) -> int:
	return sum((resolution // stride) ** 2 for stride in _DETECTION_STRIDES)