From aa40de4ecf2df68ea28d306951bf022b342d9d64 Mon Sep 17 00:00:00 2001 From: Tomasz Stanczyk Date: Thu, 14 May 2026 22:46:18 +0200 Subject: [PATCH 1/6] Add McByte mask manager --- src/trackers/core/mcbyte/mask_manager.py | 72 +++++++++ src/trackers/core/mcbyte/masks/__init__.py | 25 +++ src/trackers/core/mcbyte/masks/base.py | 64 ++++++++ src/trackers/core/mcbyte/masks/dummy.py | 88 +++++++++++ src/trackers/core/mcbyte/tracker.py | 62 +++++++- tests/core/test_mcbyte_mask_manager.py | 168 +++++++++++++++++++++ 6 files changed, 477 insertions(+), 2 deletions(-) create mode 100644 src/trackers/core/mcbyte/mask_manager.py create mode 100644 src/trackers/core/mcbyte/masks/__init__.py create mode 100644 src/trackers/core/mcbyte/masks/base.py create mode 100644 src/trackers/core/mcbyte/masks/dummy.py create mode 100644 tests/core/test_mcbyte_mask_manager.py diff --git a/src/trackers/core/mcbyte/mask_manager.py b/src/trackers/core/mcbyte/mask_manager.py new file mode 100644 index 00000000..20f214b7 --- /dev/null +++ b/src/trackers/core/mcbyte/mask_manager.py @@ -0,0 +1,72 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +from __future__ import annotations + +import numpy as np + +from trackers.core.mcbyte.masks.base import ( + MaskGenerator, + MaskOutput, + MaskPropagator, + TrackletSnapshot, +) + + +class MaskManager: + """Manage McByte mask generation and propagation. + + The manager follows the original McByte timing: masks for the current frame + are prepared before association, but they are initialized/updated from tracker + outputs produced on the previous frame. + """ + + def __init__( + self, + mask_generator: MaskGenerator, + mask_propagator: MaskPropagator | None = None, + ) -> None: + self.mask_generator = mask_generator + self.mask_propagator = mask_propagator + self._initialized = False + + def reset(self) -> None: + self._initialized = False + if self.mask_propagator is not None: + self.mask_propagator.reset() + + def get_updated_masks( + self, + frame: np.ndarray, + previous_frame: np.ndarray | None, + previous_tracklets: list[TrackletSnapshot], + ) -> MaskOutput | None: + """Return masks for the current frame. + + No masks are returned until at least one previous frame and previous + tracker output are available. + """ + if previous_frame is None or len(previous_tracklets) == 0: + return None + + if not self._initialized or self.mask_propagator is None: + mask_output = self.mask_generator.generate(previous_frame, previous_tracklets) + self._initialized = True + + if self.mask_propagator is not None: + self.mask_propagator.initialize(previous_frame, mask_output) + propagated_output = self.mask_propagator.propagate(frame) + return propagated_output if propagated_output is not None else mask_output + + return mask_output + + propagated_output = self.mask_propagator.propagate(frame) + if propagated_output is not None: + return propagated_output + + mask_output = self.mask_generator.generate(previous_frame, previous_tracklets) + self.mask_propagator.initialize(previous_frame, mask_output) + return mask_output \ No newline at end of file diff --git a/src/trackers/core/mcbyte/masks/__init__.py b/src/trackers/core/mcbyte/masks/__init__.py new file mode 100644 index 00000000..bb75a957 --- /dev/null +++ b/src/trackers/core/mcbyte/masks/__init__.py @@ -0,0 +1,25 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +from trackers.core.mcbyte.masks.base import ( + MaskGenerator, + MaskOutput, + MaskPropagator, + TrackletSnapshot, +) +from trackers.core.mcbyte.masks.dummy import ( + DummyBoxMaskGenerator, + DummyIdentityMaskPropagator, +) + +__all__ = [ + "DummyBoxMaskGenerator", + "DummyIdentityMaskPropagator", + "MaskGenerator", + "MaskOutput", + "MaskPropagator", + "TrackletSnapshot", +] \ No newline at end of file diff --git a/src/trackers/core/mcbyte/masks/base.py b/src/trackers/core/mcbyte/masks/base.py new file mode 100644 index 00000000..f7894446 --- /dev/null +++ b/src/trackers/core/mcbyte/masks/base.py @@ -0,0 +1,64 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class TrackletSnapshot: + """Minimal tracker state needed by mask components.""" + + tracker_id: int + xyxy: np.ndarray + + +@dataclass(frozen=True) +class MaskOutput: + """Mask information produced before McByte association.""" + + masks: np.ndarray | None + tracklet_mask_dict: dict[int, int] + mask_avg_prob_dict: dict[int, float] | None = None + + +class MaskGenerator(ABC): + """Generate masks from tracklet boxes.""" + + @abstractmethod + def generate( + self, + frame: np.ndarray, + tracklets: list[TrackletSnapshot], + ) -> MaskOutput: + """Generate masks for the given tracklet snapshots.""" + + +class MaskPropagator(ABC): + """Propagate masks from one frame to the next.""" + + @abstractmethod + def reset(self) -> None: + """Reset propagation state.""" + + @abstractmethod + def initialize( + self, + frame: np.ndarray, + mask_output: MaskOutput, + ) -> None: + """Initialize propagation state.""" + + @abstractmethod + def propagate( + self, + frame: np.ndarray, + ) -> MaskOutput | None: + """Propagate masks to the current frame.""" \ No newline at end of file diff --git a/src/trackers/core/mcbyte/masks/dummy.py b/src/trackers/core/mcbyte/masks/dummy.py new file mode 100644 index 00000000..b622feda --- /dev/null +++ b/src/trackers/core/mcbyte/masks/dummy.py @@ -0,0 +1,88 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +from __future__ import annotations + +import numpy as np + +from trackers.core.mcbyte.masks.base import ( + MaskGenerator, + MaskOutput, + MaskPropagator, + TrackletSnapshot, +) + + +class DummyBoxMaskGenerator(MaskGenerator): + """Generate rectangular binary masks from tracklet bounding boxes.""" + + def generate( + self, + frame: np.ndarray, + tracklets: list[TrackletSnapshot], + ) -> MaskOutput: + height, width = frame.shape[:2] + masks = np.zeros((len(tracklets), height, width), dtype=bool) + tracklet_mask_dict: dict[int, int] = {} + + for mask_index, tracklet in enumerate(tracklets): + x1, y1, x2, y2 = tracklet.xyxy.astype(int) + + x1 = int(np.clip(x1, 0, width)) + x2 = int(np.clip(x2, 0, width)) + y1 = int(np.clip(y1, 0, height)) + y2 = int(np.clip(y2, 0, height)) + + masks[mask_index, y1:y2, x1:x2] = True + tracklet_mask_dict[tracklet.tracker_id] = mask_index + + return MaskOutput( + masks=masks, + tracklet_mask_dict=tracklet_mask_dict, + mask_avg_prob_dict=None, + ) + + +class DummyIdentityMaskPropagator(MaskPropagator): + """Return the last initialized mask output unchanged.""" + + def __init__(self) -> None: + self._mask_output: MaskOutput | None = None + + def reset(self) -> None: + self._mask_output = None + + def initialize( + self, + frame: np.ndarray, + mask_output: MaskOutput, + ) -> None: + self._mask_output = MaskOutput( + masks=None if mask_output.masks is None else mask_output.masks.copy(), + tracklet_mask_dict=mask_output.tracklet_mask_dict.copy(), + mask_avg_prob_dict=( + None + if mask_output.mask_avg_prob_dict is None + else mask_output.mask_avg_prob_dict.copy() + ), + ) + + def propagate( + self, + frame: np.ndarray, + ) -> MaskOutput | None: + if self._mask_output is None: + return None + + return MaskOutput( + masks=None if self._mask_output.masks is None else self._mask_output.masks.copy(), + tracklet_mask_dict=self._mask_output.tracklet_mask_dict.copy(), + mask_avg_prob_dict=( + None + if self._mask_output.mask_avg_prob_dict is None + else self._mask_output.mask_avg_prob_dict.copy() + ), + ) \ No newline at end of file diff --git a/src/trackers/core/mcbyte/tracker.py b/src/trackers/core/mcbyte/tracker.py index 9edecd9e..f12c1b2c 100644 --- a/src/trackers/core/mcbyte/tracker.py +++ b/src/trackers/core/mcbyte/tracker.py @@ -21,6 +21,12 @@ BaseStateEstimator, XCYCWHStateEstimator, ) +from trackers.core.mcbyte.mask_manager import MaskManager +from trackers.core.mcbyte.masks.base import MaskOutput, TrackletSnapshot +from trackers.core.mcbyte.masks.dummy import ( + DummyBoxMaskGenerator, + DummyIdentityMaskPropagator, +) class McByteTracker(BaseTracker): @@ -42,6 +48,8 @@ def __init__( instant_first_frame_activation: bool = True, state_estimator_class: type[BaseStateEstimator] = XCYCWHStateEstimator, iou: BaseIoU | None = None, + enable_mask_manager: bool = False, + mask_manager: MaskManager | None = None, ) -> None: # Calculate maximum frames without update based on lost_track_buffer and # frame_rate. This scales the buffer based on the frame rate to ensure @@ -62,6 +70,17 @@ def __init__( self.enable_cmc = enable_cmc self.cmc = CMC(CMCConfig(method=cmc_method, downscale=cmc_downscale)) if enable_cmc else None + self.mask_manager = mask_manager + if self.mask_manager is None and enable_mask_manager: + self.mask_manager = MaskManager( + mask_generator=DummyBoxMaskGenerator(), + mask_propagator=DummyIdentityMaskPropagator(), + ) + + self._previous_frame: np.ndarray | None = None + self._previous_tracklets: list[TrackletSnapshot] = [] + self._last_mask_output: MaskOutput | None = None + def update( self, detections: sv.Detections, @@ -90,9 +109,24 @@ def update( """ self.frame_id += 1 + # For the convenience and better understanding. McByte processes uses previous + # frame and current frame. It is better to keep the method argument as "frame", + # as in case of the other trackers. + current_frame = frame + + if self.mask_manager is not None and current_frame is not None: + self._last_mask_output = self.mask_manager.get_updated_masks( + frame=current_frame, + previous_frame=self._previous_frame, + previous_tracklets=self._previous_tracklets, + ) + else: + self._last_mask_output = None + if len(self.tracks) == 0 and len(detections) == 0: result = sv.Detections.empty() result.tracker_id = np.array([], dtype=int) + self._store_previous_mask_inputs(current_frame, result) return result out_det_indices: list[int] = [] @@ -132,9 +166,9 @@ def update( unconfirmed_tracks.append(track) # CMC: apply to all predicted tracks before association - if self.enable_cmc and self.cmc is not None and frame is not None: + if self.enable_cmc and self.cmc is not None and current_frame is not None: mask_boxes = high_boxes if len(high_boxes) > 0 else None - H = self.cmc.estimate(frame, mask_boxes) + H = self.cmc.estimate(current_frame, mask_boxes) CMC.apply_batch(H, self.tracks) # Step 1: associate high-confidence detections to confirmed + lost tracks. # Lost tracks are included here (following the original ByteTrack), and @@ -230,12 +264,36 @@ def update( if not out_det_indices: result = sv.Detections.empty() result.tracker_id = np.array([], dtype=int) + self._store_previous_mask_inputs(current_frame, result) return result idx = np.array(out_det_indices) result = cast(sv.Detections, detections[idx]) result.tracker_id = np.array(out_tracker_ids, dtype=int) + self._store_previous_mask_inputs(current_frame, result) return result + + def _store_previous_mask_inputs( + self, + frame: np.ndarray | None, + detections: sv.Detections, + ) -> None: + """Store current tracker output for mask preparation on the next frame.""" + self._previous_frame = None if frame is None else frame.copy() + self._previous_tracklets = [] + + if detections.tracker_id is None: + return + + for xyxy, tracker_id in zip(detections.xyxy, detections.tracker_id): + if tracker_id < 0: + continue + self._previous_tracklets.append( + TrackletSnapshot( + tracker_id=int(tracker_id), + xyxy=xyxy.copy(), + ) + ) def _get_iou_matrix(self, tracklets: list[McByteTracklet], detections: np.ndarray) -> np.ndarray: if len(tracklets) == 0: diff --git a/tests/core/test_mcbyte_mask_manager.py b/tests/core/test_mcbyte_mask_manager.py new file mode 100644 index 00000000..46a0a8e6 --- /dev/null +++ b/tests/core/test_mcbyte_mask_manager.py @@ -0,0 +1,168 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +from __future__ import annotations + +import numpy as np + +from trackers.core.mcbyte.mask_manager import MaskManager +from trackers.core.mcbyte.masks import ( + DummyBoxMaskGenerator, + DummyIdentityMaskPropagator, + TrackletSnapshot, +) + + +def _make_frame(h: int = 100, w: int = 120) -> np.ndarray: + return np.zeros((h, w, 3), dtype=np.uint8) + + +def test_dummy_box_mask_generator_returns_expected_shape() -> None: + generator = DummyBoxMaskGenerator() + frame = _make_frame() + + output = generator.generate( + frame=frame, + tracklets=[ + TrackletSnapshot( + tracker_id=7, + xyxy=np.array([10, 20, 30, 40], dtype=np.float32), + ) + ], + ) + + assert output.masks is not None + assert output.masks.shape == (1, 100, 120) + assert output.tracklet_mask_dict == {7: 0} + + +def test_dummy_box_mask_generator_fills_detection_box() -> None: + generator = DummyBoxMaskGenerator() + frame = _make_frame() + + output = generator.generate( + frame=frame, + tracklets=[ + TrackletSnapshot( + tracker_id=7, + xyxy=np.array([10, 20, 30, 40], dtype=np.float32), + ) + ], + ) + + assert output.masks is not None + assert output.masks[0, 20:40, 10:30].all() + assert not output.masks[0, :10, :10].any() + assert output.masks.sum() == 20 * 20 + + +def test_mask_manager_returns_none_without_previous_inputs() -> None: + manager = MaskManager( + mask_generator=DummyBoxMaskGenerator(), + mask_propagator=DummyIdentityMaskPropagator(), + ) + + output = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=None, + previous_tracklets=[], + ) + + assert output is None + + +def test_mask_manager_generates_masks_from_previous_frame_inputs() -> None: + manager = MaskManager( + mask_generator=DummyBoxMaskGenerator(), + mask_propagator=None, + ) + + output = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=_make_frame(), + previous_tracklets=[ + TrackletSnapshot( + tracker_id=3, + xyxy=np.array([5, 6, 25, 30], dtype=np.float32), + ) + ], + ) + + assert output is not None + assert output.masks is not None + assert output.tracklet_mask_dict == {3: 0} + assert output.masks.shape == (1, 100, 120) + assert output.masks.sum() == 20 * 24 + + +def test_mask_manager_uses_propagator_after_initialization() -> None: + manager = MaskManager( + mask_generator=DummyBoxMaskGenerator(), + mask_propagator=DummyIdentityMaskPropagator(), + ) + + previous_tracklets = [ + TrackletSnapshot( + tracker_id=3, + xyxy=np.array([5, 6, 25, 30], dtype=np.float32), + ) + ] + + first_output = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=_make_frame(), + previous_tracklets=previous_tracklets, + ) + + second_output = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=_make_frame(), + previous_tracklets=[ + TrackletSnapshot( + tracker_id=99, + xyxy=np.array([50, 50, 70, 70], dtype=np.float32), + ) + ], + ) + + assert first_output is not None + assert second_output is not None + assert second_output.tracklet_mask_dict == first_output.tracklet_mask_dict + + +def test_mask_manager_reset_clears_state() -> None: + manager = MaskManager( + mask_generator=DummyBoxMaskGenerator(), + mask_propagator=DummyIdentityMaskPropagator(), + ) + + output_before_reset = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=_make_frame(), + previous_tracklets=[ + TrackletSnapshot( + tracker_id=3, + xyxy=np.array([5, 6, 25, 30], dtype=np.float32), + ) + ], + ) + + manager.reset() + + output_after_reset = manager.get_updated_masks( + frame=_make_frame(), + previous_frame=_make_frame(), + previous_tracklets=[ + TrackletSnapshot( + tracker_id=9, + xyxy=np.array([40, 40, 60, 60], dtype=np.float32), + ) + ], + ) + + assert output_before_reset is not None + assert output_after_reset is not None + assert output_after_reset.tracklet_mask_dict == {9: 0} \ No newline at end of file From ac528b55e512d6f330c6cad9f1e83f0c1bc1e67b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 May 2026 14:57:24 +0000 Subject: [PATCH 2/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/trackers/core/mcbyte/mask_manager.py | 4 ++-- src/trackers/core/mcbyte/masks/__init__.py | 2 +- src/trackers/core/mcbyte/masks/base.py | 2 +- src/trackers/core/mcbyte/masks/dummy.py | 10 +++------- src/trackers/core/mcbyte/tracker.py | 18 +++++++++--------- tests/core/test_mcbyte_mask_manager.py | 2 +- 6 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/trackers/core/mcbyte/mask_manager.py b/src/trackers/core/mcbyte/mask_manager.py index 20f214b7..4871638a 100644 --- a/src/trackers/core/mcbyte/mask_manager.py +++ b/src/trackers/core/mcbyte/mask_manager.py @@ -20,7 +20,7 @@ class MaskManager: """Manage McByte mask generation and propagation. The manager follows the original McByte timing: masks for the current frame - are prepared before association, but they are initialized/updated from tracker + are prepared before association, but they are initialized/updated from tracker outputs produced on the previous frame. """ @@ -69,4 +69,4 @@ def get_updated_masks( mask_output = self.mask_generator.generate(previous_frame, previous_tracklets) self.mask_propagator.initialize(previous_frame, mask_output) - return mask_output \ No newline at end of file + return mask_output diff --git a/src/trackers/core/mcbyte/masks/__init__.py b/src/trackers/core/mcbyte/masks/__init__.py index bb75a957..b873fcb9 100644 --- a/src/trackers/core/mcbyte/masks/__init__.py +++ b/src/trackers/core/mcbyte/masks/__init__.py @@ -22,4 +22,4 @@ "MaskOutput", "MaskPropagator", "TrackletSnapshot", -] \ No newline at end of file +] diff --git a/src/trackers/core/mcbyte/masks/base.py b/src/trackers/core/mcbyte/masks/base.py index f7894446..7b1d1122 100644 --- a/src/trackers/core/mcbyte/masks/base.py +++ b/src/trackers/core/mcbyte/masks/base.py @@ -61,4 +61,4 @@ def propagate( self, frame: np.ndarray, ) -> MaskOutput | None: - """Propagate masks to the current frame.""" \ No newline at end of file + """Propagate masks to the current frame.""" diff --git a/src/trackers/core/mcbyte/masks/dummy.py b/src/trackers/core/mcbyte/masks/dummy.py index b622feda..433c887c 100644 --- a/src/trackers/core/mcbyte/masks/dummy.py +++ b/src/trackers/core/mcbyte/masks/dummy.py @@ -64,9 +64,7 @@ def initialize( masks=None if mask_output.masks is None else mask_output.masks.copy(), tracklet_mask_dict=mask_output.tracklet_mask_dict.copy(), mask_avg_prob_dict=( - None - if mask_output.mask_avg_prob_dict is None - else mask_output.mask_avg_prob_dict.copy() + None if mask_output.mask_avg_prob_dict is None else mask_output.mask_avg_prob_dict.copy() ), ) @@ -81,8 +79,6 @@ def propagate( masks=None if self._mask_output.masks is None else self._mask_output.masks.copy(), tracklet_mask_dict=self._mask_output.tracklet_mask_dict.copy(), mask_avg_prob_dict=( - None - if self._mask_output.mask_avg_prob_dict is None - else self._mask_output.mask_avg_prob_dict.copy() + None if self._mask_output.mask_avg_prob_dict is None else self._mask_output.mask_avg_prob_dict.copy() ), - ) \ No newline at end of file + ) diff --git a/src/trackers/core/mcbyte/tracker.py b/src/trackers/core/mcbyte/tracker.py index f12c1b2c..83fd4e50 100644 --- a/src/trackers/core/mcbyte/tracker.py +++ b/src/trackers/core/mcbyte/tracker.py @@ -12,6 +12,12 @@ from scipy.optimize import linear_sum_assignment from trackers.core.base import BaseTracker +from trackers.core.mcbyte.mask_manager import MaskManager +from trackers.core.mcbyte.masks.base import MaskOutput, TrackletSnapshot +from trackers.core.mcbyte.masks.dummy import ( + DummyBoxMaskGenerator, + DummyIdentityMaskPropagator, +) from trackers.core.mcbyte.tracklet import McByteTracklet from trackers.core.mcbyte.utils import _fuse_score, get_alive_tracklets from trackers.utils.cmc import CMC, CMCConfig, CMCMethod @@ -21,12 +27,6 @@ BaseStateEstimator, XCYCWHStateEstimator, ) -from trackers.core.mcbyte.mask_manager import MaskManager -from trackers.core.mcbyte.masks.base import MaskOutput, TrackletSnapshot -from trackers.core.mcbyte.masks.dummy import ( - DummyBoxMaskGenerator, - DummyIdentityMaskPropagator, -) class McByteTracker(BaseTracker): @@ -109,7 +109,7 @@ def update( """ self.frame_id += 1 - # For the convenience and better understanding. McByte processes uses previous + # For the convenience and better understanding. McByte processes uses previous # frame and current frame. It is better to keep the method argument as "frame", # as in case of the other trackers. current_frame = frame @@ -122,7 +122,7 @@ def update( ) else: self._last_mask_output = None - + if len(self.tracks) == 0 and len(detections) == 0: result = sv.Detections.empty() result.tracker_id = np.array([], dtype=int) @@ -272,7 +272,7 @@ def update( result.tracker_id = np.array(out_tracker_ids, dtype=int) self._store_previous_mask_inputs(current_frame, result) return result - + def _store_previous_mask_inputs( self, frame: np.ndarray | None, diff --git a/tests/core/test_mcbyte_mask_manager.py b/tests/core/test_mcbyte_mask_manager.py index 46a0a8e6..95414fa7 100644 --- a/tests/core/test_mcbyte_mask_manager.py +++ b/tests/core/test_mcbyte_mask_manager.py @@ -165,4 +165,4 @@ def test_mask_manager_reset_clears_state() -> None: assert output_before_reset is not None assert output_after_reset is not None - assert output_after_reset.tracklet_mask_dict == {9: 0} \ No newline at end of file + assert output_after_reset.tracklet_mask_dict == {9: 0} From 721cd64d73282ed8127ba164c3cdf58c55bbea1d Mon Sep 17 00:00:00 2001 From: Tomasz Stanczyk Date: Mon, 18 May 2026 17:36:23 +0200 Subject: [PATCH 3/6] Address selected review comments for McByte mask manager --- src/trackers/core/mcbyte/mask_manager.py | 22 +++++----- src/trackers/core/mcbyte/tracker.py | 56 +++++++++++++++++++++--- tests/core/test_mcbyte_mask_manager.py | 10 ++--- tests/core/test_mcbyte_tracker.py | 38 ++++++++++++++++ 4 files changed, 103 insertions(+), 23 deletions(-) diff --git a/src/trackers/core/mcbyte/mask_manager.py b/src/trackers/core/mcbyte/mask_manager.py index 4871638a..3cff7da9 100644 --- a/src/trackers/core/mcbyte/mask_manager.py +++ b/src/trackers/core/mcbyte/mask_manager.py @@ -48,25 +48,25 @@ def get_updated_masks( No masks are returned until at least one previous frame and previous tracker output are available. + + If a propagator is configured, masks are initialized from the previous + frame and propagated to the current frame. If propagation is unavailable + or fails, ``None`` is returned to avoid using stale or misaligned masks. """ if previous_frame is None or len(previous_tracklets) == 0: return None - if not self._initialized or self.mask_propagator is None: + if self.mask_propagator is None: + return None + + if not self._initialized: mask_output = self.mask_generator.generate(previous_frame, previous_tracklets) + self.mask_propagator.initialize(previous_frame, mask_output) self._initialized = True - if self.mask_propagator is not None: - self.mask_propagator.initialize(previous_frame, mask_output) - propagated_output = self.mask_propagator.propagate(frame) - return propagated_output if propagated_output is not None else mask_output - - return mask_output - propagated_output = self.mask_propagator.propagate(frame) if propagated_output is not None: return propagated_output - mask_output = self.mask_generator.generate(previous_frame, previous_tracklets) - self.mask_propagator.initialize(previous_frame, mask_output) - return mask_output + self._initialized = False + return None diff --git a/src/trackers/core/mcbyte/tracker.py b/src/trackers/core/mcbyte/tracker.py index 83fd4e50..917115b3 100644 --- a/src/trackers/core/mcbyte/tracker.py +++ b/src/trackers/core/mcbyte/tracker.py @@ -30,6 +30,39 @@ class McByteTracker(BaseTracker): + """McByte-style multi-object tracker. + + This tracker currently provides the initial McByte integration skeleton, + built on top of IoU association, Kalman-filter-based tracklets, optional camera + motion compensation, and optional mask-manager infrastructure. + + Args: + lost_track_buffer: Time buffer, in frames at 30 FPS, for keeping lost + tracks alive before deletion. This value is scaled by ``frame_rate``. + frame_rate: Video frame rate used to scale ``lost_track_buffer``. + track_activation_threshold: Minimum confidence required to spawn a new + track. + minimum_consecutive_frames: Number of successful updates required before + assigning a stable track ID. + minimum_iou_threshold_first_assoc: Minimum similarity threshold for the + first association stage. + minimum_iou_threshold_second_assoc: Minimum similarity threshold for the + second association stage. + minimum_iou_threshold_unconfirmed_assoc: Minimum similarity threshold for + matching unconfirmed tracks. + high_conf_det_threshold: Confidence threshold used to split detections + into high- and low-confidence groups. + enable_cmc: Whether to enable camera motion compensation. + cmc_method: Camera motion compensation method. + cmc_downscale: Downscale factor used by camera motion compensation. + instant_first_frame_activation: Whether tracks spawned on the first frame + receive confirmed IDs immediately. + state_estimator_class: State estimator class used by McByte tracklets. + iou: IoU implementation used for association. + enable_mask_manager: Whether to create the default dummy mask manager. + mask_manager: Optional custom mask manager instance. + """ + tracker_id = "mcbyte" def __init__( @@ -279,22 +312,29 @@ def _store_previous_mask_inputs( detections: sv.Detections, ) -> None: """Store current tracker output for mask preparation on the next frame.""" - self._previous_frame = None if frame is None else frame.copy() + self._previous_frame = None self._previous_tracklets = [] - if detections.tracker_id is None: + if self.mask_manager is None or frame is None or detections.tracker_id is None: return + previous_tracklets = [] for xyxy, tracker_id in zip(detections.xyxy, detections.tracker_id): if tracker_id < 0: continue - self._previous_tracklets.append( + previous_tracklets.append( TrackletSnapshot( tracker_id=int(tracker_id), xyxy=xyxy.copy(), ) ) + if len(previous_tracklets) == 0: + return + + self._previous_frame = frame.copy() + self._previous_tracklets = previous_tracklets + def _get_iou_matrix(self, tracklets: list[McByteTracklet], detections: np.ndarray) -> np.ndarray: if len(tracklets) == 0: tracklet_boxes = np.empty((0, 4)) @@ -374,12 +414,18 @@ def _spawn_new_tracks( out_tracker_ids.append(tracklet.tracker_id) def reset(self) -> None: - """Reset tracker state by clearing all tracks and resetting ID counter. - Call this method when switching to a new video or scene. + """Reset tracker state by clearing all tracks, resetting ID counter, camera + motion compensation and mask manager. Call this method when switching to a new + video or scene. """ self.tracks = [] self.frame_id = 0 McByteTracklet.count_id = 0 + self._previous_frame = None + self._previous_tracklets = [] + self._last_mask_output = None + if self.mask_manager is not None: + self.mask_manager.reset() if self.cmc is not None: self.cmc.reset() diff --git a/tests/core/test_mcbyte_mask_manager.py b/tests/core/test_mcbyte_mask_manager.py index 95414fa7..2134a96b 100644 --- a/tests/core/test_mcbyte_mask_manager.py +++ b/tests/core/test_mcbyte_mask_manager.py @@ -59,7 +59,7 @@ def test_dummy_box_mask_generator_fills_detection_box() -> None: assert output.masks.sum() == 20 * 20 -def test_mask_manager_returns_none_without_previous_inputs() -> None: +def test_mask_manager_returns_none_without_previous_frame_or_tracklets() -> None: manager = MaskManager( mask_generator=DummyBoxMaskGenerator(), mask_propagator=DummyIdentityMaskPropagator(), @@ -74,7 +74,7 @@ def test_mask_manager_returns_none_without_previous_inputs() -> None: assert output is None -def test_mask_manager_generates_masks_from_previous_frame_inputs() -> None: +def test_mask_manager_returns_none_without_propagator() -> None: manager = MaskManager( mask_generator=DummyBoxMaskGenerator(), mask_propagator=None, @@ -91,11 +91,7 @@ def test_mask_manager_generates_masks_from_previous_frame_inputs() -> None: ], ) - assert output is not None - assert output.masks is not None - assert output.tracklet_mask_dict == {3: 0} - assert output.masks.shape == (1, 100, 120) - assert output.masks.sum() == 20 * 24 + assert output is None def test_mask_manager_uses_propagator_after_initialization() -> None: diff --git a/tests/core/test_mcbyte_tracker.py b/tests/core/test_mcbyte_tracker.py index 12417995..30b657d7 100644 --- a/tests/core/test_mcbyte_tracker.py +++ b/tests/core/test_mcbyte_tracker.py @@ -41,3 +41,41 @@ def test_mcbyte_instantiates_and_updates_with_frame_and_sparse_opt_flow_cmc_retu assert result.tracker_id is not None assert result.tracker_id[0] >= 0 assert len(tracker.tracks) == 1 + +def test_mcbyte_reset_clears_mask_state() -> None: + """reset() clears tracker and mask-manager temporal state.""" + tracker = McByteTracker( + enable_cmc=False, + enable_mask_manager=True, + minimum_consecutive_frames=1, + ) + + frame = _make_frame() + + tracker.update(_detection((100.0, 100.0, 200.0, 200.0)), frame) + tracker.update(_detection((100.0, 100.0, 200.0, 200.0)), frame) + + assert tracker._previous_frame is not None + assert len(tracker._previous_tracklets) == 1 + assert tracker._last_mask_output is not None + + tracker.reset() + + assert tracker._previous_frame is None + assert tracker._previous_tracklets == [] + assert tracker._last_mask_output is None + +def test_mcbyte_does_not_store_previous_frame_without_mask_manager() -> None: + """McByteTracker avoids frame copies when mask manager is disabled.""" + tracker = McByteTracker( + enable_cmc=False, + enable_mask_manager=False, + minimum_consecutive_frames=1, + ) + + frame = _make_frame() + + tracker.update(_detection((100.0, 100.0, 200.0, 200.0)), frame) + + assert tracker._previous_frame is None + assert tracker._previous_tracklets == [] \ No newline at end of file From c6843d056c83f3218cc5b0507ef0f05357651e9d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 May 2026 16:06:41 +0000 Subject: [PATCH 4/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/trackers/core/mcbyte/tracker.py | 6 +++--- tests/core/test_mcbyte_tracker.py | 4 +++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/trackers/core/mcbyte/tracker.py b/src/trackers/core/mcbyte/tracker.py index 917115b3..2a101b44 100644 --- a/src/trackers/core/mcbyte/tracker.py +++ b/src/trackers/core/mcbyte/tracker.py @@ -33,7 +33,7 @@ class McByteTracker(BaseTracker): """McByte-style multi-object tracker. This tracker currently provides the initial McByte integration skeleton, - built on top of IoU association, Kalman-filter-based tracklets, optional camera + built on top of IoU association, Kalman-filter-based tracklets, optional camera motion compensation, and optional mask-manager infrastructure. Args: @@ -414,8 +414,8 @@ def _spawn_new_tracks( out_tracker_ids.append(tracklet.tracker_id) def reset(self) -> None: - """Reset tracker state by clearing all tracks, resetting ID counter, camera - motion compensation and mask manager. Call this method when switching to a new + """Reset tracker state by clearing all tracks, resetting ID counter, camera + motion compensation and mask manager. Call this method when switching to a new video or scene. """ self.tracks = [] diff --git a/tests/core/test_mcbyte_tracker.py b/tests/core/test_mcbyte_tracker.py index 30b657d7..38e28e59 100644 --- a/tests/core/test_mcbyte_tracker.py +++ b/tests/core/test_mcbyte_tracker.py @@ -42,6 +42,7 @@ def test_mcbyte_instantiates_and_updates_with_frame_and_sparse_opt_flow_cmc_retu assert result.tracker_id[0] >= 0 assert len(tracker.tracks) == 1 + def test_mcbyte_reset_clears_mask_state() -> None: """reset() clears tracker and mask-manager temporal state.""" tracker = McByteTracker( @@ -65,6 +66,7 @@ def test_mcbyte_reset_clears_mask_state() -> None: assert tracker._previous_tracklets == [] assert tracker._last_mask_output is None + def test_mcbyte_does_not_store_previous_frame_without_mask_manager() -> None: """McByteTracker avoids frame copies when mask manager is disabled.""" tracker = McByteTracker( @@ -78,4 +80,4 @@ def test_mcbyte_does_not_store_previous_frame_without_mask_manager() -> None: tracker.update(_detection((100.0, 100.0, 200.0, 200.0)), frame) assert tracker._previous_frame is None - assert tracker._previous_tracklets == [] \ No newline at end of file + assert tracker._previous_tracklets == [] From 588b8526c008cb1b7d5484057abcf7792d85143d Mon Sep 17 00:00:00 2001 From: Tomasz Stanczyk Date: Wed, 20 May 2026 19:43:35 +0200 Subject: [PATCH 5/6] Limit McByte dummy mask classes to explicit imports --- src/trackers/core/mcbyte/masks/__init__.py | 6 ------ tests/core/test_mcbyte_mask_manager.py | 7 ++++--- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/src/trackers/core/mcbyte/masks/__init__.py b/src/trackers/core/mcbyte/masks/__init__.py index b873fcb9..ea911a75 100644 --- a/src/trackers/core/mcbyte/masks/__init__.py +++ b/src/trackers/core/mcbyte/masks/__init__.py @@ -10,14 +10,8 @@ MaskPropagator, TrackletSnapshot, ) -from trackers.core.mcbyte.masks.dummy import ( - DummyBoxMaskGenerator, - DummyIdentityMaskPropagator, -) __all__ = [ - "DummyBoxMaskGenerator", - "DummyIdentityMaskPropagator", "MaskGenerator", "MaskOutput", "MaskPropagator", diff --git a/tests/core/test_mcbyte_mask_manager.py b/tests/core/test_mcbyte_mask_manager.py index 2134a96b..7f2213a7 100644 --- a/tests/core/test_mcbyte_mask_manager.py +++ b/tests/core/test_mcbyte_mask_manager.py @@ -9,10 +9,11 @@ import numpy as np from trackers.core.mcbyte.mask_manager import MaskManager -from trackers.core.mcbyte.masks import ( +from trackers.core.mcbyte.masks import TrackletSnapshot + +from trackers.core.mcbyte.masks.dummy import ( DummyBoxMaskGenerator, - DummyIdentityMaskPropagator, - TrackletSnapshot, + DummyIdentityMaskPropagator ) From dc5382d42be0d43df1c540529772d35628af42b2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 20 May 2026 17:44:49 +0000 Subject: [PATCH 6/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/core/test_mcbyte_mask_manager.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/core/test_mcbyte_mask_manager.py b/tests/core/test_mcbyte_mask_manager.py index 7f2213a7..5e0c8e2b 100644 --- a/tests/core/test_mcbyte_mask_manager.py +++ b/tests/core/test_mcbyte_mask_manager.py @@ -10,11 +10,7 @@ from trackers.core.mcbyte.mask_manager import MaskManager from trackers.core.mcbyte.masks import TrackletSnapshot - -from trackers.core.mcbyte.masks.dummy import ( - DummyBoxMaskGenerator, - DummyIdentityMaskPropagator -) +from trackers.core.mcbyte.masks.dummy import DummyBoxMaskGenerator, DummyIdentityMaskPropagator def _make_frame(h: int = 100, w: int = 120) -> np.ndarray: