Reference for ultralytics/trackers/utils/reid.py#

class ReID:
    """ReID encoder. Routes `.pt` to the YOLO predictor path; everything else to `AutoBackend`."""

    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): Use half precision when the backend supports it.
        """
        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)
            # Initialize predictor with embed=[idx] so subsequent calls return embeddings.
            self.model(embed=[len(self.model.model.model) - 2], 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]

@torch.no_grad()
def __call__(self, img: np.ndarray, dets: np.ndarray) -> list[np.ndarray | None]:
    """Extract embeddings for detected objects."""
    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]

@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

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)