Change ball calibration from rgb to hsv

reduce complexity by selecting a hue range between +10 and -10 of the selected value

ball.yaml

@@ -1,8 +1,8 @@
-bounds_hsv:
-- - 6
-  - 122
-  - 140
-- - 71
+bounds:
+- - 17
+  - 37
+  - 25
+- - 37
   - 255
   - 255
 invert_frame: false

foosball/models.py

@@ -99,16 +99,10 @@ class GoalsDetectionResult:
 
 @dataclass
 class BallConfig:
-    bounds_hsv: [HSV, HSV]
+    bounds: [HSV, HSV]
     invert_frame: bool = False
     invert_mask: bool = False
 
-    def bounds(self, mode="hsv"):
-        if mode == "hsv":
-            return self.bounds_hsv
-        else:
-            return [hsv2rgb(x) for x in self.bounds_hsv]
-
     def store(self):
         filename = f"ball.yaml"
         print(f"Store config {filename}" + (" " * 50), end="\n\n")
@@ -121,14 +115,23 @@ def load(filename='ball.yaml'):
             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_hsv=np.array(c['bounds_hsv']))
+                return BallConfig(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):
+            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
+                    self.invert_frame == other.invert_frame)
+        return False
+
     def to_dict(self):
         return {
-            "bounds_hsv": [x.tolist() for x in self.bounds_hsv],
+            "bounds": [x.tolist() for x in self.bounds],
             "invert_frame": self.invert_frame,
             "invert_mask": self.invert_mask
         }
@@ -142,7 +145,6 @@ class GoalConfig:
 
     def store(self):
         filename = f"goal.yaml"
-        [lower, upper] = self.bounds
         print(f"Store config {filename}" + (" " * 50), end="\n\n")
         with open(filename, "w") as f:
             yaml.dump(self.to_dict(), f)
@@ -153,6 +155,14 @@ def load(filename='goal.yaml'):
             c = yaml.safe_load(f)
             return GoalConfig(**c)
 
+    def __eq__(self, other):
+        """Overrides the default implementation"""
+        if isinstance(other, GoalConfig):
+            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)
+        return False
+
     def to_dict(self):
         return {
             "bounds": self.bounds,

foosball/sink/opencv.py

@@ -7,7 +7,7 @@
 
 from . import Sink
 from foosball.tracking import BallConfig
-from foosball.models import rgb2hsv, hsv2rgb, GoalConfig
+from foosball.models import GoalConfig
 
 GOAL = "goal"
 BALL = "ball"
@@ -81,18 +81,15 @@ def add_config_input(calibration, config):
 
 
 def add_ball_config_input(bounds: BallConfig):
-    [lower_hsv, upper_hsv] = bounds.bounds_hsv
-    lower_rgb = hsv2rgb(lower_hsv)
-    upper_rgb = hsv2rgb(upper_hsv)
+    [lower_hsv, upper_hsv] = bounds.bounds
     cv2.createTrackbar(f'invert_frame', BALL, 1 if bounds.invert_frame else 0, 1, lambda v: None)
     cv2.createTrackbar(f'invert_mask', BALL, 1 if bounds.invert_mask else 0, 1, lambda v: None)
     # create trackbars for color change
-    cv2.createTrackbar(slider_label('R', 'low'), BALL, lower_rgb[0], 255, lambda v: None)
-    cv2.createTrackbar(slider_label('G', 'low'), BALL, lower_rgb[1], 255, lambda v: None)
-    cv2.createTrackbar(slider_label('B', 'low'), BALL, lower_rgb[2], 255, lambda v: None)
-    cv2.createTrackbar(slider_label('R', 'high'), BALL, upper_rgb[0], 255, lambda v: None)
-    cv2.createTrackbar(slider_label('G', 'high'), BALL, upper_rgb[1], 255, lambda v: None)
-    cv2.createTrackbar(slider_label('B', 'high'), BALL, upper_rgb[2], 255, lambda v: None)
+    cv2.createTrackbar('Hue', BALL, avg(lower_hsv[0], upper_hsv[0]), 179, lambda v: None)
+    cv2.createTrackbar(slider_label('S', 'low'), BALL, lower_hsv[1], 255, lambda v: None)
+    cv2.createTrackbar(slider_label('V', 'low'), BALL, lower_hsv[2], 255, lambda v: None)
+    cv2.createTrackbar(slider_label('S', 'high'), BALL, upper_hsv[1], 255, lambda v: None)
+    cv2.createTrackbar(slider_label('V', 'high'), BALL, upper_hsv[2], 255, lambda v: None)
     # cv2.createButton("Reset", reset_bounds, (name, lower_rgb, upper_rgb))
 
 
@@ -113,25 +110,25 @@ def reset_config(calibration, config):
         reset_ball_config(config)
 
 
+def avg(x, y):
+    return int((x + y) / 2)
+
 def reset_ball_config(bounds: BallConfig):
-    [lower_hsv, upper_hsv] = bounds.bounds_hsv
+    [lower_hsv, upper_hsv] = bounds.bounds
     print(f"Reset config {BALL}", end="\n\n\n")
-    lower_rgb = hsv2rgb(lower_hsv)
-    upper_rgb = hsv2rgb(upper_hsv)
 
     cv2.setTrackbarPos('invert_frame', BALL, 1 if bounds.invert_frame else 0)
     cv2.setTrackbarPos('invert_mask', BALL, 1 if bounds.invert_mask else 0)
 
-    cv2.setTrackbarPos(slider_label('R', 'low'), BALL, lower_rgb[0])
-    cv2.setTrackbarPos(slider_label('G', 'low'), BALL, lower_rgb[1])
-    cv2.setTrackbarPos(slider_label('B', 'low'), BALL, lower_rgb[2])
-    cv2.setTrackbarPos(slider_label('R', 'high'), BALL, upper_rgb[0])
-    cv2.setTrackbarPos(slider_label('G', 'high'), BALL, upper_rgb[1])
-    cv2.setTrackbarPos(slider_label('B', 'high'), BALL, upper_rgb[2])
+    cv2.setTrackbarPos('Hue', BALL, avg(lower_hsv[0], upper_hsv[0]))
+    cv2.setTrackbarPos(slider_label('S', 'low'), BALL, lower_hsv[1])
+    cv2.setTrackbarPos(slider_label('V', 'low'), BALL, lower_hsv[2])
+    cv2.setTrackbarPos(slider_label('S', 'high'), BALL, upper_hsv[1])
+    cv2.setTrackbarPos(slider_label('V', 'high'), BALL, upper_hsv[2])
 
 
 def reset_goal_config(config: GoalConfig):
-    [lower, upper] = config.bounds_hsv
+    [lower, upper] = config.bounds
     print(f"Reset config {GOAL}", end="\n\n\n")
 
     cv2.setTrackbarPos('invert_frame', GOAL, 1 if config.invert_frame else 0)
@@ -143,14 +140,12 @@ def reset_goal_config(config: GoalConfig):
 
 def store_ball_config(config: BallConfig):
     filename = f"ball.yaml"
-    [lower_hsv, upper_hsv] = config.bounds_hsv
+    [lower, upper] = config.bounds
     print(f"Store config {filename}" + (" " * 50), end="\n\n")
-    lower_rgb = hsv2rgb(lower_hsv)
-    upper_rgb = hsv2rgb(upper_hsv)
     with open(filename, "w") as f:
         yaml.dump({
-            "lower_rgb": lower_rgb.tolist(),
-            "upper_rgb": upper_rgb.tolist(),
+            "lower": lower.tolist(),
+            "upper": upper.tolist(),
             "invert_frame": config.invert_frame,
             "invert_mask": config.invert_mask
         }, f)
@@ -176,22 +171,27 @@ def get_slider_config(calibration):
         return get_slider_ball_config()
 
 
+def int2bool(x: int) -> bool:
+    return True if x == 1 else False
+
+
 def get_slider_ball_config():
     # get current positions of four trackbars
     invert_frame = cv2.getTrackbarPos('invert_frame', BALL)
     invert_mask = cv2.getTrackbarPos('invert_mask', BALL)
 
-    rl = cv2.getTrackbarPos(slider_label('R', 'low'), BALL)
-    rh = cv2.getTrackbarPos(slider_label('R', 'high'), BALL)
+    hue = cv2.getTrackbarPos('Hue', BALL)
+    hl = max(0, hue - 10)
+    hh = min(179, hue + 10)
 
-    gl = cv2.getTrackbarPos(slider_label('G', 'low'), BALL)
-    gh = cv2.getTrackbarPos(slider_label('G', 'high'), BALL)
+    sl = cv2.getTrackbarPos(slider_label('S', 'low'), BALL)
+    sh = cv2.getTrackbarPos(slider_label('S', 'high'), BALL)
 
-    bl = cv2.getTrackbarPos(slider_label('B', 'low'), BALL)
-    bh = cv2.getTrackbarPos(slider_label('B', 'high'), BALL)
-    lower = rgb2hsv(np.array([rl, gl, bl]))
-    upper = rgb2hsv(np.array([rh, gh, bh]))
-    return BallConfig(bounds_hsv=[lower, upper], invert_mask=invert_mask, invert_frame=invert_frame)
+    vl = cv2.getTrackbarPos(slider_label('V', 'low'), BALL)
+    vh = cv2.getTrackbarPos(slider_label('V', 'high'), BALL)
+    lower = np.array([hl, sl, vl])
+    upper = np.array([hh, sh, vh])
+    return BallConfig(bounds=[lower, upper], invert_mask=int2bool(invert_mask), invert_frame=int2bool(invert_frame))
 
 
 def get_slider_goals_config():
@@ -202,4 +202,4 @@ def get_slider_goals_config():
     lower = cv2.getTrackbarPos('lower', GOAL)
     upper = cv2.getTrackbarPos('upper', GOAL)
 
-    return GoalConfig(bounds=[lower, upper], invert_mask=invert_mask, invert_frame=invert_frame)
+    return GoalConfig(bounds=[lower, upper], invert_mask=int2bool(invert_mask), invert_frame=int2bool(invert_frame))

foosball/tracking/__init__.py

@@ -22,14 +22,14 @@ def yellow_ball() -> BallConfig:
     lower = rgb2hsv(np.array([140, 86, 73]))
     upper = rgb2hsv(np.array([0, 255, 94]))
 
-    return BallConfig(bounds_hsv=[lower, upper], invert_frame=False, invert_mask=False)
+    return BallConfig(bounds=[lower, upper], invert_frame=False, invert_mask=False)
 
 
 def orange_ball() -> BallConfig:
     lower = rgb2hsv(np.array([166, 94, 72]))
     upper = rgb2hsv(np.array([0, 249, 199]))
 
-    return BallConfig(bounds_hsv=[lower, upper], invert_frame=False, invert_mask=False)
+    return BallConfig(bounds=[lower, upper], invert_frame=False, invert_mask=False)
 
 
 def get_goal_config() -> GoalConfig:

foosball/tracking/ai.py

@@ -164,7 +164,7 @@ def adjust_calibration(self):
         # see if some sliders changed
         if self.calibration in ["goal", "ball"]:
             new_config = get_slider_config(self.calibration)
-            if new_config.to_dict() != self.calibration_config().to_dict():
+            if new_config != self.calibration_config():
                 self.set_calibration_config(new_config)
                 self.tracking.config_input(self.calibration_config())
 

foosball/tracking/colordetection.py

@@ -30,7 +30,7 @@ def detect_goals(frame, config: GoalConfig) -> GoalsDetectionResult:
 
 
 def filter_color_range(frame, bounds: BallConfig) -> Frame:
-    [lower, upper] = bounds.bounds("hsv")
+    [lower, upper] = bounds.bounds
     f = frame if not bounds.invert_frame else cv2.bitwise_not(frame)
 
     blurred = cv2.GaussianBlur(f, (1, 1), 0)

foosball/tracking/tracker.py

@@ -57,9 +57,9 @@ def get_info(self, ball_track: Track) -> Info:
             ("Tracker", f"{'off' if self.off else 'on'}")
         ]
         if self.ball_calibration:
-            [lower_rgb, upper_rgb] = self.calibration_bounds().bounds("rgb")
-            info.append((f"lower", f'({",".join(map(str,lower_rgb))})'))
-            info.append((f"upper", f'({",".join(map(str,upper_rgb))})'))
+            [lower, upper] = self.calibration_bounds().bounds
+            info.append((f"lower", f'({",".join(map(str,lower))})'))
+            info.append((f"upper", f'({",".join(map(str,upper))})'))
             info.append((f"invert frame", f'{self.calibration_bounds().invert_frame}'))
             info.append((f"invert mask", f'{self.calibration_bounds().invert_mask}'))
         return info