trackforge 0.5.0__cp315-cp315-win_amd64.whl

trackforge/__init__.py

@@ -0,0 +1,6 @@
+from .trackforge import *  # noqa: F403
+from .ultralytics import from_ultralytics, track_ultralytics
+
+__doc__ = trackforge.__doc__  # noqa: F405
+if hasattr(trackforge, "__all__"):  # noqa: F405
+    __all__ = [*trackforge.__all__, "from_ultralytics", "track_ultralytics"]  # noqa: F405

trackforge/trackforge.pyi

@@ -0,0 +1,315 @@
+from typing import List, Optional, Tuple
+
+__all__ = [
+    "BOTSORT",
+    "BYTETRACK",
+    "DEEPOCSORT",
+    "DEEPSORT",
+    "OCSORT",
+    "SORT",
+    "TRACKTRACK",
+]
+
+class BYTETRACK:
+    """
+    ByteTrack tracker implementation.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import BYTETRACK
+
+    tracker = BYTETRACK(
+        track_thresh=0.5,
+        track_buffer=30,
+        match_thresh=0.8,
+        det_thresh=0.6,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+        ([200.0, 200.0, 60.0, 120.0], 0.85, 0),
+    ]
+
+    tracks = tracker.update(detections)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        track_thresh: float = 0.5,
+        track_buffer: int = 30,
+        match_thresh: float = 0.8,
+        det_thresh: float = 0.6,
+        second_match_thresh: float = 0.5,
+    ) -> None: ...
+    def update(
+        self, output_results: List[Tuple[List[float], float, int]]
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class SORT:
+    """
+    SORT (Simple Online and Realtime Tracking) tracker implementation.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import SORT
+
+    tracker = SORT(max_age=1, min_hits=3, iou_threshold=0.3)
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+
+    tracks = tracker.update(detections)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        max_age: int = 1,
+        min_hits: int = 3,
+        iou_threshold: float = 0.3,
+    ) -> None: ...
+    def update(
+        self, detections: List[Tuple[List[float], float, int]]
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class OCSORT:
+    """
+    OC-SORT (Observation-Centric SORT) tracker implementation.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import OCSORT
+
+    tracker = OCSORT(
+        max_age=30,
+        min_hits=3,
+        iou_threshold=0.3,
+        delta_t=3,
+        inertia=0.2,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+
+    tracks = tracker.update(detections)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        max_age: int = 30,
+        min_hits: int = 3,
+        iou_threshold: float = 0.3,
+        delta_t: int = 3,
+        inertia: float = 0.2,
+    ) -> None: ...
+    def update(
+        self, detections: List[Tuple[List[float], float, int]]
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class DEEPSORT:
+    """
+    DeepSORT tracker implementation with appearance feature matching.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import DEEPSORT
+    import numpy as np
+
+    tracker = DEEPSORT(
+        max_age=70,
+        n_init=3,
+        max_iou_distance=0.7,
+        max_cosine_distance=0.2,
+        nn_budget=100,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+    embeddings = [np.random.rand(128).tolist()]
+
+    tracks = tracker.update(detections, embeddings)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        max_age: int = 70,
+        n_init: int = 3,
+        max_iou_distance: float = 0.7,
+        max_cosine_distance: float = 0.2,
+        nn_budget: int = 100,
+    ) -> None: ...
+    def update(
+        self,
+        detections: List[Tuple[List[float], float, int]],
+        embeddings: List[List[float]],
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class DEEPOCSORT:
+    """
+    Deep OC-SORT tracker: OC-SORT motion with appearance association.
+
+    Blends a cosine distance to a per-track feature gallery into the OC-SORT
+    motion cost. Pass embeddings for appearance-aware tracking, or omit them to
+    track on motion only.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import DEEPOCSORT
+
+    tracker = DEEPOCSORT(
+        max_age=30,
+        min_hits=3,
+        iou_threshold=0.3,
+        delta_t=3,
+        inertia=0.2,
+        appearance_weight=0.5,
+        max_cosine_distance=0.2,
+        nn_budget=100,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+    embeddings = [[0.1, 0.2, 0.3]]
+
+    tracks = tracker.update(detections, embeddings)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        max_age: int = 30,
+        min_hits: int = 3,
+        iou_threshold: float = 0.3,
+        delta_t: int = 3,
+        inertia: float = 0.2,
+        appearance_weight: float = 0.5,
+        max_cosine_distance: float = 0.2,
+        nn_budget: int = 100,
+    ) -> None: ...
+    def update(
+        self,
+        detections: List[Tuple[List[float], float, int]],
+        embeddings: List[List[float]] = ...,
+        camera_motion: Optional[List[float]] = ...,
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class BOTSORT:
+    """
+    BoT-SORT tracker: ByteTrack with camera motion and appearance fusion.
+
+    Adds camera motion compensation and an appearance-fused first association stage
+    on top of ByteTrack's two-stage cascade. Pass embeddings for appearance-aware
+    tracking, or omit them to track on motion only.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import BOTSORT
+
+    tracker = BOTSORT(
+        track_thresh=0.5,
+        track_buffer=30,
+        match_thresh=0.8,
+        det_thresh=0.6,
+        proximity_thresh=0.5,
+        appearance_thresh=0.25,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+    embeddings = [[0.1, 0.2, 0.3]]
+
+    tracks = tracker.update(detections, embeddings)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        track_thresh: float = 0.5,
+        track_buffer: int = 30,
+        match_thresh: float = 0.8,
+        det_thresh: float = 0.6,
+        second_match_thresh: float = 0.5,
+        proximity_thresh: float = 0.5,
+        appearance_thresh: float = 0.25,
+    ) -> None: ...
+    def update(
+        self,
+        detections: List[Tuple[List[float], float, int]],
+        embeddings: List[List[float]] = ...,
+        camera_motion: Optional[List[float]] = ...,
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...
+
+class TRACKTRACK:
+    """
+    TrackTrack tracker: track-perspective association with track-aware initialization.
+
+    Each track picks its own best detection and a pair matches only when the choice
+    is mutual, repeating with a cost gate that tightens each round. New tracks start
+    only from confident detections that do not overlap existing tracks. Pass
+    embeddings for appearance-aware tracking, or omit them to track on motion only.
+
+    **Usage Example:**
+
+    ```python
+    from trackforge import TRACKTRACK
+
+    tracker = TRACKTRACK(
+        det_thresh=0.6,
+        match_thresh=0.7,
+        track_buffer=30,
+        min_hits=3,
+    )
+
+    detections = [
+        ([100.0, 100.0, 50.0, 100.0], 0.9, 0),
+    ]
+
+    tracks = tracker.update(detections)
+    for track_id, box, score, class_id, det_ind in tracks:
+        print(f"Track ID: {track_id}, Box: {box}")
+    ```
+    """
+
+    def __init__(
+        self,
+        det_thresh: float = 0.6,
+        match_thresh: float = 0.7,
+        track_buffer: int = 30,
+        min_hits: int = 3,
+        init_thresh: float = 0.7,
+        tai_thresh: float = 0.55,
+        penalty_low: float = 0.2,
+        reduce_step: float = 0.05,
+    ) -> None: ...
+    def update(
+        self,
+        detections: List[Tuple[List[float], float, int]],
+        embeddings: List[List[float]] = ...,
+        camera_motion: Optional[List[float]] = ...,
+    ) -> List[Tuple[int, List[float], float, int, Optional[int]]]: ...

trackforge/ultralytics.py

@@ -0,0 +1,83 @@
+"""Helpers for feeding Ultralytics YOLO results into trackforge trackers.
+
+They only read the result objects, so ``ultralytics`` stays an optional dependency
+and is never imported here.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple
+
+Detection = Tuple[List[float], float, int]
+Embedder = Callable[[Any, List[List[float]]], Sequence[Sequence[float]]]
+
+__all__ = ["from_ultralytics", "track_ultralytics"]
+
+
+def from_ultralytics(result: Any) -> List[Detection]:
+    """Convert one Ultralytics result into trackforge detections.
+
+    Args:
+        result: A ``Results`` object from ``model.predict(...)`` or ``model(...)``, or
+            the one-element list those calls return for a single image.
+
+    Returns:
+        A list of ``(tlwh, score, class_id)`` detections, in the same order as
+        ``result.boxes``, so a track's ``det_ind`` indexes into ``result.boxes``.
+
+    Raises:
+        ValueError: If a list holding more than one result is passed.
+    """
+    if isinstance(result, (list, tuple)):
+        if len(result) != 1:
+            raise ValueError(
+                f"expected the result of a single image, got {len(result)} results"
+            )
+        result = result[0]
+
+    boxes = getattr(result, "boxes", None)
+    if boxes is None or len(boxes) == 0:
+        return []
+
+    return [
+        ([x1, y1, x2 - x1, y2 - y1], float(score), int(cls))
+        for (x1, y1, x2, y2), score, cls in zip(
+            boxes.xyxy.tolist(), boxes.conf.tolist(), boxes.cls.tolist()
+        )
+    ]
+
+
+def track_ultralytics(
+    model: Any,
+    source: Any,
+    tracker: Any,
+    embedder: Optional[Embedder] = None,
+    **predict_kwargs: Any,
+) -> Iterator[Tuple[Any, list]]:
+    """Run an Ultralytics model over a source and track every frame.
+
+    Args:
+        model: A loaded ``ultralytics.YOLO`` (or RT-DETR) model.
+        source: Anything ``model.predict`` accepts: a video path, webcam index,
+            stream URL, image folder, and so on.
+        tracker: A trackforge tracker such as ``trackforge.BYTETRACK()``.
+        embedder: Optional ``embedder(frame, tlwh_boxes) -> embeddings`` callable for
+            appearance trackers. ``DEEPSORT`` needs one; ``DEEPOCSORT``, ``BOTSORT``
+            and ``TRACKTRACK`` use it when given.
+        **predict_kwargs: Passed to ``model.predict``, for example ``classes=[0]``
+            or ``conf=0.3``.
+
+    Yields:
+        ``(result, tracks)`` per frame, where ``tracks`` is the tracker output and
+        ``result.orig_img`` is the BGR frame.
+    """
+    predict_kwargs.setdefault("stream", True)
+    predict_kwargs.setdefault("verbose", False)
+    for result in model.predict(source, **predict_kwargs):
+        detections = from_ultralytics(result)
+        if embedder is None:
+            tracks = tracker.update(detections)
+        else:
+            boxes = [d[0] for d in detections]
+            tracks = tracker.update(detections, embedder(result.orig_img, boxes))
+        yield result, tracks