From 126329bafd9e42d34f3aeb5aad5fe9c97ae3d909 Mon Sep 17 00:00:00 2001 From: Pavel Lutskov Date: Sat, 23 Jun 2018 13:08:12 +0200 Subject: [PATCH] Saner color calibration. Finder API changed --- pykick/colorpicker.py | 80 +++++++++++++++++++++------------------- pykick/detection_demo.py | 34 ++++++++++------- pykick/finders.py | 57 +++++++++++++++++++++++----- pykick/nao_defaults.json | 34 ++++++++--------- pykick/striker.py | 8 +--- pykick/utils.py | 20 ---------- 6 files changed, 129 insertions(+), 104 deletions(-) diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index ec0314a..eaa15ba 100644 --- a/pykick/colorpicker.py +++ b/pykick/colorpicker.py @@ -7,15 +7,15 @@ import argparse import cv2 from .imagereaders import VideoReader, NaoImageReader, PictureReader -from .finders import GoalFinder -from .utils import read_config, imresize, field_mask +from .finders import GoalFinder, BallFinder, FieldFinder +from .utils import read_config, imresize class Colorpicker(object): - WINDOW_CAPTURE_NAME = 'Video Capture' - WINDOW_DETECTION_NAME = 'Object Detection' + WINDOW_CAPTURE_NAME = 'Object Detection (or not)' + WINDOW_DETECTION_NAME = 'Primary Mask' - def __init__(self, markers=()): + def __init__(self, target=None): parameters = ['low_h', 'low_s', 'low_v', 'high_h', 'high_s', 'high_v'] maxes = [180, 255, 255, 180, 255, 255] checkers = [ @@ -34,12 +34,22 @@ class Colorpicker(object): 'high_s': 255, 'high_v': 255 } - self.markers = {} - if 'goal' in markers: - self.markers['goal'] = GoalFinder( + if target == 'goal': + Marker = GoalFinder + elif target == 'ball': + Marker = BallFinder + elif target == 'field': + Marker = FieldFinder + else: + Marker = None + + if Marker is not None: + self.marker = Marker( tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))), tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v'))) ) + else: + self.marker = None cv2.namedWindow(self.WINDOW_CAPTURE_NAME) cv2.namedWindow(self.WINDOW_DETECTION_NAME) @@ -64,35 +74,29 @@ class Colorpicker(object): def _hsv_updated(self, param): cv2.setTrackbarPos(param, self.WINDOW_DETECTION_NAME, self.settings[param]) - for marker in self.markers: - self.markers[marker].hsv_lower = tuple( - map(self.settings.get, ('low_h', 'low_s', 'low_v')) - ) - self.markers[marker].hsv_upper = tuple( - map(self.settings.get, ('high_h', 'high_s', 'high_v')) - ) - - def show_frame(self, frame, width=None, manual=False): - # hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - # frame_threshold = cv2.inRange( - # hsv, - # tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))), - # tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v'))) - # ) - frame = imresize(frame, width=width) - # frame_threshold = imresize(frame_threshold, width=width) - - frame_threshold = field_mask( - frame, - tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))), - tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v'))) + self.marker.hsv_lower = tuple( + map(self.settings.get, ('low_h', 'low_s', 'low_v')) + ) + self.marker.hsv_upper = tuple( + map(self.settings.get, ('high_h', 'high_s', 'high_v')) ) - for marker in self.markers: - self.markers[marker].draw(frame) + def show_frame(self, frame, width=None, manual=False): + frame = imresize(frame, width=width) + if self.marker is not None: + thr = self.marker.primary_mask(frame) + stuff = self.marker.find(frame) + frame = self.marker.draw(frame, stuff) + else: + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + thr = cv2.inRange( + hsv, + tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))), + tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v'))) + ) cv2.imshow(self.WINDOW_CAPTURE_NAME, frame) - cv2.imshow(self.WINDOW_DETECTION_NAME, frame_threshold) + cv2.imshow(self.WINDOW_DETECTION_NAME, thr) return cv2.waitKey(0 if manual else 1) def save(self, filename, color): @@ -102,10 +106,10 @@ class Colorpicker(object): except IOError: conf = {} conf.update( - {color: - [list(map(self.settings.get, ['low_h', 'low_s', 'low_v'])), - list(map(self.settings.get, ['high_h', 'high_s', 'high_v']))] - } + {color: [ + list(map(self.settings.get, ['low_h', 'low_s', 'low_v'])), + list(map(self.settings.get, ['high_h', 'high_s', 'high_v'])) + ]} ) with open(filename, 'w') as f: json.dump(conf, f, indent=4) @@ -191,7 +195,7 @@ if __name__ == '__main__': ) args = parser.parse_args() - cp = Colorpicker() + cp = Colorpicker(args.target) if args.input_config: cp.load(args.input_config, args.target) diff --git a/pykick/detection_demo.py b/pykick/detection_demo.py index a74c829..e4220ce 100644 --- a/pykick/detection_demo.py +++ b/pykick/detection_demo.py @@ -7,7 +7,7 @@ import cv2 from .utils import read_config, imresize, field_mask from .imagereaders import NaoImageReader, VideoReader, PictureReader -from .finders import BallFinder, GoalFinder +from .finders import BallFinder, GoalFinder, FieldFinder if __name__ == '__main__': @@ -83,9 +83,12 @@ if __name__ == '__main__': goal_finder = GoalFinder(defaults['goal'][0], defaults['goal'][1]) ball_finder = BallFinder(defaults['ball'][0], defaults['ball'][1], - defaults['min_radius']) - window_name = 'Live detection' - cv2.namedWindow(window_name) + defaults['ball_min_radius']) + field_finder = FieldFinder(defaults['field'][0], defaults['field'][1]) + ball_window = 'Ball detection' + goal_window = 'Goal detection' + cv2.namedWindow(ball_window) + cv2.namedWindow(goal_window) try: if args.still: @@ -95,16 +98,21 @@ if __name__ == '__main__': if not args.still: frame = rdr.get_frame() frame = imresize(frame, width=args.width) - field_hsv = defaults['field'] - field = field_mask(frame, field_hsv[0], field_hsv[1]) + + field = field_finder.find(frame) not_field = cv2.bitwise_not(field) - goal = goal_finder.find_goal_contour( - cv2.bitwise_and(frame, frame, mask=not_field)) - ball = ball_finder.find_colored_ball( - cv2.bitwise_and(frame, frame, mask=field)) - goal_finder.draw(frame, goal) - ball_finder.draw(frame, ball) - cv2.imshow(window_name, frame) + + ball_frame = field_finder.draw(frame, field) + goal_frame = field_finder.draw(frame, not_field) + + ball = ball_finder.find(ball_frame) + goal = goal_finder.find(goal_frame) + + ball_frame = ball_finder.draw(ball_frame, ball) + goal_frame = goal_finder.draw(goal_frame, goal) + + cv2.imshow(ball_window, ball_frame) + cv2.imshow(goal_window, goal_frame) key = cv2.waitKey(0 if args.manual else 1) if key == ord('q') or key == 27: diff --git a/pykick/finders.py b/pykick/finders.py index a7e2f96..08e7424 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -7,13 +7,51 @@ import cv2 import numpy as np +class FieldFinder(object): + + def __init__(self, hsv_lower, hsv_upper): + self.hsv_lower = tuple(hsv_lower) + self.hsv_upper = tuple(hsv_upper) + + def primary_mask(self, frame): + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + blurred = cv2.GaussianBlur(hsv, (25, 25), 20) + thr = cv2.inRange(blurred, tuple(self.hsv_lower), tuple(self.hsv_upper)) + thr = cv2.erode(thr, None, iterations=6) + thr = cv2.dilate(thr, None, iterations=10) + return thr + + def find(self, frame): + thr = self.primary_mask(frame) + cnts, _ = cv2.findContours(thr.copy(), cv2.RETR_EXTERNAL, + cv2.CHAIN_APPROX_SIMPLE) + if not cnts: + return None + field = max(cnts, key=cv2.contourArea) + field = cv2.convexHull(field) + mask = np.zeros(thr.shape, dtype=np.uint8) + cv2.drawContours(mask, (field,), -1, 255, -1) + return mask + + def draw(self, frame, field): + if field is not None: + frame = cv2.bitwise_and(frame, frame, mask=field) + return frame + + class GoalFinder(object): def __init__(self, hsv_lower, hsv_upper): - self.hsv_lower = tuple(hsv_lower) self.hsv_upper = tuple(hsv_upper) + def primary_mask(self, frame): + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + thr = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) + thr = cv2.erode(thr, None, iterations=2) + thr = cv2.dilate(thr, None, iterations=2) + return thr + def goal_similarity(self, contour): contour = contour.squeeze(axis=1) hull = cv2.convexHull(contour).squeeze(axis=1) @@ -36,13 +74,8 @@ class GoalFinder(object): print('Goal candidate:', shape_sim, area_sim, final_score) return final_score - def find_goal_contour(self, frame): - hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - thr = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) - - # The ususal - thr = cv2.erode(thr, None, iterations=2) - thr = cv2.dilate(thr, None, iterations=2) + def find(self, frame): + thr = self.primary_mask(frame) cnts, _ = cv2.findContours(thr, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) areas = np.array([cv2.contourArea(cnt) for cnt in cnts]) @@ -70,7 +103,7 @@ class GoalFinder(object): best = min(similarities) print('Final goal score:', best) print() - if best > 0.35: + if best > 0.45: return None # Find the contour with the shape closest to that of the goal goal = good_cnts[similarities.index(best)] @@ -85,7 +118,9 @@ class GoalFinder(object): def draw(self, frame, goal): if goal is not None: + frame = frame.copy() cv2.drawContours(frame, (goal,), -1, (0, 255, 0), 2) + return frame class BallFinder(object): @@ -97,7 +132,7 @@ class BallFinder(object): self.min_radius = min_radius self.history = deque(maxlen=64) - def find_colored_ball(self, frame): + def find(self, frame): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) # construct a mask for the color, then perform a series of @@ -139,8 +174,10 @@ class BallFinder(object): def draw(self, frame, ball): if ball is not None: + frame = frame.copy() center, radius = ball cv2.circle(frame, center, radius, (255, 255, 0), 1) + return frame # cv2.circle(frame, center, 5, (0, 255, 0), -1) # loop over the set of tracked points diff --git a/pykick/nao_defaults.json b/pykick/nao_defaults.json index 1564dc8..a0ec935 100644 --- a/pykick/nao_defaults.json +++ b/pykick/nao_defaults.json @@ -11,34 +11,34 @@ 255 ] ], - "res": 1, + "cam": 1, "goal": [ [ 0, 0, - 159 + 114 ], [ 180, - 62, + 49, + 255 + ] + ], + "res": 1, + "ball_min_radius": 0.01, + "field": [ + [ + 19, + 57, + 84 + ], + [ + 65, + 255, 255 ] ], "fps": 30, "ip": "192.168.0.11", - "field": [ - [ - 11, - 74, - 28 - ], - [ - 69, - 255, - 255 - ] - ], - "cam": 1, - "ball_min_radius": 0.01, "port": 9559 } \ No newline at end of file diff --git a/pykick/striker.py b/pykick/striker.py index fd02af7..b78cb31 100644 --- a/pykick/striker.py +++ b/pykick/striker.py @@ -68,9 +68,7 @@ class Striker(object): def get_ball_angles_from_camera(self, cam): """Detect the ball and return its angles in camera coordinates.""" - ball = self.ball_finder.find_colored_ball( - cam.get_frame() - ) + ball = self.ball_finder.find(cam.get_frame()) if ball is None: return None @@ -194,9 +192,7 @@ class Striker(object): self.mover.wait() return False - goal_contour = self.goal_finder.find_goal_contour( - self.upper_camera.get_frame() - ) + goal_contour = self.goal_finder.find(self.upper_camera.get_frame()) if goal_contour is not None: goal_center_x = self.goal_finder.goal_center(goal_contour) gcx_rel, _ = self.upper_camera.to_relative(goal_center_x, 0) diff --git a/pykick/utils.py b/pykick/utils.py index 4b9fe06..7ee628b 100644 --- a/pykick/utils.py +++ b/pykick/utils.py @@ -29,23 +29,3 @@ def imresize(frame, width=None, height=None): sf = 0 sz = (width, height) return cv2.resize(frame, sz, fx=sf, fy=sf) - - -def field_mask(frame, hsv_lower, hsv_upper): - hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - blurred = cv2.GaussianBlur(hsv, (25, 25), 20) - thr = cv2.inRange(blurred, tuple(hsv_lower), tuple(hsv_upper)) - thr = cv2.erode(thr, None, iterations=4) - thr = cv2.dilate(thr, None, iterations=8) - cnts, _ = cv2.findContours(thr.copy(), cv2.RETR_EXTERNAL, - cv2.CHAIN_APPROX_SIMPLE) - field = max(cnts, key=cv2.contourArea) - field = cv2.convexHull(field) - # print(field) - mask = np.zeros(thr.shape, dtype=np.uint8) - print(mask.dtype) - thr = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR) - cv2.drawContours(mask, (field,), -1, 255, -1) - - # The ususal - return mask