diff --git a/fixtracks/utils/enums.py b/fixtracks/utils/enums.py index d379786..b3edbe4 100644 --- a/fixtracks/utils/enums.py +++ b/fixtracks/utils/enums.py @@ -1,7 +1,29 @@ from enum import Enum +from PySide6.QtGui import QColor + class DetectionData(Enum): ID = 0 FRAME = 1 COORDINATES = 2 - TRACK_ID = 3 \ No newline at end of file + TRACK_ID = 3 + +class Tracks(Enum): + TRACKONE = 1 + TRACKTWO = 2 + UNASSIGNED = -1 + + def toColor(self): + track_colors = { + Tracks.TRACKONE: QColor.fromString("orange"), + Tracks.TRACKTWO: QColor.fromString("green"), + Tracks.UNASSIGNED: QColor.fromString("red") + } + return track_colors.get(self, QColor(128, 128, 128)) # Default to black if not found + + @staticmethod + def fromValue(value): + for track in Tracks: + if track.value == value: + return track + return Tracks.UNASSIGNED