diff --git a/ball.yaml b/ball.yaml index 595f0ec..04ebc6d 100644 --- a/ball.yaml +++ b/ball.yaml @@ -1,8 +1,8 @@ -bounds_hsv: -- - 6 - - 122 - - 140 -- - 71 +bounds: +- - 17 + - 37 + - 25 +- - 37 - 255 - 255 invert_frame: false diff --git a/foosball/models.py b/foosball/models.py index 37647d5..5f2103a 100644 --- a/foosball/models.py +++ b/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, diff --git a/foosball/sink/opencv.py b/foosball/sink/opencv.py index d3851b9..de56c35 100644 --- a/foosball/sink/opencv.py +++ b/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)) diff --git a/foosball/tracking/__init__.py b/foosball/tracking/__init__.py index 8270217..2c1ca4d 100644 --- a/foosball/tracking/__init__.py +++ b/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: diff --git a/foosball/tracking/ai.py b/foosball/tracking/ai.py index b53f20b..dcf77ee 100644 --- a/foosball/tracking/ai.py +++ b/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()) diff --git a/foosball/tracking/colordetection.py b/foosball/tracking/colordetection.py index 4a51784..b3412c3 100644 --- a/foosball/tracking/colordetection.py +++ b/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) diff --git a/foosball/tracking/tracker.py b/foosball/tracking/tracker.py index 088e72e..b686738 100644 --- a/foosball/tracking/tracker.py +++ b/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