reid 0.1.0.dev0__py3-none-any.whl

reid/__init__.py

@@ -0,0 +1,47 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""re-ID: appearance re-identification models and gallery evaluation.
+
+The public surface is intentionally narrow. Advanced helpers live in
+submodules: checkpoint loading in :mod:`reid.loaders`, the curated alias
+catalog and ``ModelCard`` in :mod:`reid.catalog`, shared dataset types and
+benchmark loaders in :mod:`reid.data`, and architecture dispatch in
+:mod:`reid.architectures`.
+"""
+
+from __future__ import annotations
+
+from reid.architectures import list_architectures
+from reid.catalog import DEFAULT_MODEL, FASTREID_MOT17_SBS50, list_aliases
+from reid.eval import (
+    ReIDEvaluator,
+    ReIDMetrics,
+    ReIDResult,
+    ReIDSplit,
+    compute_reid_metrics,
+    load_market1501,
+    load_msmt17,
+)
+from reid.model import ReIDModel
+from reid.preprocessing import Interpolation, ReIDPreprocessing
+
+__all__ = [
+    "DEFAULT_MODEL",
+    "FASTREID_MOT17_SBS50",
+    "Interpolation",
+    "ReIDEvaluator",
+    "ReIDMetrics",
+    "ReIDModel",
+    "ReIDPreprocessing",
+    "ReIDResult",
+    "ReIDSplit",
+    "compute_reid_metrics",
+    "list_aliases",
+    "list_architectures",
+    "load_market1501",
+    "load_msmt17",
+]

reid/architectures/__init__.py

@@ -0,0 +1,23 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""ReID architecture public API."""
+
+from __future__ import annotations
+
+from reid.architectures.registry import (
+    build_architecture,
+    checkpoint_remap_for_architecture,
+    default_preprocessing_for_architecture,
+    list_architectures,
+)
+
+__all__ = [
+    "build_architecture",
+    "checkpoint_remap_for_architecture",
+    "default_preprocessing_for_architecture",
+    "list_architectures",
+]

reid/architectures/fastreid_sbs.py

@@ -0,0 +1,141 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+#
+# Adapted from JDAI-CV/fast-reid (Apache-2.0)
+# Copyright 2019 JD.com Inc. JD AI
+# Source: https://github.com/JDAI-CV/fast-reid
+#   - GeM: fastreid/layers/pooling.py (GeneralizedMeanPooling / GeneralizedMeanPoolingP)
+#   - SBS head layout: fastreid/modeling/heads/embedding_head.py
+#   - last_stride=1 ResNeSt semantics: fastreid/modeling/backbones/resnest.py
+# Also used via BoT-SORT's FastReIDInterface (NirAharon/BoT-SORT, MIT).
+# ------------------------------------------------------------------------
+
+"""FastReID Strong Baseline (SBS) inference stack for BoT-SORT checkpoints.
+
+Uses ``timm`` ``resnest50d`` with a small FastReID ``last_stride=1`` patch on
+``layer4[0]``, then Generalized Mean Pooling and a BatchNorm neck — the same
+inference path as FastReID ``EmbeddingHead`` in eval.
+"""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import timm
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from reid.preprocessing import ReIDPreprocessing
+
+FASTREID_SBS_ARCHITECTURE = "fastreid_sbs_resnest50"
+FASTREID_SBS_EMBED_DIM = 2048
+
+
+def _patch_resnest50d_for_fastreid_last_stride(backbone: nn.Module) -> nn.Module:
+    """Align ``timm`` ``resnest50d`` (``output_stride=16``) with FastReID ResNeSt.
+
+    FastReID ``LAST_STRIDE=1`` keeps layer4 spatial stride at 1 via:
+    - downsample avg-pool kernel 1 (identity spatial size), and
+    - average-downsampling (AVD) after the Split-Attention conv with stride 1.
+
+    ``timm``'s dilated / stride-16 path instead uses ``AvgPool2dSame(2, stride=1)``
+    and leaves ``avd_last`` unset on ``layer4[0]``, which shifts embeddings even
+    when checkpoint keys load 100%.
+    """
+    backbone_any = cast(Any, backbone)
+    block = backbone_any.layer4[0]
+    block.downsample[0] = nn.AvgPool2d(kernel_size=1, stride=1, padding=0)
+    block.avd_last = nn.AvgPool2d(kernel_size=3, stride=1, padding=1)
+    return backbone
+
+
+class GeneralizedMeanPooling(nn.Module):
+    """GeM pooling used by FastReID SBS (learnable exponent *p*).
+
+    ``p`` defaults to 3.0 only until checkpoint load; MOT17 SBS-S50 stores the
+    trained value under ``heads.pool_layer.p`` (≈1.72). See
+    :func:`remap_fastreid_sbs_state_dict`.
+    """
+
+    def __init__(self, p: float = 3.0, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1) * p)
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.clamp(min=self.eps).pow(self.p)
+        x = F.adaptive_avg_pool2d(x, 1)
+        return x.pow(1.0 / self.p)
+
+
+class FastReIDSBSResNeSt50(nn.Module):
+    """ResNeSt50 SBS re-ID encoder matching BoT-SORT / FastReID MOT17 checkpoints.
+
+    Inference path: patched ``timm`` ResNeSt-50 → GeM (:attr:`pool`) → BNNeck
+    (:attr:`bottleneck`). Weights for all three stages load from a BoT-SORT
+    ``.pth`` via :func:`remap_fastreid_sbs_state_dict`.
+    """
+
+    @classmethod
+    def default_preprocessing(cls) -> ReIDPreprocessing:
+        """Default crop/tensor recipe for FastReID SBS (BoT-SORT FastReIDInterface).
+
+        Config (aligned with BoT-SORT / FastReID MOT17 SBS inference):
+        - input size: 384x128 (H, W)
+        - resize: stretch (cv2.resize, no letterbox)
+        - colour: BGR crops converted to RGB
+        - normalize: ImageNet mean/std
+        """
+        return ReIDPreprocessing(
+            input_size=(384, 128),
+            resize_mode="stretch",
+            to_rgb=True,
+        )
+
+    def __init__(self) -> None:
+        super().__init__()
+        backbone = timm.create_model(
+            "resnest50d",
+            pretrained=False,
+            num_classes=0,
+            global_pool="",
+            output_stride=16,
+        )
+        self.backbone = _patch_resnest50d_for_fastreid_last_stride(backbone)
+        self.pool = GeneralizedMeanPooling()
+        self.bottleneck = nn.BatchNorm1d(FASTREID_SBS_EMBED_DIM)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.backbone(x)
+        x = self.pool(x).flatten(1)
+        return self.bottleneck(x)
+
+
+def build_fastreid_sbs_resnest50(*, num_classes: int = 0, pretrained: bool = False) -> FastReIDSBSResNeSt50:
+    """Build the BoT-SORT FastReID SBS ResNeSt50 encoder (weights loaded separately)."""
+    del num_classes, pretrained
+    return FastReIDSBSResNeSt50()
+
+
+def remap_fastreid_sbs_state_dict(state_dict: dict) -> dict:
+    """Map BoT-SORT / FastReID SBS checkpoint keys onto :class:`FastReIDSBSResNeSt50`.
+
+    Renames FastReID ``heads.*`` keys to ``pool.*`` / ``bottleneck.*`` and passes
+    ``backbone.*`` through unchanged. Skips ``heads.weight`` (classifier). GeM
+    ``heads.pool_layer.p`` becomes ``pool.p`` and replaces the 3.0 init default.
+    """
+    mapped: dict = {}
+    for key, value in state_dict.items():
+        key = key[7:] if key.startswith("module.") else key
+        if key.startswith("backbone."):
+            mapped[key] = value
+        elif key == "heads.pool_layer.p":
+            mapped["pool.p"] = value
+        elif key.startswith("heads.bottleneck.0."):
+            mapped["bottleneck." + key[len("heads.bottleneck.0.") :]] = value
+        # Skip heads.weight (identity classifier; unused at inference).
+    return mapped

reid/architectures/osnet.py

@@ -0,0 +1,314 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+#
+# Adapted from KaiyangZhou/deep-person-reid
+# Copyright (c) 2018-2021 Kaiyang Zhou
+# Licensed under the MIT License
+# Source: https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/models/osnet.py
+# Paper: Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.
+# ------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from reid.preprocessing import ReIDPreprocessing
+
+# --------------------------------------------------------------------------- #
+# Basic layers
+# --------------------------------------------------------------------------- #
+
+
+class _ConvBnRelu(nn.Module):
+    """Conv2d → BatchNorm2d → ReLU building block."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        groups: int = 1,
+    ) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False,
+            groups=groups,
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.bn(self.conv(x)))
+
+
+class _Conv1x1(_ConvBnRelu):
+    """1x1 conv + BN + ReLU."""
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1, groups: int = 1) -> None:
+        super().__init__(in_channels, out_channels, 1, stride=stride, padding=0, groups=groups)
+
+
+class _Conv1x1Linear(nn.Module):
+    """1x1 conv + BN (no non-linearity)."""
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 1, stride=stride, padding=0, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.bn(self.conv(x))
+
+
+class _LightConv3x3(nn.Module):
+    """Lightweight 3x3 conv: 1x1 linear + depthwise 3x3 + BN + ReLU."""
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride=1, padding=1, bias=False, groups=out_channels)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.bn(self.conv2(self.conv1(x))))
+
+
+# --------------------------------------------------------------------------- #
+# Omni-scale building blocks
+# --------------------------------------------------------------------------- #
+
+
+class _ChannelGate(nn.Module):
+    """Channel-wise gating via squeeze-and-excitation."""
+
+    def __init__(self, in_channels: int, reduction: int = 16) -> None:
+        super().__init__()
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc1 = nn.Conv2d(in_channels, in_channels // reduction, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.fc2 = nn.Conv2d(in_channels // reduction, in_channels, 1, bias=True)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        g = self.sigmoid(self.fc2(self.relu(self.fc1(self.global_avgpool(x)))))
+        return x * g
+
+
+class OSBlock(nn.Module):
+    """Omni-scale feature learning residual block.
+
+    Aggregates features from four parallel streams with different effective
+    receptive-field depths (1, 2, 3, 4 stacked lightweight 3x3 convs) gated
+    by a shared channel-wise gate before the residual addition.
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, bottleneck_reduction: int = 4) -> None:
+        super().__init__()
+        mid = out_channels // bottleneck_reduction
+        self.conv1 = _Conv1x1(in_channels, mid)
+        self.conv2a = _LightConv3x3(mid, mid)
+        self.conv2b = nn.Sequential(_LightConv3x3(mid, mid), _LightConv3x3(mid, mid))
+        self.conv2c = nn.Sequential(_LightConv3x3(mid, mid), _LightConv3x3(mid, mid), _LightConv3x3(mid, mid))
+        self.conv2d = nn.Sequential(
+            _LightConv3x3(mid, mid),
+            _LightConv3x3(mid, mid),
+            _LightConv3x3(mid, mid),
+            _LightConv3x3(mid, mid),
+        )
+        self.gate = _ChannelGate(mid)
+        self.conv3 = _Conv1x1Linear(mid, out_channels)
+        self.downsample = _Conv1x1Linear(in_channels, out_channels) if in_channels != out_channels else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        x1 = self.conv1(x)
+        x2 = (
+            self.gate(self.conv2a(x1))
+            + self.gate(self.conv2b(x1))
+            + self.gate(self.conv2c(x1))
+            + self.gate(self.conv2d(x1))
+        )
+        x3 = self.conv3(x2)
+        if self.downsample is not None:
+            identity = self.downsample(identity)
+        return F.relu(x3 + identity)
+
+
+# --------------------------------------------------------------------------- #
+# Full OSNet architecture
+# --------------------------------------------------------------------------- #
+
+
+class OSNet(nn.Module):
+    """Omni-Scale Network for instance re-identification.
+
+    Architecture from:
+        Zhou et al. *Omni-Scale Feature Learning for Person Re-Identification*. ICCV, 2019.
+        Zhou et al. *Learning Generalisable Omni-Scale Representations for Person Re-Identification*. TPAMI, 2021.
+
+    At inference (``model.eval()``), and whenever no classifier was built
+    (``num_classes <= 0``), ``forward()`` returns the ``feature_dim`` embedding.
+    Classifier logits are returned only in ``train()`` mode with a head.
+
+    Args:
+        num_classes: Number of identity classes for the classification head.
+            Use ``0`` (or less) to omit the head for inference-only models.
+        blocks: List of block classes for each stage.
+        layers: Number of blocks per stage.
+        channels: Channel widths ``[stem, stage1, stage2, stage3]``.
+        feature_dim: Dimensionality of the final embedding vector.
+    """
+
+    @classmethod
+    def default_preprocessing(cls) -> ReIDPreprocessing:
+        """Default crop/tensor recipe for OSNet (torchreid / ImageNet norms).
+
+        Config (matches common OSNet x1.0 zoo checkpoints):
+        - input size: 256x128 (H, W)
+        - resize: stretch (no letterbox)
+        - colour: BGR crops converted to RGB
+        - normalize: ImageNet mean/std
+        """
+        return ReIDPreprocessing(
+            input_size=(256, 128),
+            resize_mode="stretch",
+            to_rgb=True,
+        )
+
+    def __init__(
+        self,
+        num_classes: int,
+        blocks: list[type[nn.Module]],
+        layers: list[int],
+        channels: list[int],
+        feature_dim: int = 512,
+    ) -> None:
+        super().__init__()
+        if len(blocks) != len(layers) or len(layers) != len(channels) - 1:
+            raise ValueError("blocks, layers, and channels must have consistent lengths")
+
+        self.feature_dim = feature_dim
+        self.conv1 = _ConvBnRelu(3, channels[0], 7, stride=2, padding=3)
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.conv2 = self._make_layer(blocks[0], layers[0], channels[0], channels[1], downsample=True)
+        self.conv3 = self._make_layer(blocks[1], layers[1], channels[1], channels[2], downsample=True)
+        self.conv4 = self._make_layer(blocks[2], layers[2], channels[2], channels[3], downsample=False)
+        self.conv5 = _Conv1x1(channels[3], channels[3])
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channels[3], feature_dim),
+            nn.BatchNorm1d(feature_dim),
+            nn.ReLU(inplace=True),
+        )
+        # Inference builds use num_classes=0 and omit the head so checkpoints
+        # that drop classifier.* can load with required_match_fraction=1.0.
+        if num_classes > 0:
+            self.classifier = nn.Linear(feature_dim, num_classes)
+        self._init_params()
+
+    @staticmethod
+    def _make_layer(
+        block: type[nn.Module],
+        num_blocks: int,
+        in_channels: int,
+        out_channels: int,
+        downsample: bool,
+    ) -> nn.Sequential:
+        layers: list[nn.Module] = [block(in_channels, out_channels)]
+        for _ in range(1, num_blocks):
+            layers.append(block(out_channels, out_channels))
+        if downsample:
+            layers.append(nn.Sequential(_Conv1x1(out_channels, out_channels), nn.AvgPool2d(2, stride=2)))
+        return nn.Sequential(*layers)
+
+    def _init_params(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _featuremaps(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        return self.conv5(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run a forward pass.
+
+        Returns the embedding at inference time (``model.eval()``) and when no
+        classifier was built. Returns classifier logits only in training mode
+        with ``num_classes > 0``.
+        """
+        x = self._featuremaps(x)
+        v = self.global_avgpool(x).view(x.size(0), -1)
+        v = self.fc(v)
+        if self.training and hasattr(self, "classifier"):
+            return self.classifier(v)
+        return v
+
+
+# --------------------------------------------------------------------------- #
+# Factory
+# --------------------------------------------------------------------------- #
+
+_CONFIGS: dict[str, dict] = {
+    "x1_0": {"channels": [64, 256, 384, 512]},
+    "x0_75": {"channels": [48, 192, 288, 384]},
+    "x0_5": {"channels": [32, 128, 192, 256]},
+    "x0_25": {"channels": [16, 64, 96, 128]},
+}
+
+
+def build_osnet(variant: str = "x1_0", num_classes: int = 0) -> OSNet:
+    """Instantiate an OSNet architecture without loading weights.
+
+    Args:
+        variant: Width multiplier variant. One of ``"x1_0"`` (default),
+            ``"x0_75"``, ``"x0_5"``, ``"x0_25"``.
+        num_classes: Identity classes for the classification head. Use ``0``
+            (default) for inference-only models with no classifier.
+
+    Returns:
+        An initialised ``OSNet`` in training mode.
+
+    Raises:
+        ValueError: If ``variant`` is not one of the supported values.
+
+    Examples:
+        >>> model = build_osnet("x1_0")
+        >>> model.feature_dim
+        512
+    """
+    if variant not in _CONFIGS:
+        raise ValueError(f"Unknown OSNet variant '{variant}'. Choose from: {list(_CONFIGS)}")
+    cfg = _CONFIGS[variant]
+    return OSNet(
+        num_classes=num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=cfg["channels"],
+    )

reid/architectures/registry.py

@@ -0,0 +1,104 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""Architecture construction, listing, remapping, and preprocessing defaults."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import timm
+import torch.nn as nn
+
+from reid.architectures.fastreid_sbs import (
+    FASTREID_SBS_ARCHITECTURE,
+    FastReIDSBSResNeSt50,
+    build_fastreid_sbs_resnest50,
+    remap_fastreid_sbs_state_dict,
+)
+from reid.architectures.osnet import OSNet, build_osnet
+from reid.preprocessing import ReIDPreprocessing
+
+_OSNET_VARIANTS: dict[str, str] = {
+    "osnet_x0_25": "x0_25",
+    "osnet_x0_5": "x0_5",
+    "osnet_x0_75": "x0_75",
+    "osnet_x1_0": "x1_0",
+}
+
+_DEFAULT_PREPROCESSING_BY_ARCH: dict[str, Callable[[], ReIDPreprocessing]] = {
+    **{name: OSNet.default_preprocessing for name in _OSNET_VARIANTS},
+    FASTREID_SBS_ARCHITECTURE: FastReIDSBSResNeSt50.default_preprocessing,
+}
+
+_TIMM_PREFIX = "timm:"
+
+
+def build_architecture(
+    architecture: str | nn.Module,
+    *,
+    num_classes: int = 0,
+    pretrained: bool = False,
+) -> nn.Module:
+    """Build a backbone from ``osnet_*``, ``timm:<name>``, or a pre-built module.
+
+    ``pretrained=True`` only affects ``timm:<name>`` models (ImageNet weights).
+    OSNet ignores it because checkpoints are always loaded explicitly.
+
+    Args:
+        architecture: Registered name, ``timm:<name>``, or an ``nn.Module``.
+        num_classes: Classifier head size (``0`` for embedding-only).
+        pretrained: Load ImageNet weights for ``timm:`` models when ``True``.
+
+    Returns:
+        Backbone module ready for weight loading or inference.
+    """
+    if not isinstance(architecture, str):
+        return architecture
+
+    if architecture.startswith(_TIMM_PREFIX):
+        name = architecture[len(_TIMM_PREFIX) :]
+        return timm.create_model(name, pretrained=pretrained, num_classes=num_classes)
+
+    if architecture in _OSNET_VARIANTS:
+        return build_osnet(variant=_OSNET_VARIANTS[architecture], num_classes=num_classes)
+
+    if architecture == FASTREID_SBS_ARCHITECTURE:
+        return build_fastreid_sbs_resnest50(
+            num_classes=num_classes,
+            pretrained=pretrained,
+        )
+
+    raise ValueError(
+        f"Unknown architecture {architecture!r}. Choose a registered name "
+        f"({list_architectures()}), a timm model as 'timm:<name>' (e.g. "
+        f"'timm:resnet50'), or pass a torch.nn.Module instance."
+    )
+
+
+def list_architectures() -> list[str]:
+    """Return registered architecture names (timm models use ``timm:<name>``)."""
+    return sorted([*_OSNET_VARIANTS, FASTREID_SBS_ARCHITECTURE])
+
+
+def checkpoint_remap_for_architecture(
+    architecture: str,
+) -> Callable[[dict], dict] | None:
+    """Return the registered checkpoint key remapper, if any."""
+    if architecture == FASTREID_SBS_ARCHITECTURE:
+        return remap_fastreid_sbs_state_dict
+    return None
+
+
+def default_preprocessing_for_architecture(
+    architecture: str | None,
+) -> ReIDPreprocessing:
+    """Return the preprocessing defaults owned by ``architecture``."""
+    if architecture is None:
+        return OSNet.default_preprocessing()
+    if architecture in _DEFAULT_PREPROCESSING_BY_ARCH:
+        return _DEFAULT_PREPROCESSING_BY_ARCH[architecture]()
+    return OSNet.default_preprocessing()