diff --git a/boxmot/__init__.py b/boxmot/__init__.py index f716c0d1af..4251df7537 100644 --- a/boxmot/__init__.py +++ b/boxmot/__init__.py @@ -4,17 +4,17 @@ from boxmot.postprocessing.gsi import gsi from boxmot.tracker_zoo import create_tracker, get_tracker_config -from boxmot.trackers.botsort.botsort import BoTSORT -from boxmot.trackers.bytetrack.byte_tracker import BYTETracker -from boxmot.trackers.deepocsort.deep_ocsort import DeepOCSort as DeepOCSORT -from boxmot.trackers.hybridsort.hybridsort import HybridSORT -from boxmot.trackers.ocsort.ocsort import OCSort as OCSORT -from boxmot.trackers.strongsort.strong_sort import StrongSORT -from boxmot.trackers.imprassoc.impr_assoc_tracker import ImprAssocTrack +from boxmot.trackers.botsort.botsort import BotSort +from boxmot.trackers.bytetrack.bytetrack import ByteTrack +from boxmot.trackers.deepocsort.deepocsort import DeepOcSort +from boxmot.trackers.hybridsort.hybridsort import HybridSort +from boxmot.trackers.ocsort.ocsort import OcSort +from boxmot.trackers.strongsort.strongsort import StrongSort +from boxmot.trackers.imprassoc.imprassoctrack import ImprAssocTrack TRACKERS = ['bytetrack', 'botsort', 'strongsort', 'ocsort', 'deepocsort', 'hybridsort', 'imprassoc'] __all__ = ("__version__", - "StrongSORT", "OCSORT", "BYTETracker", "BoTSORT", "DeepOCSORT", "HybridSORT", "ImprAssocTrack" + "StrongSort", "OcSort", "ByteTrack", "BotSort", "DeepOcSort", "HybridSort", "ImprAssocTrack" "create_tracker", "get_tracker_config", "gsi") diff --git a/boxmot/tracker_zoo.py b/boxmot/tracker_zoo.py index 3ecd32886b..be05e13c95 100644 --- a/boxmot/tracker_zoo.py +++ b/boxmot/tracker_zoo.py @@ -41,13 +41,13 @@ def create_tracker(tracker_type, tracker_config=None, reid_weights=None, device= # Map tracker types to their corresponding classes tracker_mapping = { - 'strongsort': 'boxmot.trackers.strongsort.strong_sort.StrongSORT', - 'ocsort': 'boxmot.trackers.ocsort.ocsort.OCSort', - 'bytetrack': 'boxmot.trackers.bytetrack.byte_tracker.BYTETracker', - 'botsort': 'boxmot.trackers.botsort.botsort.BoTSORT', - 'deepocsort': 'boxmot.trackers.deepocsort.deep_ocsort.DeepOCSort', - 'hybridsort': 'boxmot.trackers.hybridsort.hybridsort.HybridSORT', - 'imprassoc': 'boxmot.trackers.imprassoc.impr_assoc_tracker.ImprAssocTrack' + 'strongsort': 'boxmot.trackers.strongsort.strongsort.StrongSort', + 'ocsort': 'boxmot.trackers.ocsort.ocsort.OcSort', + 'bytetrack': 'boxmot.trackers.bytetrack.bytetrack.ByteTrack', + 'botsort': 'boxmot.trackers.botsort.botsort.BotSort', + 'deepocsort': 'boxmot.trackers.deepocsort.deepocsort.DeepOcSort', + 'hybridsort': 'boxmot.trackers.hybridsort.hybridsort.HybridSort', + 'imprassoc': 'boxmot.trackers.imprassoc.imprassoctrack.ImprAssocTrack' } # Check if the tracker type exists in the mapping diff --git a/boxmot/trackers/botsort/botsort.py b/boxmot/trackers/botsort/botsort.py index 40d878189a..3601d47a09 100644 --- a/boxmot/trackers/botsort/botsort.py +++ b/boxmot/trackers/botsort/botsort.py @@ -17,7 +17,7 @@ -class BoTSORT(BaseTracker): +class BotSort(BaseTracker): """ BoTSORT Tracker: A tracking algorithm that combines appearance and motion-based tracking. diff --git a/boxmot/trackers/bytetrack/byte_tracker.py b/boxmot/trackers/bytetrack/bytetrack.py similarity index 99% rename from boxmot/trackers/bytetrack/byte_tracker.py rename to boxmot/trackers/bytetrack/bytetrack.py index 32f2af99b3..a3357327ce 100644 --- a/boxmot/trackers/bytetrack/byte_tracker.py +++ b/boxmot/trackers/bytetrack/bytetrack.py @@ -115,7 +115,7 @@ def xyxy(self): return ret -class BYTETracker(BaseTracker): +class ByteTrack(BaseTracker): """ BYTETracker: A tracking algorithm based on ByteTrack, which utilizes motion-based tracking. diff --git a/boxmot/trackers/deepocsort/deep_ocsort.py b/boxmot/trackers/deepocsort/deepocsort.py similarity index 99% rename from boxmot/trackers/deepocsort/deep_ocsort.py rename to boxmot/trackers/deepocsort/deepocsort.py index aa8237f68e..158a233582 100644 --- a/boxmot/trackers/deepocsort/deep_ocsort.py +++ b/boxmot/trackers/deepocsort/deepocsort.py @@ -222,7 +222,7 @@ def mahalanobis(self, bbox): return self.kf.md_for_measurement(self.bbox_to_z_func(bbox)) -class DeepOCSort(BaseTracker): +class DeepOcSort(BaseTracker): """ DeepOCSort Tracker: A tracking algorithm that utilizes a combination of appearance and motion-based tracking. diff --git a/boxmot/trackers/hybridsort/hybridsort.py b/boxmot/trackers/hybridsort/hybridsort.py index 257fadf54c..ed52af1217 100644 --- a/boxmot/trackers/hybridsort/hybridsort.py +++ b/boxmot/trackers/hybridsort/hybridsort.py @@ -329,7 +329,7 @@ def get_state(self): return convert_x_to_bbox(self.kf.x) -class HybridSORT(BaseTracker): +class HybridSort(BaseTracker): """ HybridSORT Tracker: A tracking algorithm that utilizes a combination of appearance and motion-based tracking and temporal consistency models (TCM) for improved tracking accuracy and robustness. diff --git a/boxmot/trackers/imprassoc/impr_assoc_tracker.py b/boxmot/trackers/imprassoc/imprassoctrack.py similarity index 100% rename from boxmot/trackers/imprassoc/impr_assoc_tracker.py rename to boxmot/trackers/imprassoc/imprassoctrack.py diff --git a/boxmot/trackers/ocsort/ocsort.py b/boxmot/trackers/ocsort/ocsort.py index f2151a5442..0131463a32 100644 --- a/boxmot/trackers/ocsort/ocsort.py +++ b/boxmot/trackers/ocsort/ocsort.py @@ -179,7 +179,7 @@ def get_state(self): return convert_x_to_bbox(self.kf.x) -class OCSort(BaseTracker): +class OcSort(BaseTracker): """ OCSort Tracker: A tracking algorithm that utilizes motion-based tracking. diff --git a/boxmot/trackers/strongsort/strong_sort.py b/boxmot/trackers/strongsort/strongsort.py similarity index 99% rename from boxmot/trackers/strongsort/strong_sort.py rename to boxmot/trackers/strongsort/strongsort.py index 86a41374a0..0e180ff43a 100644 --- a/boxmot/trackers/strongsort/strong_sort.py +++ b/boxmot/trackers/strongsort/strongsort.py @@ -13,7 +13,7 @@ from boxmot.trackers.basetracker import BaseTracker -class StrongSORT(object): +class StrongSort(object): """ StrongSORT Tracker: A tracking algorithm that utilizes a combination of appearance and motion-based tracking. diff --git a/examples/torchvision_det_boxmot.ipynb b/examples/torchvision_det_boxmot.ipynb index 2ffdcfc8dd..72fbf821b5 100644 --- a/examples/torchvision_det_boxmot.ipynb +++ b/examples/torchvision_det_boxmot.ipynb @@ -12,7 +12,7 @@ "import numpy as np\n", "from pathlib import Path\n", "\n", - "from boxmot import BoTSORT\n", + "from boxmot import BotSort\n", "\n", "# Load a pre-trained Faster R-CNN model\n", "device = torch.device('cpu') # Use 'cuda' if you have a GPU\n", @@ -20,7 +20,7 @@ "detector.eval().to(device)\n", "\n", "# Initialize the tracker\n", - "tracker = BoTSORT(\n", + "tracker = BotSort(\n", " reid_weights=Path('osnet_x0_25_msmt17.pt'), # Path to ReID model\n", " device=device, # Use CPU for inference\n", " half=False\n", diff --git a/examples/torchvision_pose_boxmot.ipynb b/examples/torchvision_pose_boxmot.ipynb index a25175e960..548d8f00ab 100644 --- a/examples/torchvision_pose_boxmot.ipynb +++ b/examples/torchvision_pose_boxmot.ipynb @@ -11,14 +11,14 @@ "import cv2\n", "import numpy as np\n", "from pathlib import Path\n", - "from boxmot import BoTSORT\n", + "from boxmot import BotSort\n", "\n", "# Load a pre-trained Keypoint R-CNN model from torchvision\n", "device = torch.device('cpu') # Change to 'cuda' if you have a GPU available\n", "pose_model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)\n", "pose_model.eval().to(device)\n", "\n", - "tracker = BoTSORT(\n", + "tracker = BotSort(\n", " reid_weights=Path('osnet_x0_25_msmt17.pt'), # ReID model to use\n", " device=device,\n", " half=False,\n", diff --git a/examples/torchvision_seg_boxmot.ipynb b/examples/torchvision_seg_boxmot.ipynb index 3a21b67945..c5276274ad 100644 --- a/examples/torchvision_seg_boxmot.ipynb +++ b/examples/torchvision_seg_boxmot.ipynb @@ -11,14 +11,14 @@ "import cv2\n", "import numpy as np\n", "from pathlib import Path\n", - "from boxmot import BoTSORT\n", + "from boxmot import BotSort\n", "\n", "# Load a pre-trained Mask R-CNN model from torchvision\n", "device = torch.device('cpu') # Change to 'cuda' if you have a GPU available\n", "segmentation_model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\n", "segmentation_model.eval().to(device)\n", "\n", - "tracker = BoTSORT(\n", + "tracker = BotSort(\n", " reid_weights=Path('osnet_x0_25_msmt17.pt'), # ReID model to use\n", " device=device,\n", " half=False,\n", diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/performance/test_tracking_p.py b/tests/performance/test_tracking_p.py index cf3fdda424..06c9d97e55 100644 --- a/tests/performance/test_tracking_p.py +++ b/tests/performance/test_tracking_p.py @@ -7,13 +7,12 @@ from numpy.testing import assert_allclose from boxmot import ( - StrongSORT, BoTSORT, DeepOCSORT, OCSORT, BYTETracker, ImprAssocTrack, get_tracker_config, create_tracker, + StrongSort, BotSort, DeepOcSort, OcSort, ByteTrack, ImprAssocTrack, get_tracker_config, create_tracker, ) +from tests.test_config import MOTION_ONLY_TRACKING_NAMES, MOTION_N_APPEARANCE_TRACKING_NAMES -MOTION_N_APPEARANCE_TRACKING_METHODS=['botsort', 'deepocsort', 'strongsort', 'imprassoc'] -MOTION_ONLY_TRACKING_METHODS=['ocsort', 'bytetrack'] -@pytest.mark.parametrize("tracker_type", MOTION_ONLY_TRACKING_METHODS) +@pytest.mark.parametrize("tracker_type", MOTION_ONLY_TRACKING_NAMES) def test_motion_tracker_update_time(tracker_type): tracker_conf = get_tracker_config(tracker_type) tracker = create_tracker( @@ -41,7 +40,7 @@ def test_motion_tracker_update_time(tracker_type): assert elapsed_time_per_iteration < max_allowed_time, f"Tracking algorithms processing time exceeds the allowed limit: {elapsed_time_per_iteration} > {max_allowed_time}" -@pytest.mark.parametrize("tracker_type", MOTION_N_APPEARANCE_TRACKING_METHODS) +@pytest.mark.parametrize("tracker_type", MOTION_N_APPEARANCE_TRACKING_NAMES) def test_motion_n_appearance_tracker_update_time(tracker_type): tracker_conf = get_tracker_config(tracker_type) tracker = create_tracker( diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000000..7880055c8d --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,12 @@ +from boxmot import ( + StrongSort, BotSort, DeepOcSort, OcSort, ByteTrack, ImprAssocTrack, get_tracker_config, create_tracker, +) + +MOTION_N_APPEARANCE_TRACKING_NAMES = ['botsort', 'deepocsort', 'strongsort', 'imprassoc'] +MOTION_ONLY_TRACKING_NAMES = ['ocsort', 'bytetrack'] + +MOTION_N_APPEARANCE_TRACKING_METHODS=[StrongSort, BotSort, DeepOcSort, ImprAssocTrack] +MOTION_ONLY_TRACKING_METHODS=[OcSort, ByteTrack] + +ALL_TRACKERS = ['botsort', 'deepocsort', 'ocsort', 'bytetrack', 'strongsort', 'imprassoc'] +PER_CLASS_TRACKERS = ['botsort', 'deepocsort', 'ocsort', 'bytetrack', 'imprassoc'] \ No newline at end of file diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/test_trackers.py b/tests/unit/test_trackers.py index 763ed23bcb..e3e2b68ba2 100644 --- a/tests/unit/test_trackers.py +++ b/tests/unit/test_trackers.py @@ -6,18 +6,12 @@ from numpy.testing import assert_allclose from boxmot import ( - StrongSORT, BoTSORT, DeepOCSORT, OCSORT, BYTETracker, ImprAssocTrack, get_tracker_config, create_tracker, + StrongSort, BotSort, DeepOcSort, OcSort, ByteTrack, ImprAssocTrack, get_tracker_config, create_tracker, ) from boxmot.trackers.ocsort.ocsort import KalmanBoxTracker as OCSortKalmanBoxTracker -from boxmot.trackers.deepocsort.deep_ocsort import KalmanBoxTracker as DeepOCSortKalmanBoxTracker - - - -MOTION_ONLY_TRACKING_METHODS=[OCSORT, BYTETracker] -MOTION_N_APPEARANCE_TRACKING_METHODS=[StrongSORT, BoTSORT, DeepOCSORT, ImprAssocTrack] -ALL_TRACKERS=['botsort', 'deepocsort', 'ocsort', 'bytetrack', 'strongsort', 'imprassoc'] -PER_CLASS_TRACKERS=['botsort', 'deepocsort', 'ocsort', 'bytetrack', 'imprassoc'] +from boxmot.trackers.deepocsort.deepocsort import KalmanBoxTracker as DeepOCSortKalmanBoxTracker +from tests.test_config import MOTION_ONLY_TRACKING_METHODS, MOTION_N_APPEARANCE_TRACKING_METHODS, ALL_TRACKERS, PER_CLASS_TRACKERS @pytest.mark.parametrize("Tracker", MOTION_N_APPEARANCE_TRACKING_METHODS) @@ -56,7 +50,7 @@ def test_tracker_output_size(tracker_type): def test_dynamic_max_obs_based_on_max_age(): max_age = 400 - ocsort = OCSORT( + ocsort = OcSort( max_age=max_age ) @@ -84,14 +78,14 @@ def create_kalman_box_tracker_deepocsort(bbox, cls, det_ind, tracker): TRACKER_CREATORS = { - OCSORT: create_kalman_box_tracker_ocsort, - DeepOCSORT: create_kalman_box_tracker_deepocsort, + OcSort: create_kalman_box_tracker_ocsort, + DeepOcSort: create_kalman_box_tracker_deepocsort, } @pytest.mark.parametrize("Tracker, init_args", [ - (OCSORT, {}), - (DeepOCSORT, { + (OcSort, {}), + (DeepOcSort, { 'reid_weights': Path(WEIGHTS / 'osnet_x0_25_msmt17.pt'), 'device': 'cpu', 'half': True