Refactor calibration

foosball/__init__.py

@@ -6,7 +6,7 @@
 
 from const import CALIBRATION_MODE, CALIBRATION_IMG_PATH, CALIBRATION_VIDEO, CALIBRATION_SAMPLE_SIZE, ARUCO_BOARD, \
     FILE, CAMERA_ID, FRAMERATE, OUTPUT, CAPTURE, DISPLAY, BALL, XPAD, YPAD, SCALE, VERBOSE, HEADLESS, OFF, \
-    MAX_PIPE_SIZE, INFO_VERBOSITY, GPU, AUDIO, WEBHOOK, BUFFER
+    MAX_PIPE_SIZE, INFO_VERBOSITY, GPU, AUDIO, WEBHOOK, BUFFER, BallPresets, CalibrationMode
 from foosball.arUcos.calibration import print_aruco_board, calibrate_camera
 from foosball.tracking.ai import AI
 
@@ -39,7 +39,7 @@ def get_argparse():
     ap.add_argument("-conf", "--config", type=file_path, default=None)
 
     calibration = ap.add_argument_group(title="Calibration", description="Options for camera or setup calibration")
-    calibration.add_argument("-c", f"--{CALIBRATION_MODE}", choices=['ball', 'goal', 'cam'], help="Calibration mode")
+    calibration.add_argument("-c", f"--{CALIBRATION_MODE}", choices=[CalibrationMode.BALL, CalibrationMode.GOAL, CalibrationMode.CAM], help="Calibration mode")
     calibration.add_argument("-ci", f"--{CALIBRATION_IMG_PATH}", type=dir_path, default=None,
                              help="Images path for calibration mode. If not given switching to live calibration")
     calibration.add_argument("-cv", f"--{CALIBRATION_VIDEO}", type=file_path, default=None,
@@ -57,7 +57,7 @@ def get_argparse():
     io.add_argument("-d", f"--{DISPLAY}", choices=['cv', 'gear'], default='cv', help="display backend cv=direct display, gear=stream")
 
     tracker = ap.add_argument_group(title="Tracker", description="Options for the ball/goal tracker")
-    tracker.add_argument("-ba", f"--{BALL}", choices=['yaml', 'orange', 'yellow'], default='yaml',
+    tracker.add_argument("-ba", f"--{BALL}", choices=[BallPresets.YAML, BallPresets.ORANGE, BallPresets.YELLOW], default=BallPresets.YAML,
                     help="Pre-configured ball color bounds. If 'yaml' is selected, a file called 'ball.yaml' "
                          "(stored by hitting 's' in ball calibration mode) will be loaded as a preset."
                          "If no file present fallback to 'yellow'")

foosball/detectors/__init__.py

@@ -1,15 +1,28 @@
 from abc import ABC, abstractmethod
 from typing import TypeVar, Generic
 
-from foosball.models import Frame
+from foosball.models import Frame, DetectedBall, DetectedGoals
 
 DetectorResult = TypeVar('DetectorResult')
+DetectorConfig = TypeVar('DetectorConfig')
+BallConfig = TypeVar('BallConfig')
+GoalConfig = TypeVar('GoalConfig')
 
 
-class Detector(ABC, Generic[DetectorResult]):
-    def __init__(self, *args, **kwargs):
-        pass
+class Detector(ABC, Generic[DetectorConfig, DetectorResult]):
+    def __init__(self, config: DetectorConfig, *args, **kwargs):
+        self.config = config
 
     @abstractmethod
     def detect(self, frame: Frame) -> DetectorResult | None:
         pass
+
+
+class BallDetector(Generic[BallConfig], Detector[BallConfig, DetectedBall], ABC):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+class GoalDetector(Generic[GoalConfig], Detector[GoalConfig, DetectedGoals], ABC):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

foosball/detectors/color.py

@@ -1,24 +1,21 @@
 import logging
 import os
-from abc import ABC
 from dataclasses import dataclass
-from typing import TypeVar, Generic
 
 import cv2
 import imutils
 import numpy as np
 import yaml
 
-from . import Detector, DetectorResult
-from ..models import Frame, DetectedGoals, Point, Goal, Blob, Goals, DetectedBall, HSV
+from const import BallPresets
+from . import BallDetector, GoalDetector
+from ..models import Frame, DetectedGoals, Point, Goal, Blob, Goals, DetectedBall, HSV, rgb2hsv
 
 logger = logging.getLogger(__name__)
 
-DetectorConfig = TypeVar('DetectorConfig')
-
 
 @dataclass
-class BallConfig:
+class BallColorConfig:
     bounds: [HSV, HSV]
     invert_frame: bool = False
     invert_mask: bool = False
@@ -29,21 +26,48 @@ def store(self):
         with open(filename, "w") as f:
             yaml.dump(self.to_dict(), f)
 
+    @staticmethod
+    def yellow():
+        lower = rgb2hsv(np.array([140, 86, 73]))
+        upper = rgb2hsv(np.array([0, 255, 94]))
+
+        return BallColorConfig(bounds=[lower, upper], invert_frame=False, invert_mask=False)
+
+    @staticmethod
+    def orange():
+        lower = rgb2hsv(np.array([166, 94, 72]))
+        upper = rgb2hsv(np.array([0, 249, 199]))
+
+        return BallColorConfig(bounds=[lower, upper], invert_frame=False, invert_mask=False)
+
+    @staticmethod
+    def preset(ball: str):
+        match ball:
+            case BallPresets.YAML:
+                return BallColorConfig.load() or BallColorConfig.yellow()
+            case BallPresets.ORANGE:
+                return BallColorConfig.orange()
+            case BallPresets.YELLOW:
+                return BallColorConfig.yellow()
+            case _:
+                logging.error("Unknown ball color. Falling back to 'yellow'")
+                return BallColorConfig.yellow()
+
     @staticmethod
     def load(filename='ball.yaml'):
         if os.path.isfile(filename):
             logging.info("Loading ball config ball.yaml")
             with open(filename, 'r') as f:
                 c = yaml.safe_load(f)
-                return BallConfig(invert_frame=c['invert_frame'], invert_mask=c['invert_mask'],
-                                  bounds=np.array(c['bounds']))
+                return BallColorConfig(invert_frame=c['invert_frame'], invert_mask=c['invert_mask'],
+                                       bounds=np.array(c['bounds']))
         else:
             logging.info("No ball config found")
         return None
 
     def __eq__(self, other):
         """Overrides the default implementation"""
-        if isinstance(other, BallConfig):
+        if isinstance(other, BallColorConfig):
             return (all([a == b for a, b in zip(self.bounds[0], other.bounds[0])]) and
                     all([a == b for a, b in zip(self.bounds[1], other.bounds[1])]) and
                     self.invert_mask == other.invert_mask and
@@ -58,13 +82,19 @@ def to_dict(self):
         }
 
 
-
 @dataclass
-class GoalConfig:
+class GoalColorConfig:
     bounds: [int, int]
     invert_frame: bool = True
     invert_mask: bool = True
 
+    @staticmethod
+    def preset():
+        default_config = GoalColorConfig(bounds=[0, 235], invert_frame=True, invert_mask=True)
+        if os.path.isfile('goal.yaml'):
+            return GoalColorConfig.load() or default_config
+        return default_config
+
     def store(self):
         filename = "goal.yaml"
         print(f"Store config {filename}" + (" " * 50), end="\n\n")
@@ -75,11 +105,11 @@ def store(self):
     def load(filename='goal.yaml'):
         with open(filename, 'r') as f:
             c = yaml.safe_load(f)
-            return GoalConfig(**c)
+            return GoalColorConfig(**c)
 
     def __eq__(self, other):
         """Overrides the default implementation"""
-        if isinstance(other, GoalConfig):
+        if isinstance(other, GoalColorConfig):
             return (all([a == b for a, b in zip(self.bounds, other.bounds)]) and
                     self.invert_mask == other.invert_mask and
                     self.invert_frame == other.invert_frame)
@@ -93,13 +123,7 @@ def to_dict(self):
         }
 
 
-class ColorDetector(Generic[DetectorConfig, DetectorResult], Detector[DetectorResult], ABC):
-    def __init__(self, config: DetectorConfig, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.config = config
-
-
-class BallDetector(ColorDetector[BallConfig, DetectedBall]):
+class BallColorDetector(BallDetector[BallColorConfig]):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -112,7 +136,7 @@ def detect(self, frame) -> DetectedBall:
             logger.error("Ball Detection not possible. Config is None")
 
 
-class GoalDetector(ColorDetector[GoalConfig, DetectedGoals]):
+class GoalColorDetector(GoalDetector[GoalColorConfig]):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -156,7 +180,7 @@ def detect(self, frame: Frame) -> DetectedGoals:
             logger.error("Goal Detection not possible. config is None")
 
 
-def filter_color_range(frame, config: BallConfig) -> Frame:
+def filter_color_range(frame, config: BallColorConfig) -> Frame:
     [lower, upper] = config.bounds
     f = frame if not config.invert_frame else cv2.bitwise_not(frame)
 
@@ -182,7 +206,7 @@ def filter_color_range(frame, config: BallConfig) -> Frame:
     return cv2.bitwise_and(f, f, mask=simple_mask)
 
 
-def filter_gray_range(frame, config: GoalConfig) -> Frame:
+def filter_gray_range(frame, config: GoalColorConfig) -> Frame:
     try:
         [lower, upper] = config.bounds
         f = frame if not config.invert_frame else cv2.bitwise_not(frame)

foosball/models.py

@@ -100,9 +100,15 @@ class DetectedGoals:
 Track = collections.deque
 
 
+class Verbosity(Enum):
+    TRACE = 0
+    DEBUG = 1
+    INFO = 2
+
+
 @dataclass
 class Info:
-    verbosity: int
+    verbosity: Verbosity
     title: str
     value: str
 
@@ -123,8 +129,9 @@ def append(self, info: Info):
     def concat(self, info_log):
         self.infos = self.infos + info_log.infos
 
-    def filter(self, info_verbosity=0):
-        return InfoLog(infos=[i for i in self.infos if info_verbosity <= i.verbosity])
+    def filter(self, infoVerbosity: Verbosity = Verbosity.TRACE):
+        return InfoLog(
+            infos=[i for i in self.infos if infoVerbosity is not None and infoVerbosity.value <= i.verbosity.value])
 
     def to_string(self):
         return " - ".join([i.to_string() for i in self.infos])