# SPDX-License-Identifier: Apache-2.0 """``SuperPoint``: the Python entry point of the Tenstorrent Blackhole SuperPoint port. from tt_superpoint import SuperPoint with SuperPoint.from_pretrained(device_id=0) as model: out = model("image.jpg") out.keypoints # (N, 2) float32 [x, y] in original image pixels out.scores # (N,) float32, descending out.descriptors # (N, 256) float32, L2-normalised The call runs the same request path as the HTTP server (``code/models/server/app.py``): one metal trace (network, device NMS, keypoint list, bilinear descriptor sampling), the same fallbacks, the same output order. Keypoints, scores and descriptors are equal to the server's response for the same image and parameters (the server rounds the JSON floats and sends float16 descriptors). """ from __future__ import annotations import logging import threading import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, Union import numpy as np import torch from . import device as _device from . import warmup as _warm from .inputs import is_single_image, split_batch, to_plane LOG = logging.getLogger("tt_superpoint") DEFAULT_MODEL_ID = "magic-leap-community/superpoint" #: The weights revision the published numbers and the server (serve.env) use. DEFAULT_REVISION = "734450e9ffe229074f5998494ddc615475cdb20a" #: Network frame (H, W). Every image is resized to it (bilinear, aspect ratio not kept). INPUT_SIZE: Tuple[int, int] = (480, 640) DESCRIPTOR_DIM = 256 #: Source sizes (W, H) whose device-resize variant is captured in from_pretrained (as serve.env). DEFAULT_PRECOMPILE_SIZES: Tuple[Tuple[int, int], ...] = tuple(_warm.DEFAULT_WARMUP_VARIANTS["sizes"]) MAX_KEYPOINTS_CAP = INPUT_SIZE[0] * INPUT_SIZE[1] MAX_NMS_RADIUS = 32 @dataclass class SuperPointOutput: """Keypoints of one image. Same keys as the Hugging Face ``SuperPointImageProcessor.post_process_keypoint_detection`` result, so ``out["keypoints"]`` also works; keypoints stay float (sub-pixel scale) here instead of HF's int32. Attributes: keypoints: ``(N, 2)`` float32 tensor, ``[x, y]`` in pixels of the original image (network-frame pixel position x ``scale``). scores: ``(N,)`` float32 tensor, keypoint scores in (0, 1], sorted descending. descriptors: ``(N, 256)`` float32 tensor, L2-normalised, in keypoint order; ``None`` when called with ``return_descriptors=False``. image_size: ``(height, width)`` of the original image. scale: ``(sx, sy)`` = original size / network size (640, 480). device_nms: True when NMS ran on the device (False for the host NMS fallback, used for ``nms_radius`` 0 or > 8). """ keypoints: torch.Tensor scores: torch.Tensor descriptors: Optional[torch.Tensor] image_size: Tuple[int, int] scale: Tuple[float, float] device_nms: bool = True timing_ms: Dict[str, float] = field(default_factory=dict, repr=False) def __len__(self) -> int: return int(self.keypoints.shape[0]) def __getitem__(self, key: str): if key not in ("keypoints", "scores", "descriptors"): raise KeyError(key) return getattr(self, key) def keys(self): return ("keypoints", "scores", "descriptors") def numpy(self) -> Dict[str, Optional[np.ndarray]]: """``{"keypoints", "scores", "descriptors"}`` as numpy arrays (float32).""" return {k: (None if getattr(self, k) is None else getattr(self, k).numpy()) for k in self.keys()} def to_dict(self) -> Dict[str, Any]: """JSON-friendly dict (lists of floats), the same fields as the HTTP response.""" d = { "num_keypoints": len(self), "keypoints": self.keypoints.tolist(), "scores": self.scores.tolist(), "image_size": {"height": self.image_size[0], "width": self.image_size[1]}, "scale": {"x": self.scale[0], "y": self.scale[1]}, } if self.descriptors is not None: d["descriptors"] = self.descriptors.tolist() return d ImageLike = Any # path, bytes, PIL.Image, numpy array, torch tensor class SuperPoint: """SuperPoint keypoint detector + descriptor on one Tenstorrent Blackhole chip. Create it with :meth:`from_pretrained`; call it on an image (or a list of images); close it with :meth:`close` or a ``with`` block. Device work is serialised by an internal lock, so one instance can be shared between threads. """ def __init__(self, tt_model, tt_input, device, *, owns_device: bool, dispatch: Optional[str], border_removal_distance: int, config: Dict[str, Any]): self._m = tt_model self._tt_in = tt_input self._device = device self._owns_device = owns_device self._lock = threading.Lock() self._closed = False self.dispatch = dispatch self.border_removal_distance = int(border_removal_distance) #: Defaults of the per-call parameters (the server's defaults). self.defaults = dict(max_keypoints=1024, keypoint_threshold=0.005, nms_radius=4, return_descriptors=True) self.config = config self._pool: Optional[ThreadPoolExecutor] = None # host decode threads of list calls (kept) self._pool_lock = threading.Lock() self._warm_done: set = set() self._fallbacks_on = False # ------------------------------------------------------------------ construction @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str = DEFAULT_MODEL_ID, *, revision: Optional[str] = None, device=None, device_id: int = 0, dispatch: str = "auto", warmup_variants: Any = None, precompile_sizes: Optional[Iterable[Tuple[int, int]]] = None, precompile_nms_radii: Iterable[int] = (), trace_region_size: int = _device.TRACE_REGION_SIZE, local_files_only: bool = False, verbose: bool = False, ) -> "SuperPoint": """Load the weights, open the device, build the model and capture its traces. Args: pretrained_model_name_or_path: Hugging Face repo id or a local directory with ``config.json`` + ``model.safetensors`` (default ``magic-leap-community/superpoint``). revision: weights revision; default: the pinned revision of the published numbers (``734450e9``) for the default repo, the repo's default branch otherwise. device: an already-open ttnn device to use (it must have ``l1_small_size >= 32768`` and ``trace_region_size >= 32 MiB``; it is not closed by :meth:`close`). When ``None``, chip ``device_id`` is opened. device_id: chip to open when ``device`` is None. dispatch: ``"auto"`` (ETH dispatch + 12x10 grid when the patched tt-metal is present, else the default dispatch with a warning), ``"eth"`` or ``"worker"``. warmup_variants: the per-request variants to prepare now, so that the first real call of each one is as fast as the later calls (see :mod:`tt_superpoint.warmup`): ``None`` / ``"default"`` (sizes 1920x1080, 1600x900, 1280x720 and 640x480, ``nms_radius`` 4, JPEG and PNG decoders, 4 list workers, the keypoint-list fallbacks), ``"minimal"`` / ``False`` (640x480, radius 4 only), ``"all"`` (the default plus every ``nms_radius`` 0..9), or a dict that replaces keys of the default, e.g. ``{"nms_radii": (3, 4), "sizes": ((1024, 768),)}``. precompile_sizes: replaces ``warmup_variants["sizes"]`` when given (older name). precompile_nms_radii: added to ``warmup_variants["nms_radii"]`` (older name). trace_region_size: trace region in bytes when this call opens the device. local_files_only: load the weights from the HF cache only (no network). verbose: False (default) sets the tt-metal and loguru log levels to errors and warnings (``TT_LOGGER_LEVEL=Error``, ``LOGURU_LEVEL=WARNING``) when they are not set and ttnn is not imported yet. True keeps the tt-metal info log. Startup takes about 5 s with cached kernels (10-60 s when the kernels compile). """ if not verbose: _quiet_native_logs() spec = _warm.resolve_spec(warmup_variants) if precompile_sizes is not None: spec["sizes"] = _warm.normalize({"sizes": precompile_sizes})["sizes"] if precompile_nms_radii: spec["nms_radii"] = _warm.normalize({"nms_radii": tuple(spec["nms_radii"]) + tuple(precompile_nms_radii)})["nms_radii"] torch.set_grad_enabled(False) hf_model = _load_weights(pretrained_model_name_or_path, revision, local_files_only) border = int(hf_model.config.border_removal_distance) cfg = {"model_id": pretrained_model_name_or_path, "revision": revision if revision is not None else (DEFAULT_REVISION if pretrained_model_name_or_path == DEFAULT_MODEL_ID else None)} import ttnn from ._port.tt import fused_host as _fused from ._port.tt.superpoint_ttnn import TtSuperPoint owns = device is None mode = None if owns: device, mode = _device.open_device(device_id, dispatch=dispatch, trace_region_size=trace_region_size) tt_model = tt_in = self = None try: g = device.compute_with_storage_grid_size() cfg["compute_grid"] = f"{g.x}x{g.y}" cfg["dispatch"] = mode or "external" tt_model = TtSuperPoint(hf_model, device, input_height=INPUT_SIZE[0], input_width=INPUT_SIZE[1], fused=True, fused_stages=_fused.ALL_STAGES) tt_in = tt_model.allocate_input(batch_size=1) del hf_model self = cls(tt_model, tt_in, device, owns_device=owns, dispatch=mode, border_removal_distance=border, config=cfg) self._warmup_base(ttnn) t0 = time.perf_counter() self.warmup(**spec) self.config["warmup_s"]["variants"] = round(time.perf_counter() - t0, 2) self.config["warmup_variants"] = spec return self except BaseException: if self is not None: self._shutdown_pool() _release(device, tt_model, tt_in, close=owns) raise def _warmup_base(self, ttnn) -> None: """The server's warm-up (models/server/app.py ``_warmup_fused``): eager compile pass -> trace capture -> one replay. Then the request path of the 480x640 network frame at the default radius on synthetic images: readback buffers of every keypoint-list bucket, the host post-processing (top-k, sort, descriptor normalisation) and the input conversions.""" m, tt_in, dev = self._m, self._tt_in, self._device h, w = INPUT_SIZE t = {} with self._lock, torch.inference_mode(): t0 = time.perf_counter() m.run_fused(tt_in, torch.zeros(1, 3, h, w, dtype=torch.float32)) # compiles every kernel ttnn.synchronize_device(dev) t["compile_forward"] = time.perf_counter() - t0 t0 = time.perf_counter() m.capture_trace(tt_in, b=1) t["trace_capture"] = time.perf_counter() - t0 res = m.run_fused(tt_in, torch.zeros(1, 3, h, w, dtype=torch.float32)) ttnn.synchronize_device(dev) if not bool(torch.isfinite(res.descriptors_nchw).all()): raise RuntimeError("warm-up forward produced non-finite outputs") if m.trace_id is None: raise RuntimeError("warm-up did not capture the metal trace") t0 = time.perf_counter() self._warm_radius(m.nms_radius_traced if m.kpc_ready else self.defaults["nms_radius"], (w, h)) # host first-use costs of the other call options and input types (no device work) img = _warm.synthetic_image(64, 48, seed=3) to_plane(torch.from_numpy(img).permute(2, 0, 1).float() / 255.0) to_plane(img, bgr=True) self._run_plane(to_plane(_warm.synthetic_image(w, h, seed=4, density=2.0)), dict(self.defaults, max_keypoints=100, return_descriptors=False)) t["request_path"] = time.perf_counter() - t0 self.config["warmup_s"] = {k: round(v, 2) for k, v in t.items()} def _warm_calls(self, size: Tuple[int, int], nms_radius: int, n: int = 3) -> None: """``n`` full calls on synthetic images of ``size`` (W, H) with ~650 keypoints each.""" p = dict(self.defaults, nms_radius=int(nms_radius)) for i in range(n): self._run_plane(np.ascontiguousarray(_warm.synthetic_image(size[0], size[1], seed=11 + i, density=2.0)[..., 0]), p) def _warm_radius(self, r: int, size: Tuple[int, int] = (INPUT_SIZE[1], INPUT_SIZE[0])) -> None: m = self._m if m.kpc_ready and m.supports_device_nms_radius(r): with self._lock, torch.inference_mode(): m.warm_kpc_readback(r) # builds the radius variant when it is not the traced radius self._warm_calls(size, r) else: self._warm_calls(size, r, n=2) # host NMS: torch only, nothing to build def warmup(self, *, size: Optional[Tuple[int, int]] = None, sizes: Iterable[Tuple[int, int]] = (), nms_radius: Optional[int] = None, nms_radii: Iterable[int] = (), decoder: Optional[str] = None, decoders: Iterable[str] = (), num_workers: int = 0, fallbacks: bool = False) -> Dict[str, float]: """Prepare more per-request variants now (``from_pretrained`` already prepared the ``warmup_variants``). Idempotent: a variant that is ready is skipped. Args: size / sizes: source image sizes ``(width, height)``. A size in the device-resize range gets its resize trace (approximately 10 ms each with cached kernels; at most 8 sizes are kept, the least recently used one is released). Other sizes use the host resize. nms_radius / nms_radii: 1..8 build the device NMS trace of that radius (approximately 50 ms each with cached kernels, 1-3 s when its kernels compile); 0 and > 8 warm the host NMS. decoder / decoders: image file formats to decode once: ``"jpeg"``, ``"png"``, ``"bmp"``, ``"webp"``, ``"tiff"``. num_workers: start this many host threads for list / iterator calls. fallbacks: warm the exact fallbacks of the keypoint list (more than 1024 NMS candidates, candidate-slot overflow) for the prepared radii. Returns: seconds spent per step (an empty dict when everything was ready). """ self._check_open() spec = _warm.variant_kwargs(size=size, sizes=sizes, nms_radius=nms_radius, nms_radii=nms_radii, decoder=decoder, decoders=decoders, num_workers=num_workers, fallbacks=fallbacks) m, h, w = self._m, INPUT_SIZE[0], INPUT_SIZE[1] spent: Dict[str, float] = {} def step(key, fn): if key in self._warm_done: return t0 = time.perf_counter() fn() self._warm_done.add(key) spent["/".join(str(k) for k in key)] = round(time.perf_counter() - t0, 3) self._warm_done.add(("nms_radius", m.nms_radius_traced if m.kpc_ready else self.defaults["nms_radius"])) for r in spec["nms_radii"]: step(("nms_radius", r), lambda r=r: self._warm_radius(r)) if len(spec["sizes"]) > m.RSZ_MAX_VARIANTS: LOG.warning("%d warm-up sizes, but only %d device-resize variants are kept", len(spec["sizes"]), m.RSZ_MAX_VARIANTS) for sw, sh in spec["sizes"]: if (sh, sw) == INPUT_SIZE: continue on_device = m.device_resize and m.supports_device_resize(sw, sh) if on_device and (sw, sh) in m._rsz_vars: continue # trace resident (LRU), nothing to do key = ("size", sw, sh) if on_device: self._warm_done.discard(key) # released by the LRU: build it again step(key, lambda sw=sw, sh=sh: self._warm_calls((sw, sh), self.defaults["nms_radius"], n=2)) self._fallbacks_on |= spec["fallbacks"] # later device radii get their fallbacks warmed too if self._fallbacks_on: radii = sorted(r for k, r in (x for x in self._warm_done if x[0] == "nms_radius") if 1 <= r <= 8) for r in radii: step(("fallbacks", r), lambda r=r: self._warm_fallbacks(r)) if spec["decoders"]: # freed Pillow image buffers are kept for the next decode (2 per decoding thread) _warm.keep_pillow_blocks(2 * (1 + max(spec["num_workers"], self._pool_workers()))) if spec["num_workers"] > 0: step(("num_workers", spec["num_workers"]), lambda: self._warm_pool(spec["num_workers"], spec["decoders"])) # last: the large tensors of the steps above can shrink the heap the decodes grow big = self._largest_warm_size() for fmt in spec["decoders"]: step(("decoder", fmt), lambda fmt=fmt: _decode_n(_warm.encoded_image(fmt, *big))) if spent: # end on the default request path: the steps above leave the host caches cold for it self._warm_calls((INPUT_SIZE[1], INPUT_SIZE[0]), self.defaults["nms_radius"], n=2) return spent def _warm_fallbacks(self, r: int) -> None: """More than 1024 NMS candidates on a dense synthetic image (host top-k + sampling trace), then the host post-processing of the resident maps (the slot-overflow fallback).""" m = self._m img = np.ascontiguousarray(_warm.synthetic_image(INPUT_SIZE[1], INPUT_SIZE[0], seed=21, density=6.0)[..., 0]) self._run_plane(img, dict(self.defaults, nms_radius=r)) self._run_plane(img, dict(self.defaults, nms_radius=r, max_keypoints=-1)) if m.kpc_ready: var = m._variant(r) with self._lock, torch.inference_mode(): m._keypoints_from_resident(self.defaults["keypoint_threshold"], self.defaults["max_keypoints"], self.border_removal_distance, True, None if var is None else var.nms_map) def _pool_workers(self) -> int: return 0 if self._pool is None else int(self._pool._max_workers) def _get_pool(self, num_workers: int) -> ThreadPoolExecutor: with self._pool_lock: if self._pool is None or self._pool._max_workers < num_workers: old, self._pool = self._pool, ThreadPoolExecutor(max_workers=num_workers, thread_name_prefix="tt_superpoint_io") if old is not None: old.shutdown(wait=False) return self._pool def _largest_warm_size(self) -> Tuple[int, int]: sizes = [(k[1], k[2]) for k in self._warm_done if k[0] == "size"] + [(INPUT_SIZE[1], INPUT_SIZE[0])] return max(sizes, key=lambda s: s[0] * s[1]) def _warm_pool(self, num_workers: int, decoders: Sequence[str] = ()) -> None: """Start the host threads of list calls; each thread decodes one image file of the largest warm size (its malloc arena grows once), then one list call on synthetic images.""" pool = self._get_pool(num_workers) data = _warm.encoded_image(decoders[0] if decoders else "jpeg", *self._largest_warm_size()) barrier = threading.Barrier(num_workers + 1) def task(): barrier.wait(30) # one task on each thread _decode_n(data) futs = [pool.submit(task) for _ in range(num_workers)] barrier.wait(30) for f in futs: f.result() imgs = [_warm.synthetic_image(INPUT_SIZE[1], INPUT_SIZE[0], seed=31 + i, density=2.0) for i in range(2)] list(self.iter(imgs, num_workers=num_workers)) def _shutdown_pool(self) -> None: with self._pool_lock: if self._pool is not None: self._pool.shutdown(wait=True) self._pool = None # ------------------------------------------------------------------ inference def _check_open(self) -> None: if self._closed: raise RuntimeError("SuperPoint model is closed") def _infer(self, plane: np.ndarray, *, max_keypoints: int, keypoint_threshold: float, nms_radius: int, return_descriptors: bool): """The server's ``_infer_fused`` on an (H, W) uint8 plane at the source size -> (kp (N, 2) network-frame xy, scores (N,), desc (N, 256) | None, device_nms).""" from ._port.tt import postprocess as _post m, tt_in, border = self._m, self._tt_in, self.border_removal_distance if not plane.flags.writeable and plane.shape == INPUT_SIZE: plane = np.array(plane) # torch.as_tensor warns on read-only arrays host_in = m.prepare_source(plane) if m.kpc_ready and m.supports_device_nms_radius(nms_radius): with self._lock, torch.inference_mode(): kp, sc, desc = m.run_fused_keypoints_kpc( tt_in, host_in, keypoint_threshold=keypoint_threshold, max_keypoints=max_keypoints, border_removal_distance=border, with_descriptors=return_descriptors, nms_radius=nms_radius, ) return kp, sc, desc, True with self._lock, torch.inference_mode(): res = m.run_fused_prepared(tt_in, host_in, nms_radius=nms_radius) with torch.inference_mode(): if res.nms_map is not None: kp, sc, desc = _post.postprocess_from_nms_map( res.nms_map, res.descriptors_nchw, keypoint_threshold=keypoint_threshold, max_keypoints=max_keypoints, border_removal_distance=border, with_descriptors=return_descriptors, )[0] return kp, sc, desc, True kp, sc, desc = _post.postprocess_keypoints( res.scores_nchw, res.descriptors_nchw, nms_radius=nms_radius, keypoint_threshold=keypoint_threshold, max_keypoints=max_keypoints, border_removal_distance=border, with_descriptors=return_descriptors, )[0] return kp, sc, desc, False def _params(self, max_keypoints, keypoint_threshold, nms_radius, return_descriptors) -> Dict[str, Any]: p = dict(self.defaults) if max_keypoints is not None: p["max_keypoints"] = int(max_keypoints) if keypoint_threshold is not None: p["keypoint_threshold"] = float(keypoint_threshold) if nms_radius is not None: p["nms_radius"] = int(nms_radius) if return_descriptors is not None: p["return_descriptors"] = bool(return_descriptors) return validate_params(**p) def _run_plane(self, plane: np.ndarray, p: Dict[str, Any], t_pre: float = 0.0) -> SuperPointOutput: self._check_open() orig_h, orig_w = plane.shape t0 = time.perf_counter() kp, sc, desc, device_nms = self._infer(plane, **p) t1 = time.perf_counter() order = torch.argsort(sc, descending=True) # same order as the server response kp, sc = kp[order], sc[order] if desc is not None: desc = torch.index_select(desc, 0, order).float() sx, sy = orig_w / INPUT_SIZE[1], orig_h / INPUT_SIZE[0] kp_orig = (kp.double() * torch.tensor([sx, sy], dtype=torch.float64)).float() t2 = time.perf_counter() return SuperPointOutput( keypoints=kp_orig, scores=sc.float(), descriptors=desc if p["return_descriptors"] else None, image_size=(int(orig_h), int(orig_w)), scale=(sx, sy), device_nms=bool(device_nms), timing_ms={"preprocess": round(t_pre * 1e3, 3), "device_forward": round((t1 - t0) * 1e3, 3), "postprocess": round((t2 - t1) * 1e3, 3)}, ) def __call__( self, images: Union[ImageLike, Sequence[ImageLike]], *, max_keypoints: Optional[int] = None, keypoint_threshold: Optional[float] = None, nms_radius: Optional[int] = None, return_descriptors: Optional[bool] = None, bgr: bool = False, num_workers: int = 4, ) -> Union[SuperPointOutput, List[SuperPointOutput]]: """Detect keypoints and compute descriptors. Args: images: one image, or a list / tuple of images, or a 4-D batch array / tensor ((B, H, W, C) numpy, (B, C, H, W) torch). An image is a file path, encoded bytes, a ``PIL.Image``, a numpy array or a torch tensor of any size (see :func:`tt_superpoint.inputs.to_plane`). The model reads channel 0 (R; the gray plane of a grayscale image), resized to 480x640. max_keypoints: keep the top-k by score (default 1024); -1 keeps every keypoint above the threshold. keypoint_threshold: minimum score after NMS, in [0, 1] (default 0.005). nms_radius: NMS radius in network-frame pixels, 0..32 (default 4; 1..8 run on the device, 0 and > 8 on the host). return_descriptors: compute the 256-d descriptors (default True). bgr: arrays are BGR (OpenCV); the R channel is then the last one. num_workers: host threads that decode / convert the next images of a list while the device runs the current one (lists only; 0 = no threads). Returns: one :class:`SuperPointOutput` for one image; a list of them, in input order, for a list or a batch. """ p = self._params(max_keypoints, keypoint_threshold, nms_radius, return_descriptors) self._check_open() if is_single_image(images): t0 = time.perf_counter() plane = to_plane(images, bgr=bgr) return self._run_plane(plane, p, time.perf_counter() - t0) return list(self.iter(split_batch(images), bgr=bgr, num_workers=num_workers, **p)) def iter(self, images: Iterable[ImageLike], *, bgr: bool = False, num_workers: int = 4, **params) -> Iterator[SuperPointOutput]: """Yield one :class:`SuperPointOutput` per image of ``images`` (any iterable, e.g. a generator over a video or a directory), in order. Up to ``num_workers`` images are decoded / converted on host threads ahead of the device; the device runs one image at a time, so the outputs are the same as calling the model on each image. Keyword arguments: the per-call parameters of :meth:`__call__`.""" p = self._params(params.pop("max_keypoints", None), params.pop("keypoint_threshold", None), params.pop("nms_radius", None), params.pop("return_descriptors", None)) if params: raise TypeError(f"unexpected keyword argument(s): {sorted(params)}") def load(img): t0 = time.perf_counter() return to_plane(img, bgr=bgr), time.perf_counter() - t0 if num_workers <= 0: for img in images: plane, t = load(img) yield self._run_plane(plane, p, t) return it = iter(images) ex = self._get_pool(num_workers) # kept between calls (started by the warm-up) pending = [] try: for img in it: pending.append(ex.submit(load, img)) if len(pending) > num_workers: break while pending: plane, t = pending.pop(0).result() nxt = next(it, _END) if nxt is not _END: pending.append(ex.submit(load, nxt)) yield self._run_plane(plane, p, t) finally: for f in pending: f.cancel() # ------------------------------------------------------------------ info / lifetime @property def device(self): """The ttnn device the model runs on.""" return self._device @property def input_size(self) -> Tuple[int, int]: """Network frame (height, width) = (480, 640).""" return INPUT_SIZE def close(self) -> None: """Release the traces and device tensors, and close the device if this object opened it. Safe to call more than once.""" if self._closed: return self._closed = True self._shutdown_pool() with self._lock: _release(self._device, self._m, self._tt_in, close=self._owns_device) self._m = self._tt_in = None def __enter__(self) -> "SuperPoint": return self def __exit__(self, *exc) -> None: self.close() def __repr__(self) -> str: state = "closed" if self._closed else f"grid {self.config.get('compute_grid')}, dispatch {self.config.get('dispatch')}" return f"SuperPoint({self.config.get('model_id')!r}, {state})" _END = object() def _decode_n(data: bytes, n: int = 3) -> None: """Decode an image file ``n`` times. The first decode imports the Pillow plug-in. The next ones grow the malloc heap of this thread to the image size: glibc serves the first large buffers with mmap and raises its mmap threshold only when such a buffer is freed.""" for _ in range(n): to_plane(data) def _quiet_native_logs() -> None: """Errors and warnings only from tt-metal (C++) and loguru (ttnn Python), unless the user set the levels. The levels are read when ttnn is imported, so this has no effect after that.""" import os import sys if "ttnn" in sys.modules: return os.environ.setdefault("TT_LOGGER_LEVEL", "Error") os.environ.setdefault("LOGURU_LEVEL", "WARNING") def validate_params(*, max_keypoints: int, keypoint_threshold: float, nms_radius: int, return_descriptors: bool) -> Dict[str, Any]: """Range checks of the per-call parameters (the HTTP server's limits).""" if not -1 <= max_keypoints <= MAX_KEYPOINTS_CAP: raise ValueError(f"max_keypoints must be in -1..{MAX_KEYPOINTS_CAP}, got {max_keypoints}") if not 0.0 <= keypoint_threshold <= 1.0: raise ValueError(f"keypoint_threshold must be in [0, 1], got {keypoint_threshold}") if not 0 <= nms_radius <= MAX_NMS_RADIUS: raise ValueError(f"nms_radius must be in 0..{MAX_NMS_RADIUS}, got {nms_radius}") return dict(max_keypoints=max_keypoints, keypoint_threshold=keypoint_threshold, nms_radius=nms_radius, return_descriptors=return_descriptors) def _load_weights(name_or_path: str, revision: Optional[str], local_files_only: bool): """HF ``SuperPointForKeypointDetection`` (fp32 reference whose weights the device model uses), from the HF cache; retries offline when the Hub is unreachable (as the server).""" import os from transformers import SuperPointForKeypointDetection if os.path.isdir(str(name_or_path)): model = SuperPointForKeypointDetection.from_pretrained(name_or_path) else: rev = revision if revision is not None else (DEFAULT_REVISION if name_or_path == DEFAULT_MODEL_ID else None) try: model = SuperPointForKeypointDetection.from_pretrained(name_or_path, revision=rev, local_files_only=local_files_only) except Exception as e: # noqa: BLE001 - network down, snapshot cached if local_files_only or rev is None: raise LOG.warning("Hub resolution failed (%s: %s); retrying from the local HF cache", type(e).__name__, e) model = SuperPointForKeypointDetection.from_pretrained(name_or_path, revision=rev, local_files_only=True) return model.eval() def _release(device, tt_model, tt_in, *, close: bool) -> None: if device is None and tt_in is None: if tt_model is not None: tt_model.release() return import ttnn try: if device is not None: ttnn.synchronize_device(device) if tt_model is not None: tt_model.release() if tt_in is not None: ttnn.deallocate(tt_in) except Exception as e: # noqa: BLE001 - never mask the original error LOG.warning("device tensor release failed: %s", e) finally: if close and device is not None: ttnn.close_device(device)