reid 0.1.0.dev0__py3-none-any.whl

reid/catalog.py

@@ -0,0 +1,222 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""Curated ReID aliases, ``ModelCard``, and ``reid_config.json`` I/O."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError
+
+from reid.architectures.fastreid_sbs import FASTREID_SBS_ARCHITECTURE
+from reid.preprocessing import ReIDPreprocessing
+
+# OSNet x1.0 trained on MSMT17 with combineall (train + query + gallery).
+# Strong general-purpose pedestrian features, but combineall includes MSMT17
+# test IDs — not suitable for MSMT17 evaluation.
+_DEFAULT_OSNET_WEIGHTS = (
+    "hf://kaiyangzhou/osnet/"
+    "osnet_x1_0_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10"
+    "_softmax_labelsmooth_flip_jitter.pth"
+)
+
+_DOMAIN_WARNING = (
+    "The default ReIDModel weights were trained on MSMT17 pedestrian images. "
+    "Performance may degrade significantly on non-person objects (vehicles, "
+    "animals, products, etc.). Pass a domain-specific checkpoint via "
+    "`ReIDModel.from_pretrained(source, architecture=...)` for other use cases."
+)
+
+# ---------------------------------------------------------------------------
+# Public constants and types
+# ---------------------------------------------------------------------------
+
+DEFAULT_MODEL = "osnet_x1_0_msmt17_combineall"
+"""Curated alias used when ``ReIDModel.from_pretrained`` gets no ``source``."""
+
+# BoT-SORT MOT17 SBS-S50 (ResNeSt50 + GeM + BNNeck), trained on MOT17 train-half GT crops.
+# Weights: https://github.com/niraharon/bot-sort#model-zoo
+_FASTREID_MOT17_WEIGHTS = "gd://1QZFWpoa80rqo7O-HXmlss8J8CnS7IUsN/mot17_sbs_S50.pth"
+_FASTREID_MOT17_WARNING = (
+    "The fastreid_mot17_sbs50 weights were trained on MOT17 pedestrian crops. "
+    "Use this encoder for MOT tracking benchmarks (BoT-SORT-ReID replication); "
+    "cross-domain retrieval on Market-1501 / MSMT17 is expected to underperform."
+)
+
+FASTREID_MOT17_SBS50 = "fastreid_mot17_sbs50"
+
+
+@dataclass
+class ModelCard:
+    """Architecture, weights source, and optional domain metadata.
+
+    Produced by curated ``ALIASES`` entries or by loading ``reid_config.json``
+    from a ``save_pretrained`` directory / HF repo. Preprocessing for curated
+    aliases is derived from the architecture at load time; saved configs may
+    carry a preprocessing override in ``reid_config.json`` (see
+    :func:`resolve_model_card`).
+    """
+
+    architecture: str
+    weights: str | None
+    domain_warning: str | None = None
+
+
+ALIASES: dict[str, ModelCard] = {
+    DEFAULT_MODEL: ModelCard(
+        architecture="osnet_x1_0",
+        weights=_DEFAULT_OSNET_WEIGHTS,
+        domain_warning=_DOMAIN_WARNING,
+    ),
+    FASTREID_MOT17_SBS50: ModelCard(
+        architecture=FASTREID_SBS_ARCHITECTURE,
+        weights=_FASTREID_MOT17_WEIGHTS,
+        domain_warning=_FASTREID_MOT17_WARNING,
+    ),
+}
+
+# ---------------------------------------------------------------------------
+# Resolution helpers
+# ---------------------------------------------------------------------------
+
+
+def list_aliases() -> list[str]:
+    """Return curated alias names (sorted keys of :data:`ALIASES`)."""
+    return sorted(ALIASES)
+
+
+def resolve_model_card(source: str) -> tuple[ModelCard, ReIDPreprocessing | None] | None:
+    """Resolve an alias or ``reid_config.json`` directory/repo to a load recipe.
+
+    Returns ``None`` when ``source`` is a bare weights *file* URL/path without a
+    config. Curated aliases return ``(card, None)``; saved / HF configs return
+    ``(card, preprocessing)`` with the preprocessing dict from
+    ``reid_config.json``.
+
+    Raises:
+        ValueError: If ``source`` is a local directory that does not contain
+            ``reid_config.json`` (directories are never bare weight files).
+    """
+    if source in ALIASES:
+        return ALIASES[source], None
+
+    # Local directory must be a save_pretrained() tree (reid_config.json).
+    if os.path.isdir(source):
+        config_path = os.path.join(source, "reid_config.json")
+        if not os.path.isfile(config_path):
+            raise ValueError(
+                f"Cannot load {source!r}: directory has no reid_config.json. "
+                "Point to a directory produced by ReIDModel.save_pretrained(), "
+                "or pass a weights *file* path together with architecture=... ."
+            )
+        return load_model_config(source)
+
+    # HF repo URL (no trailing filename) — try to fetch reid_config.json.
+    # hf://org/repo has exactly 2 non-empty parts after "hf://".
+    # hf://org/repo/file.pth has 3+ parts and is a bare weights file.
+    if source.startswith("hf://"):
+        rest = source[len("hf://") :]
+        parts = [p for p in rest.split("/") if p]
+        if len(parts) == 2:
+            try:
+                return _load_hf_repo_config(source)
+            except (HfHubHTTPError, EntryNotFoundError, OSError, ValueError) as exc:
+                raise RuntimeError(f"Failed to resolve Hugging Face repo config for {source!r}.") from exc
+
+    return None
+
+
+def load_model_config(directory_or_repo: str) -> tuple[ModelCard, ReIDPreprocessing]:
+    """Load a ``ModelCard`` and preprocessing from ``reid_config.json``.
+
+    Local directories may also include ``weights.safetensors``.
+    """
+    if directory_or_repo.startswith("hf://"):
+        card, preprocessing = _load_hf_repo_config(directory_or_repo)
+        if preprocessing is None:
+            raise ValueError(f"No preprocessing found in config for {directory_or_repo!r}.")
+        return card, preprocessing
+
+    config_path = os.path.join(directory_or_repo, "reid_config.json")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(
+            f"No reid_config.json found in {directory_or_repo!r}. "
+            "Use ReIDModel.save_pretrained() to create a self-describing "
+            "model directory."
+        )
+    weights_path = os.path.join(directory_or_repo, "weights.safetensors")
+    weights_source = weights_path if os.path.exists(weights_path) else None
+    return _parse_config_file(config_path, weights_source=weights_source)
+
+
+def save_model_config(
+    card: ModelCard,
+    directory: str,
+    *,
+    preprocessing: ReIDPreprocessing,
+) -> None:
+    """Write ``reid_config.json`` to ``directory``."""
+    if not isinstance(card.architecture, str):
+        raise ValueError(
+            "Cannot save a ModelCard whose architecture is an nn.Module. "
+            "Use a named architecture string (e.g. 'osnet_x1_0', "
+            "'timm:resnet50') to enable save_pretrained()."
+        )
+
+    config = {
+        "architecture": card.architecture,
+        "preprocessing": preprocessing.to_dict(),
+    }
+    config_path = os.path.join(directory, "reid_config.json")
+    with open(config_path, "w") as f:
+        json.dump(config, f, indent=2)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_hf_repo_config(hf_repo_url: str) -> tuple[ModelCard, ReIDPreprocessing | None]:
+    """Download and parse ``reid_config.json`` from an HF repo URL."""
+    rest = hf_repo_url[len("hf://") :]
+    parts = [p for p in rest.split("/") if p]
+    repo_id = "/".join(parts[:2])
+    revision = None
+    if "@" in repo_id:
+        repo_id, revision = repo_id.split("@", 1)
+
+    config_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="reid_config.json",
+        revision=revision,
+    )
+    weights_source = hf_repo_url.rstrip("/") + "/weights.safetensors"
+    return _parse_config_file(config_path, weights_source=weights_source)
+
+
+def _parse_config_file(
+    config_path: str,
+    *,
+    weights_source: str | None,
+) -> tuple[ModelCard, ReIDPreprocessing]:
+    """Parse a ``reid_config.json`` file into a ``ModelCard`` and preprocessing."""
+    with open(config_path) as f:
+        data = json.load(f)
+
+    architecture = data["architecture"]
+    preprocessing = ReIDPreprocessing.from_dict(data.get("preprocessing", {}))
+    return (
+        ModelCard(
+            architecture=architecture,
+            weights=weights_source,
+        ),
+        preprocessing,
+    )

reid/data/__init__.py

@@ -0,0 +1,27 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""Shared ReID dataset types, benchmark loaders, and crop partitioners."""
+
+from __future__ import annotations
+
+from reid.data.benchmarks import load_market1501, load_msmt17
+from reid.data.identity_crops import (
+    build_identity_crop_holdout_split,
+    build_identity_crop_retrieval_split,
+    load_identity_crops,
+)
+from reid.data.splits import MARKET1501_GALLERY_JUNK_PIDS, ReIDSplit
+
+__all__ = [
+    "MARKET1501_GALLERY_JUNK_PIDS",
+    "ReIDSplit",
+    "build_identity_crop_holdout_split",
+    "build_identity_crop_retrieval_split",
+    "load_identity_crops",
+    "load_market1501",
+    "load_msmt17",
+]

reid/data/benchmarks.py

@@ -0,0 +1,163 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""ReID benchmark dataset loaders (Market-1501, MSMT17)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from reid.data.splits import MARKET1501_GALLERY_JUNK_PIDS, ReIDSplit
+
+# --------------------------------------------------------------------------- #
+# MSMT17
+# --------------------------------------------------------------------------- #
+
+
+def _to_zero_indexed_camid(camid_1indexed: int) -> int:
+    """Convert a 1-indexed MSMT17 camera id to the 0-indexed value used here.
+
+    MSMT17 camera ids are 1-indexed both in image filenames (the camera field)
+    and in the optional camid column of the list files, so both paths share this
+    single conversion to stay consistent.
+    """
+    return camid_1indexed - 1
+
+
+def _parse_msmt17_camid(filename: str) -> int:
+    """Extract 0-indexed camid from an MSMT17 image filename.
+
+    Filename format: ``<pid>_<idx>_<camid>_<scene>_<frame>_<track>.jpg``
+    Example: ``0001_019_07_0303morning_0020_1.jpg`` → camid = 6 (0-indexed).
+    """
+    stem = Path(filename).stem
+    return _to_zero_indexed_camid(int(stem.split("_")[2]))
+
+
+def _load_msmt17_list(list_file: Path, image_root: Path) -> ReIDSplit:
+    """Load one MSMT17 query/gallery list file."""
+    if not list_file.exists():
+        raise FileNotFoundError(f"MSMT17 list file not found: {list_file}")
+    paths, pids, camids = [], [], []
+    for line in list_file.read_text().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        parts = line.split()
+        if len(parts) < 2:
+            continue
+        rel_path = parts[0]
+        pid = int(parts[1])
+        camid = _to_zero_indexed_camid(int(parts[2])) if len(parts) >= 3 else _parse_msmt17_camid(Path(rel_path).name)
+        paths.append(str(image_root / rel_path))
+        pids.append(pid)
+        camids.append(camid)
+    return ReIDSplit(
+        image_paths=paths,
+        pids=np.array(pids, dtype=np.int32),
+        camids=np.array(camids, dtype=np.int32),
+        gallery_junk_pids=frozenset({-1}),
+    )
+
+
+def load_msmt17(root: str | Path) -> tuple[ReIDSplit, ReIDSplit]:
+    """Load MSMT17 query and gallery splits from ``root``.
+
+    Expects ``test/``, ``list_query.txt``, and ``list_gallery.txt``. List files
+    may be 2-column (``path pid``; camid parsed from the filename) or 3-column
+    (``path pid camid``). In both cases the camera id is treated as 1-indexed
+    and converted to 0-indexed, so the two formats stay consistent.
+
+    Args:
+        root: Path to the MSMT17 directory.
+
+    Returns:
+        Query and gallery as ``(ReIDSplit, ReIDSplit)``.
+
+    Raises:
+        FileNotFoundError: If ``root``, ``list_query.txt``, or ``list_gallery.txt``
+            does not exist.
+    """
+    root = Path(root)
+    if not root.exists():
+        raise FileNotFoundError(f"MSMT17 root not found: {root}")
+
+    test_root = root / "test"
+    query = _load_msmt17_list(root / "list_query.txt", test_root)
+    gallery = _load_msmt17_list(root / "list_gallery.txt", test_root)
+    return query, gallery
+
+
+# --------------------------------------------------------------------------- #
+# Market-1501
+# --------------------------------------------------------------------------- #
+
+
+def _parse_market_filename(filename: str) -> tuple[int, int]:
+    """Parse ``(pid, camid)`` from a Market-1501 filename.
+
+    ``-1_*`` → junk (``pid=-1``); ``0000_*`` → distractor (``pid=0``).
+    """
+    stem = Path(filename).stem
+    parts = stem.split("_")
+    pid = int(parts[0])
+    camid = int(parts[1][1]) - 1  # "c1" → 0
+    return pid, camid
+
+
+def _load_market1501_dir(
+    subdir: Path,
+    *,
+    gallery_junk_pids: frozenset[int],
+) -> ReIDSplit:
+    """Load one Market-1501 image directory."""
+    if not subdir.exists():
+        raise FileNotFoundError(f"Market-1501 sub-directory not found: {subdir}")
+    paths, pids, camids = [], [], []
+    for img_path in sorted(subdir.glob("*.jpg")):
+        pid, camid = _parse_market_filename(img_path.name)
+        paths.append(str(img_path))
+        pids.append(pid)
+        camids.append(camid)
+    return ReIDSplit(
+        image_paths=paths,
+        pids=np.array(pids, dtype=np.int32),
+        camids=np.array(camids, dtype=np.int32),
+        gallery_junk_pids=gallery_junk_pids,
+    )
+
+
+def load_market1501(root: str | Path) -> tuple[ReIDSplit, ReIDSplit]:
+    """Load Market-1501 query and gallery splits from ``root``.
+
+    Expects ``query/`` and ``bounding_box_test/``. Person and camera IDs are
+    parsed from filenames.
+
+    Args:
+        root: Path to the Market-1501 directory.
+
+    Returns:
+        Query and gallery as ``(ReIDSplit, ReIDSplit)``.
+
+    Raises:
+        FileNotFoundError: If ``root``, ``query/``, or ``bounding_box_test/``
+            does not exist.
+    """
+    root = Path(root)
+    if not root.exists():
+        raise FileNotFoundError(f"Market-1501 root not found: {root}")
+
+    query = _load_market1501_dir(
+        root / "query",
+        gallery_junk_pids=frozenset({-1}),
+    )
+    gallery = _load_market1501_dir(
+        root / "bounding_box_test",
+        gallery_junk_pids=MARKET1501_GALLERY_JUNK_PIDS,
+    )
+    return query, gallery

reid/data/identity_crops.py

@@ -0,0 +1,271 @@
+# ------------------------------------------------------------------------
+# re-ID
+# Copyright (c) 2026 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+"""Load and partition per-identity crop folders into ``ReIDSplit`` objects.
+
+The canonical flow is ``source format -> loader -> ReIDSplit -> optional
+partitioning -> Dataset/evaluator``. :func:`load_identity_crops` turns an
+identity-folder layout (``<root>/<identity>/<crop>.jpg``) into a single
+``ReIDSplit``. The ``build_identity_crop_*`` partitioners then consume that split to
+produce query/gallery (and optional train) splits; they never touch the
+filesystem.
+
+Query and gallery crops receive synthetic ``camid`` 0/1 so single-camera MOT
+crops still form valid cross-camera retrieval pairs under the gallery-eval junk
+rule.
+"""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+
+import numpy as np
+
+from reid.data.splits import ReIDSplit
+
+# Image suffixes scanned inside each identity folder. Kept private to this
+# module; there is no shared ``constants`` module by design.
+_IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp")
+
+_QUERY_CAMID = 0
+_GALLERY_CAMID = 1
+
+
+def load_identity_crops(
+    root: str | Path,
+    *,
+    identities: list[str] | None = None,
+) -> ReIDSplit:
+    """Scan an identity-folder crop dataset into a single :class:`ReIDSplit`.
+
+    Expects ``<root>/<identity>/<crop>.(jpg|jpeg|png|bmp)``. Identity folders are
+    sorted by name and assigned contiguous person IDs (``0..N-1``). Camera IDs
+    are set to ``-1`` because on-disk crops carry no camera information; the
+    ``build_identity_crop_*`` partitioners assign synthetic camera IDs where retrieval
+    needs them.
+
+    Args:
+        root: Dataset root with one subdirectory per identity.
+        identities: Optional whitelist of identity folder names. Folders outside
+            the whitelist are skipped before person IDs are assigned.
+
+    Returns:
+        A :class:`ReIDSplit` over every crop found, with contiguous ``pids`` and
+        ``camids`` all equal to ``-1``.
+
+    Raises:
+        FileNotFoundError: If ``root`` is not a directory, or if no crops are
+            found (after applying ``identities``).
+    """
+    root = Path(root)
+    if not root.is_dir():
+        raise FileNotFoundError(f"Crop dataset root not found: {root}")
+
+    allowed = set(identities) if identities is not None else None
+
+    image_paths: list[str] = []
+    pids: list[int] = []
+    next_pid = 0
+    for identity_dir in sorted(p for p in root.iterdir() if p.is_dir()):
+        if allowed is not None and identity_dir.name not in allowed:
+            continue
+        images = sorted(p for p in identity_dir.iterdir() if p.suffix.lower() in _IMAGE_EXTENSIONS)
+        if not images:
+            continue
+        for image_path in images:
+            image_paths.append(str(image_path))
+            pids.append(next_pid)
+        next_pid += 1
+
+    if not image_paths:
+        raise FileNotFoundError(f"No identity crops found under {root}. Expected <root>/<identity>/<crop>.jpg.")
+
+    return ReIDSplit(
+        image_paths=image_paths,
+        pids=np.array(pids, dtype=np.int32),
+        camids=np.full(len(image_paths), -1, dtype=np.int32),
+    )
+
+
+def build_identity_crop_retrieval_split(
+    split: ReIDSplit,
+    *,
+    queries_per_id: int = 1,
+    min_crops: int = 2,
+    seed: int = 0,
+) -> tuple[ReIDSplit, ReIDSplit]:
+    """Partition a crop split into query and gallery splits for retrieval eval.
+
+    Every identity with at least ``min_crops`` crops contributes ``queries_per_id``
+    crops to the query split (synthetic ``camid`` 0) and the remainder to the
+    gallery split (synthetic ``camid`` 1). Crops are shuffled per identity with
+    ``seed`` before the query/gallery cut so query selection is deterministic but
+    not tied to filename order.
+
+    Args:
+        split: A crop split, typically from :func:`load_identity_crops`.
+        queries_per_id: Query crops taken from each eligible identity.
+        min_crops: Minimum crops required for an identity to be eligible.
+        seed: Shuffle seed for per-identity query selection.
+
+    Returns:
+        ``(query, gallery)`` splits.
+
+    Raises:
+        ValueError: If ``min_crops``/``queries_per_id`` are incompatible (see
+            :func:`_check_query_params`), or if no identity is eligible.
+    """
+    _check_query_params(queries_per_id=queries_per_id, min_crops=min_crops)
+
+    indices_by_pid = _indices_by_pid(split)
+    eligible_pids = sorted(pid for pid, indices in indices_by_pid.items() if len(indices) >= min_crops)
+    if not eligible_pids:
+        raise ValueError(f"No identity has >= {min_crops} crops; cannot build a retrieval split.")
+
+    return _query_gallery_from_pids(
+        split,
+        indices_by_pid,
+        eligible_pids,
+        queries_per_id=queries_per_id,
+        seed=seed,
+    )
+
+
+def build_identity_crop_holdout_split(
+    split: ReIDSplit,
+    *,
+    holdout_fraction: float = 0.2,
+    queries_per_id: int = 1,
+    min_crops: int = 2,
+    seed: int = 0,
+) -> tuple[ReIDSplit, ReIDSplit, ReIDSplit]:
+    """Hold out whole identities for retrieval; return train, query, gallery.
+
+    A fraction of the identities that have at least ``min_crops`` crops is drawn
+    (deterministically via ``seed``) and turned into query/gallery splits exactly
+    as in :func:`build_identity_crop_retrieval_split`. Every remaining crop, including
+    crops of identities too small to be eligible for holdout, forms the train
+    split with its original camera IDs preserved. The train and evaluation splits
+    are identity-disjoint, so retrieval never scores an identity the model saw
+    during training.
+
+    Args:
+        split: A crop split, typically from :func:`load_identity_crops`.
+        holdout_fraction: Fraction of eligible identities reserved for retrieval.
+        queries_per_id: Query crops per held-out identity.
+        min_crops: Minimum crops required for an identity to be eligible for
+            holdout.
+        seed: Shuffle seed for the holdout draw and per-identity query selection.
+
+    Returns:
+        ``(train, query, gallery)`` splits.
+
+    Raises:
+        ValueError: If ``holdout_fraction`` is outside ``(0, 1]``, if
+            ``min_crops``/``queries_per_id`` are incompatible, or if no identity
+            is eligible for holdout.
+    """
+    _check_query_params(queries_per_id=queries_per_id, min_crops=min_crops)
+    if not 0.0 < holdout_fraction <= 1.0:
+        raise ValueError(f"holdout_fraction must be in (0, 1], got {holdout_fraction}.")
+
+    indices_by_pid = _indices_by_pid(split)
+    eligible_pids = sorted(pid for pid, indices in indices_by_pid.items() if len(indices) >= min_crops)
+    if not eligible_pids:
+        raise ValueError(f"No identity has >= {min_crops} crops; cannot build a retrieval holdout.")
+
+    shuffled = eligible_pids.copy()
+    random.Random(seed).shuffle(shuffled)
+    num_holdout = max(1, round(len(shuffled) * holdout_fraction))
+    holdout_pids = set(shuffled[:num_holdout])
+
+    train_indices = [index for index, pid in enumerate(split.pids.tolist()) if pid not in holdout_pids]
+    train = _subset_split(split, train_indices)
+    query, gallery = _query_gallery_from_pids(
+        split,
+        indices_by_pid,
+        sorted(holdout_pids),
+        queries_per_id=queries_per_id,
+        seed=seed,
+    )
+    return train, query, gallery
+
+
+def _check_query_params(*, queries_per_id: int, min_crops: int) -> None:
+    """Reject incompatible ``queries_per_id`` / ``min_crops`` combinations.
+
+    An eligible identity must yield at least one query and one gallery crop, so
+    ``min_crops`` has to exceed ``queries_per_id``. The combination is rejected
+    explicitly rather than silently widened.
+    """
+    if queries_per_id < 1:
+        raise ValueError(f"queries_per_id must be >= 1, got {queries_per_id}.")
+    if min_crops <= queries_per_id:
+        raise ValueError(
+            f"min_crops must be > queries_per_id to leave gallery crops, "
+            f"got min_crops={min_crops}, queries_per_id={queries_per_id}."
+        )
+
+
+def _indices_by_pid(split: ReIDSplit) -> dict[int, list[int]]:
+    """Group crop indices by person ID, preserving the split's order."""
+    grouped: dict[int, list[int]] = {}
+    for index, pid in enumerate(split.pids.tolist()):
+        grouped.setdefault(pid, []).append(index)
+    return grouped
+
+
+def _query_gallery_from_pids(
+    split: ReIDSplit,
+    indices_by_pid: dict[int, list[int]],
+    pids: list[int],
+    *,
+    queries_per_id: int,
+    seed: int,
+) -> tuple[ReIDSplit, ReIDSplit]:
+    """Cut each identity in ``pids`` into query and gallery crops."""
+    rng = random.Random(seed)
+    query_indices: list[int] = []
+    gallery_indices: list[int] = []
+    for pid in pids:
+        indices = indices_by_pid[pid].copy()
+        rng.shuffle(indices)
+        query_indices.extend(indices[:queries_per_id])
+        gallery_indices.extend(indices[queries_per_id:])
+
+    query = _subset_split(split, query_indices, camid=_QUERY_CAMID)
+    gallery = _subset_split(split, gallery_indices, camid=_GALLERY_CAMID)
+    return query, gallery
+
+
+def _subset_split(
+    split: ReIDSplit,
+    indices: list[int],
+    *,
+    camid: int | None = None,
+) -> ReIDSplit:
+    """Select ``indices`` from ``split`` into a new :class:`ReIDSplit`.
+
+    When ``camid`` is given, every crop in the subset is stamped with that
+    synthetic camera ID and the default junk rule; otherwise the original camera
+    IDs and junk rule are preserved.
+    """
+    image_paths = [split.image_paths[index] for index in indices]
+    pids = np.array([int(split.pids[index]) for index in indices], dtype=np.int32)
+    if camid is None:
+        camids = np.array([int(split.camids[index]) for index in indices], dtype=np.int32)
+        return ReIDSplit(
+            image_paths=image_paths,
+            pids=pids,
+            camids=camids,
+            gallery_junk_pids=split.gallery_junk_pids,
+        )
+    return ReIDSplit(
+        image_paths=image_paths,
+        pids=pids,
+        camids=np.full(len(indices), camid, dtype=np.int32),
+    )