# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license """Shared ReID encoder used by BoT-SORT, Deep OC-SORT, and TrackTrack. * `.pt` YOLO checkpoints — loaded via `YOLO()`; embeddings are pulled from the second-to-last layer through the predictor's `embed=[...]` argument (works with classification and ReID backbones). * Compatible exported models (`.torchscript`, `.onnx`, `.engine`, OpenVINO model directories, …) — loaded via `AutoBackend`; the model is expected to output the embedding tensor directly. """ from __future__ import annotations import numpy as np import torch from ultralytics.nn.autobackend import AutoBackend from ultralytics.utils.ops import xywh2xyxy from ultralytics.utils.plotting import save_one_box from ultralytics.utils.torch_utils import smart_inference_mode REID_ASSETS = frozenset(f"yolo26{k}-reid.onnx" for k in "nsmlx") class ReID: """ReID encoder. Routes `.pt` to the YOLO predictor path; everything else to `AutoBackend`.""" @smart_inference_mode(False) def __init__(self, model: str, imgsz: int = 224, device: str | torch.device | None = None, fp16: bool = False): """Initialize encoder for re-identification. Args: model (str): Path to a ReID model. `.pt` runs through the YOLO predictor (embed-layer extraction); other extensions go through `AutoBackend`. imgsz (int): Square input size used for crop preprocessing on the AutoBackend path. Overridden by the model's own static input size when one is detected. device (str | torch.device | None): Inference device; defaults to CUDA if available. fp16 (bool): Request half precision on the AutoBackend path when the backend supports it. Ignored for `.pt` models; models exported with FP16 inputs run in half precision regardless. """ self.imgsz = imgsz self.batch_size = None self.device = ( torch.device(device) if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu") ) self.is_pt = str(model).endswith(".pt") if self.is_pt: from ultralytics import YOLO self.model = YOLO(model) # Set up the embedding predictor that __call__ reuses for crops. self.model.embed(np.zeros((32, 32, 3), dtype=np.uint8), device=self.device, verbose=False, save=False) self.fp16 = False else: from pathlib import Path if Path(str(model)).name in REID_ASSETS: from ultralytics.utils.downloads import attempt_download_asset model = attempt_download_asset(str(model)) self.model = AutoBackend(str(model), device=self.device, fp16=fp16, verbose=False) self.fp16 = self.model.fp16 # Get model's input size for a fixed batch and crop size or detect dynamic batch and crop sizes. session = getattr(self.model, "session", None) shape = session.get_inputs()[0].shape if session is not None else () if len(shape) == 4: if isinstance(shape[0], int) and shape[0] > 0: self.batch_size = shape[0] if isinstance(shape[2], int) and shape[2] > 0: self.imgsz = shape[2] @staticmethod def _crop_detections(img: np.ndarray, dets: np.ndarray) -> list[np.ndarray]: """Crop detection regions from image, converting xywh to xyxy first. Args: img (np.ndarray): BGR image. dets (np.ndarray): Detections in xywh format (first 4 columns used). Returns: (list[np.ndarray]): Cropped image patches. """ return [save_one_box(det, img, save=False) for det in xywh2xyxy(torch.from_numpy(dets[:, :4]))] def _crops_to_tensor(self, crops: list[np.ndarray]) -> torch.Tensor: """Stack a list of valid image crops into a normalized BCHW float tensor at self.imgsz.""" batch = torch.empty(len(crops), 3, self.imgsz, self.imgsz, dtype=torch.float32) for i, c in enumerate(crops): t = torch.from_numpy(np.ascontiguousarray(c[..., ::-1])).permute(2, 0, 1).unsqueeze(0).float() / 255.0 batch[i] = torch.nn.functional.interpolate( t, size=(self.imgsz, self.imgsz), mode="bilinear", align_corners=False )[0] batch = batch.to(self.device) return batch.half() if self.fp16 else batch @torch.no_grad() def __call__(self, img: np.ndarray, dets: np.ndarray) -> list[np.ndarray | None]: """Extract embeddings for detected objects. Args: img (np.ndarray): BGR image containing the detections. dets (np.ndarray): Detections in xywh format (first 4 columns used). Returns: (list[np.ndarray | None]): One embedding per detection, or None where the crop is empty. """ crops = self._crop_detections(img, dets) valid = [bool(c.size) for c in crops] valid_crops = [crop for crop, keep in zip(crops, valid) if keep] if not valid_crops: return [None] * len(crops) if self.is_pt: feats = self.model.predictor(valid_crops) if len(feats) != len(valid_crops) and feats[0].shape[0] == len(valid_crops): feats = feats[0] # batched prediction with non-PyTorch backend valid_feats = [f.cpu().numpy() for f in feats] else: batch = self._crops_to_tensor(valid_crops) bs, n = self.batch_size, batch.shape[0] if bs is None or n == bs: feats = self.model(batch) else: # fixed-batch model (e.g. static ONNX): run in chunks of bs, padding the last partial chunk outs = [] for s in range(0, n, bs): chunk = batch[s : s + bs] if chunk.shape[0] < bs: chunk = torch.cat([chunk, chunk[-1:].expand(bs - chunk.shape[0], *chunk.shape[1:])], 0) outs.append(self.model(chunk)) feats = torch.cat(outs, 0)[:n] valid_feats = [f.cpu().numpy() for f in feats] valid_feats = iter(valid_feats) return [next(valid_feats) if keep else None for keep in valid] def build_encoder(with_reid: bool, model: str | None, device: str | torch.device | None = None): """Return a ReID encoder, the native-features pass-through, or None. Args: with_reid (bool): Whether ReID is enabled at all. model (str | None): `"auto"` returns a callable that converts pre-extracted backbone features to numpy arrays; any other value loads a `ReID` model from that path. Ignored when `with_reid` is False. device (str | torch.device | None): Inference device for the ReID model; defaults to CUDA if available. Returns: (Callable | None): A `(img, dets) -> list[np.ndarray | None]` encoder, or None when ReID is disabled. """ if not with_reid: return None if model == "auto": def _auto_encoder(feats, _dets): if isinstance(feats, np.ndarray): return [f for f in feats] return [f.cpu().numpy() for f in feats] return _auto_encoder return ReID(model, device=device) def smooth_feature( feat: np.ndarray, smooth: np.ndarray | None, alpha: float ) -> tuple[np.ndarray | None, np.ndarray | None]: """L2-normalize `feat` and blend it into `smooth` via exponential moving average. Args: feat (np.ndarray): New (un-normalized) appearance feature. smooth (np.ndarray | None): Current smoothed feature, or None on the first update. alpha (float): EMA weight on the existing `smooth` (``1.0`` keeps it unchanged). Returns: curr (np.ndarray | None): The normalized float32 feature, or None when `feat` has zero norm. smooth (np.ndarray | None): The updated, normalized float32 feature. """ feat = np.asarray(feat, dtype=np.float32) # the stored state is float32 whatever dtype the ReID backend returned norm = np.linalg.norm(feat) if norm < 1e-12: # zero-norm feature has no appearance info; signal the caller to keep its current features return None, smooth feat = feat / norm if smooth is None: return feat, feat.copy() smooth = alpha * smooth + (1 - alpha) * feat return feat, smooth / np.linalg.norm(smooth)