ovkit 0.1.0__py3-none-any.whl

ovkit/__init__.py

@@ -0,0 +1,53 @@
+"""ovkit — a simple Python inference API for OpenVINO.
+
+One import, one :class:`Model` class, a callable object, and clean
+:class:`Results` — plus OpenVINO's strengths (AUTO/NPU devices, async, INT8).
+
+Example
+-------
+>>> from ovkit import Model
+>>> model = Model("rtdetr_r50")          # name -> auto download/convert/cache
+>>> results = model("img.jpg", conf=0.25)
+>>> for r in results:
+...     print(r.boxes.xyxy, r.boxes.conf, r.boxes.cls)
+...     r.save("out.jpg")
+"""
+
+from __future__ import annotations
+
+from .core.errors import (
+    ConversionError,
+    DownloadError,
+    GatedModelError,
+    LicenseError,
+    MirrorMissingError,
+    ModelNotFoundError,
+    OfflineError,
+    OVKitError,
+    TaskDetectionError,
+)
+from .core.model import Model
+from .core.registry import list_models
+from .core.results import Boxes, Keypoints, Masks, Probs, Results
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Model",
+    "Results",
+    "Boxes",
+    "Masks",
+    "Keypoints",
+    "Probs",
+    "list_models",
+    "OVKitError",
+    "ModelNotFoundError",
+    "OfflineError",
+    "DownloadError",
+    "GatedModelError",
+    "MirrorMissingError",
+    "ConversionError",
+    "TaskDetectionError",
+    "LicenseError",
+    "__version__",
+]

ovkit/__main__.py

@@ -0,0 +1,166 @@
+"""``ovkit`` command-line interface: ``run``, ``list``, ``info``, ``download``, ``devices``."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import __version__
+from .core.convert import to_ir
+from .core.download import fetch
+from .core.errors import OVKitError
+from .core.registry import list_models, resolve
+
+
+def _cmd_list(_: argparse.Namespace) -> int:
+    names = list_models()
+    if not names:
+        print("No models registered.")
+        return 0
+    aliases: list[tuple[str, str]] = []
+    models: list[tuple[str, str, str]] = []
+    for name in names:
+        entry = resolve(name)
+        if entry is None:
+            continue
+        if entry.name != name:  # capability alias -> its target
+            aliases.append((name, entry.name))
+            continue
+        desc = entry.description or ""
+        if len(desc) > 60:
+            desc = desc[:57] + "..."
+        models.append((name, str(entry.task), desc))
+    if aliases:
+        print("aliases (capability -> model):")
+        for alias, target in aliases:
+            print(f"  {alias:24s} -> {target}")
+        print()
+    print(f"models ({len(models)}):")
+    for name, task, desc in models:
+        print(f"  {name:44s} {task:18s} {desc}")
+    return 0
+
+
+def _cmd_info(args: argparse.Namespace) -> int:
+    entry = resolve(args.name)
+    if entry is None:
+        print(f"'{args.name}' is not a registered model.", file=sys.stderr)
+        return 1
+    print(f"name       : {entry.name}")
+    print(f"task       : {entry.task}")
+    if entry.description:
+        print(f"description: {entry.description}")
+    print(f"license    : {entry.license}")
+    print(f"source     : {entry.src} ({entry.repo or entry.url})")
+    print(f"precision  : {entry.precision}")
+    if entry.filename:
+        print(f"filename   : {entry.filename}")
+    if entry.imgsz:
+        print(f"imgsz      : {entry.imgsz}")
+    return 0
+
+
+def _cmd_download(args: argparse.Namespace) -> int:
+    entry = resolve(args.name)
+    if entry is None:
+        print(f"'{args.name}' is not a registered model.", file=sys.stderr)
+        return 1
+    print(f"Fetching {entry.name} from {entry.src}...")
+    source = fetch(entry)
+    print(f"Downloaded source: {source}")
+    if not args.no_convert:
+        ir = to_ir(source, entry.name, entry.precision)
+        print(f"IR ready: {ir}")
+    return 0
+
+
+def _cmd_devices(_: argparse.Namespace) -> int:
+    from .core.backend import available_devices
+
+    for dev in available_devices():
+        print(dev)
+    return 0
+
+
+def _cmd_run(args: argparse.Namespace) -> int:
+    """One-shot inference from the shell: ``ovkit run detect img.jpg``."""
+    from pathlib import Path
+
+    from .core.model import Model
+
+    model = Model(args.model, device=args.device)
+    results = model(args.source, conf=args.conf)
+    if not isinstance(results, list):  # raw (.npy/.wav) input -> tensor dict
+        for name, arr in results.items():
+            print(f"{name}: shape={tuple(arr.shape)} dtype={arr.dtype}")
+        return 0
+
+    for r in results:
+        parts = [f"task={r.task}"]
+        if r.text:
+            parts.append(f'text="{r.text}"')
+        if r.boxes is not None:
+            parts.append(f"{len(r.boxes)} boxes")
+            for x1, y1, x2, y2, c, cl in r.boxes.data[:20]:
+                print(
+                    f"  {r.name_for(int(cl)):16s} {c:.2f} [{int(x1)},{int(y1)},{int(x2)},{int(y2)}]"
+                )
+        if r.probs is not None:
+            top = ", ".join(
+                f"{r.name_for(int(i))} {r.probs.data[int(i)]:.2f}" for i in r.probs.top5
+            )
+            parts.append(f"top-5: {top}")
+        if r.masks is not None:
+            parts.append(f"masks {tuple(r.masks.data.shape)}")
+        if r.keypoints is not None:
+            parts.append(f"keypoints {tuple(r.keypoints.data.shape)}")
+        print(" | ".join(parts))
+
+    save = args.save
+    if save is None and results and Path(str(args.source)).is_file():
+        save = f"{Path(str(args.source)).stem}_out.jpg"
+    if save and results:
+        results[0].save(save)
+        print(f"saved -> {save}")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point. Returns a process exit code."""
+    parser = argparse.ArgumentParser(prog="ovkit", description="ovkit model utilities")
+    parser.add_argument("--version", action="version", version=f"ovkit {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_list = sub.add_parser("list", help="list registered models")
+    p_list.set_defaults(func=_cmd_list)
+
+    p_info = sub.add_parser("info", help="show details for a model")
+    p_info.add_argument("name")
+    p_info.set_defaults(func=_cmd_info)
+
+    p_dl = sub.add_parser("download", help="download (and convert) a model")
+    p_dl.add_argument("name")
+    p_dl.add_argument("--no-convert", action="store_true", help="skip IR conversion")
+    p_dl.set_defaults(func=_cmd_download)
+
+    p_dev = sub.add_parser("devices", help="list OpenVINO devices")
+    p_dev.set_defaults(func=_cmd_devices)
+
+    p_run = sub.add_parser("run", help="run a model on an image/folder/video from the shell")
+    p_run.add_argument("model", help="alias, registered name, or model path")
+    p_run.add_argument("source", help="image / folder / video path (or .npy/.wav)")
+    p_run.add_argument("--conf", type=float, default=0.25, help="confidence threshold")
+    p_run.add_argument("--device", default="AUTO", help="AUTO | CPU | GPU | NPU")
+    p_run.add_argument("--save", metavar="PATH", help="annotated output (default: <src>_out.jpg)")
+    p_run.set_defaults(func=_cmd_run)
+
+    args = parser.parse_args(argv)
+    try:
+        return args.func(args)
+    except OVKitError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

ovkit/core/__init__.py

@@ -0,0 +1,8 @@
+"""Core runtime: backend, model, results, registry, download, convert, tasks."""
+
+from __future__ import annotations
+
+from .model import Model
+from .results import Boxes, Keypoints, Masks, Probs, Results
+
+__all__ = ["Model", "Results", "Boxes", "Masks", "Keypoints", "Probs"]

ovkit/core/backend.py

@@ -0,0 +1,166 @@
+"""Thin OpenVINO runtime wrapper: device abstraction, sync + async inference.
+
+A :class:`Backend` owns a compiled model for a chosen device and exposes both a
+single-shot :meth:`infer` (synchronous) and a throughput-oriented
+:meth:`infer_batch` built on ``ov.AsyncInferQueue`` for streams/folders/video.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+#: A single OpenVINO Core, shared process-wide (creating many is wasteful).
+_CORE = None
+
+
+def core() -> Any:
+    """Return the shared :class:`openvino.Core`, creating it on first use."""
+    global _CORE
+    if _CORE is None:
+        import openvino as ov
+
+        _CORE = ov.Core()
+    return _CORE
+
+
+def available_devices() -> list[str]:
+    """Return device names visible to OpenVINO (e.g. ``["CPU", "GPU", "NPU"]``)."""
+    return list(core().available_devices)
+
+
+class Backend:
+    """A compiled model bound to a device, with sync and async inference.
+
+    Parameters
+    ----------
+    model:
+        Path to an IR ``.xml`` / ONNX file, or an already-built ``ov.Model``.
+    device:
+        OpenVINO device string. ``"AUTO"`` (default) lets OpenVINO pick.
+    """
+
+    def __init__(self, model: str | Path | Any, device: str = "AUTO") -> None:
+        self.device = device
+        c = core()
+        src = str(model) if isinstance(model, (str, Path)) else model
+        self.compiled = c.compile_model(src, device)
+        self.inputs = self.compiled.inputs
+        self.outputs = self.compiled.outputs
+
+    # -- introspection ------------------------------------------------------
+
+    @property
+    def input_shape(self) -> tuple[int, ...]:
+        """Partial shape of the first input as a tuple (``-1`` for dynamic)."""
+        ps = self.compiled.inputs[0].get_partial_shape()
+        dims: list[int] = []
+        for d in ps:
+            dims.append(int(d.get_length()) if d.is_static else -1)
+        return tuple(dims)
+
+    def _adapt_image_channels(self, arr: np.ndarray) -> np.ndarray:
+        """Match a single NCHW image tensor to the model's channel count.
+
+        Many OMZ models (OCR, some classifiers) take 1-channel grayscale; the
+        preprocessor always produces 3 channels, so reconcile the two here —
+        every image adapter funnels through :meth:`infer`. Only the 3<->1 image
+        case is touched; anything else passes through unchanged.
+        """
+        shape = self.input_shape
+        if arr.ndim != 4 or len(shape) != 4 or shape[1] not in (1, 3):
+            return arr
+        exp, got = shape[1], arr.shape[1]
+        if got == exp:
+            return arr
+        if exp == 1 and got == 3:  # to grayscale (channel-order agnostic)
+            return arr.mean(axis=1, keepdims=True).astype(arr.dtype)
+        if exp == 3 and got == 1:  # broadcast gray -> 3 channels
+            return np.repeat(arr, 3, axis=1)
+        return arr
+
+    def output_signatures(self) -> list[tuple[str, tuple[int, ...]]]:
+        """Return ``(name, shape)`` for each output (``-1`` for dynamic dims)."""
+        sigs: list[tuple[str, tuple[int, ...]]] = []
+        for out in self.compiled.outputs:
+            ps = out.get_partial_shape()
+            shape = tuple(int(d.get_length()) if d.is_static else -1 for d in ps)
+            try:
+                name = out.get_any_name()
+            except RuntimeError:
+                name = ""
+            sigs.append((name, shape))
+        return sigs
+
+    def rt_info(self, *keys: str) -> str | None:
+        """Read a runtime-info value from the underlying model, or ``None``."""
+        try:
+            model = self.compiled.get_runtime_model()
+            info = model.get_rt_info(list(keys))
+            return str(info)
+        except Exception:
+            return None
+
+    # -- inference ----------------------------------------------------------
+
+    def infer(self, inputs: np.ndarray | dict[Any, np.ndarray]) -> dict[str, np.ndarray]:
+        """Run one synchronous inference and return ``{output_name: ndarray}``."""
+        if isinstance(inputs, np.ndarray):
+            inputs = self._adapt_image_channels(inputs)
+        result = self.compiled(inputs)
+        return self._named(result)
+
+    def infer_batch(
+        self,
+        feeds: Iterable[np.ndarray | dict[Any, np.ndarray]],
+        callback: Callable[[int, dict[str, np.ndarray]], None] | None = None,
+        jobs: int = 0,
+    ) -> Iterator[dict[str, np.ndarray]]:
+        """Run inference over ``feeds`` using an async queue (throughput mode).
+
+        Yields result dicts in completion order. When ``callback`` is given it
+        is invoked as ``callback(index, result)``; otherwise results are
+        collected and yielded. ``jobs`` sets the number of in-flight requests
+        (``0`` lets OpenVINO choose the optimal number).
+        """
+        import openvino as ov
+
+        queue = ov.AsyncInferQueue(self.compiled, jobs)
+        collected: dict[int, dict[str, np.ndarray]] = {}
+
+        def _on_done(request: Any, userdata: int) -> None:
+            named = self._named({out: request.get_tensor(out).data for out in self.outputs})
+            if callback is not None:
+                callback(userdata, named)
+            else:
+                collected[userdata] = named
+
+        queue.set_callback(_on_done)
+        count = 0
+        for i, feed in enumerate(feeds):
+            if isinstance(feed, np.ndarray):
+                feed = self._adapt_image_channels(feed)
+            queue.start_async(feed, userdata=i)
+            count += 1
+        queue.wait_all()
+
+        if callback is None:
+            for i in range(count):
+                if i in collected:
+                    yield collected[i]
+
+    def _named(self, result: Any) -> dict[str, np.ndarray]:
+        named: dict[str, np.ndarray] = {}
+        for idx, out in enumerate(self.compiled.outputs):
+            try:
+                name = out.get_any_name()
+            except RuntimeError:
+                name = f"output_{idx}"
+            try:
+                named[name] = np.asarray(result[out])
+            except (KeyError, TypeError):
+                named[name] = np.asarray(result[idx])
+        return named

ovkit/core/constants.py

@@ -0,0 +1,158 @@
+"""Shared constants: cache locations, license policy, well-known class lists."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+# --- cache / environment ---------------------------------------------------
+
+#: Environment variable pointing at the ovkit cache root.
+ENV_HOME = "OVKIT_HOME"
+#: Environment variable; when ``"1"`` no network access is attempted.
+ENV_OFFLINE = "OVKIT_OFFLINE"
+
+
+def cache_root() -> Path:
+    """Return the ovkit cache root, honoring ``$OVKIT_HOME``.
+
+    Defaults to ``~/.cache/ovkit``. The directory is created on demand.
+    """
+    root = os.environ.get(ENV_HOME)
+    base = Path(root).expanduser() if root else Path.home() / ".cache" / "ovkit"
+    return base
+
+
+def is_offline() -> bool:
+    """Return ``True`` when offline mode is requested via ``$OVKIT_OFFLINE``."""
+    return os.environ.get(ENV_OFFLINE, "").strip() in {"1", "true", "True", "yes"}
+
+
+# --- license policy --------------------------------------------------------
+
+#: SPDX ids accepted for models registered in the manifest. Anything outside
+#: this set (AGPL, non-commercial weights, ...) must not ship with ovkit.
+PERMISSIVE_LICENSES: frozenset[str] = frozenset(
+    {
+        "apache-2.0",
+        "mit",
+        "bsd-2-clause",
+        "bsd-3-clause",
+        "bsd",
+        "isc",
+        "unlicense",
+        "cc0-1.0",
+        "mpl-2.0",
+    }
+)
+
+
+def is_permissive(license_id: str | None) -> bool:
+    """Return ``True`` if ``license_id`` is a known permissive SPDX id."""
+    if not license_id:
+        return False
+    return license_id.strip().lower() in PERMISSIVE_LICENSES
+
+
+# --- class name tables -----------------------------------------------------
+
+#: 80 COCO class names (detection/segmentation), indexed by class id.
+COCO80: tuple[str, ...] = (
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+)
+
+#: Registry of named class tables referenced from the manifest ``classes`` key.
+CLASS_TABLES: dict[str, tuple[str, ...]] = {
+    "coco80": COCO80,
+}
+
+
+def class_names(key: str | None, num_classes: int | None = None) -> dict[int, str]:
+    """Resolve a manifest ``classes`` key to an ``{id: name}`` mapping.
+
+    Falls back to ``class_<i>`` names when ``key`` is unknown. ``num_classes``,
+    when given, sizes the fallback table.
+    """
+    if key and key in CLASS_TABLES:
+        return {i: n for i, n in enumerate(CLASS_TABLES[key])}
+    n = num_classes if num_classes is not None else 0
+    return {i: f"class_{i}" for i in range(n)}

ovkit/core/convert.py

@@ -0,0 +1,64 @@
+"""Convert source models (ONNX / IR) to OpenVINO IR, with a conversion cache.
+
+Conversion runs at most once per ``(name, precision)``: the resulting IR is
+written to the model cache and reused on subsequent loads.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .download import model_cache_dir
+from .errors import ConversionError
+
+
+def _ir_paths(name: str, precision: str) -> tuple[Path, Path]:
+    base = model_cache_dir(name) / "ir" / precision
+    return base / "model.xml", base / "model.bin"
+
+
+def cached_ir(name: str, precision: str) -> Path | None:
+    """Return the cached IR ``.xml`` for ``(name, precision)`` if it exists."""
+    xml, _ = _ir_paths(name, precision)
+    return xml if xml.is_file() else None
+
+
+def to_ir(source: Path, name: str, precision: str = "fp16") -> Path:
+    """Convert ``source`` to OpenVINO IR and return the cached ``.xml`` path.
+
+    ``source`` may already be IR (``.xml``) — in that case it is passed through
+    unchanged. ONNX sources are converted with ``openvino.convert_model`` and
+    serialized (compressing weights to fp16 when ``precision == "fp16"``).
+    The result is cached so conversion happens only once.
+    """
+    source = Path(source)
+
+    # Already IR: load directly, no conversion/caching needed.
+    if source.suffix == ".xml":
+        return source
+
+    cached = cached_ir(name, precision)
+    if cached is not None:
+        return cached
+
+    try:
+        import openvino as ov
+    except ImportError as exc:  # pragma: no cover - dependency missing
+        raise ConversionError("openvino is required to convert models to IR.") from exc
+
+    if source.suffix not in {".onnx", ".pb", ".pdmodel"} and not source.is_file():
+        raise ConversionError(f"Cannot convert '{source}': unsupported or missing source.")
+
+    xml_path, _ = _ir_paths(name, precision)
+    xml_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        ov_model = ov.convert_model(str(source))
+        compress = precision == "fp16"
+        ov.save_model(ov_model, str(xml_path), compress_to_fp16=compress)
+    except Exception as exc:  # pragma: no cover - conversion variety
+        raise ConversionError(f"Failed to convert '{source.name}' to IR: {exc}") from exc
+
+    if not xml_path.is_file():  # pragma: no cover - defensive
+        raise ConversionError(f"Conversion produced no IR at {xml_path}.")
+    return xml_path