basetrack.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. from collections import OrderedDict
  3. import numpy as np
  4. class TrackState:
  5. """Enumeration of possible object tracking states."""
  6. New = 0
  7. Tracked = 1
  8. Lost = 2
  9. Removed = 3
  10. class BaseTrack:
  11. """Base class for object tracking, handling basic track attributes and operations."""
  12. _count = 0
  13. track_id = 0
  14. is_activated = False
  15. state = TrackState.New
  16. history = OrderedDict()
  17. features = []
  18. curr_feature = None
  19. score = 0
  20. start_frame = 0
  21. frame_id = 0
  22. time_since_update = 0
  23. # Multi-camera
  24. location = (np.inf, np.inf)
  25. @property
  26. def end_frame(self):
  27. """Return the last frame ID of the track."""
  28. return self.frame_id
  29. @staticmethod
  30. def next_id():
  31. """Increment and return the global track ID counter."""
  32. BaseTrack._count += 1
  33. return BaseTrack._count
  34. def activate(self, *args):
  35. """Activate the track with the provided arguments."""
  36. raise NotImplementedError
  37. def predict(self):
  38. """Predict the next state of the track."""
  39. raise NotImplementedError
  40. def update(self, *args, **kwargs):
  41. """Update the track with new observations."""
  42. raise NotImplementedError
  43. def mark_lost(self):
  44. """Mark the track as lost."""
  45. self.state = TrackState.Lost
  46. def mark_removed(self):
  47. """Mark the track as removed."""
  48. self.state = TrackState.Removed
  49. @staticmethod
  50. def reset_id():
  51. """Reset the global track ID counter."""
  52. BaseTrack._count = 0