Merge pull request #1683 from mikel-brostrom/standardize-tracker-naming

follow tracker name convention
mikel-brostrom · Oct 1, 2024 · 9b1ab41 · 9b1ab41
2 parents e1ec94c + e105138
commit 9b1ab41
Show file tree

Hide file tree

Showing 18 changed files with 51 additions and 46 deletions.
diff --git 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
@@ -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
@@ -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 → boxmot/trackers/bytetrack/bytetrack.py b/boxmot/trackers/bytetrack/byte_tracker.py → 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 → boxmot/trackers/deepocsort/deepocsort.py b/boxmot/trackers/deepocsort/deep_ocsort.py → 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
@@ -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/.../trackers/imprassoc/impr_assoc_tracker.py → boxmot/trackers/imprassoc/imprassoctrack.py b/.../trackers/imprassoc/impr_assoc_tracker.py → boxmot/trackers/imprassoc/imprassoctrack.py
diff --git 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 → boxmot/trackers/strongsort/strongsort.py b/boxmot/trackers/strongsort/strong_sort.py → 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
@@ -12,15 +12,15 @@
  "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",
  "detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
  "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
@@ -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
@@ -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
diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py
diff --git 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
@@ -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']
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
diff --git 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