diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index d2fd4dd..ec0314a 100644 --- a/pykick/colorpicker.py +++ b/pykick/colorpicker.py @@ -8,7 +8,7 @@ import cv2 from .imagereaders import VideoReader, NaoImageReader, PictureReader from .finders import GoalFinder -from .utils import read_config, imresize +from .utils import read_config, imresize, field_mask class Colorpicker(object): @@ -72,25 +72,28 @@ class Colorpicker(object): map(self.settings.get, ('high_h', 'high_s', 'high_v')) ) - def show_frame(self, frame, width=None): - hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - frame_threshold = cv2.inRange( - hsv, + 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'))) ) - frame = imresize(frame, width=width) - frame_threshold = imresize(frame_threshold, width=width) - if 'goal' in self.markers: - self.markers['goal'].draw(frame) - - if 'ball' in self.markers: - self.markers['ball'].draw(frame) + for marker in self.markers: + self.markers[marker].draw(frame) cv2.imshow(self.WINDOW_CAPTURE_NAME, frame) cv2.imshow(self.WINDOW_DETECTION_NAME, frame_threshold) - return cv2.waitKey(1) + return cv2.waitKey(0 if manual else 1) def save(self, filename, color): try: @@ -154,6 +157,13 @@ if __name__ == '__main__': help='only take one image from video stream', action='store_true' ) + parser.add_argument( + '--manual', + help='switch frames manually', + action='store_true', + default=False + ) + parser.add_argument( '--nao-cam', choices=[0, 1], @@ -175,15 +185,15 @@ if __name__ == '__main__': default=640 ) parser.add_argument( - '--color', - help='specify which color is being calibrated', - default='white' + '--target', + help='specify for what target is being calibrated', + default='field' ) args = parser.parse_args() cp = Colorpicker() if args.input_config: - cp.load(args.input_config, args.color) + cp.load(args.input_config, args.target) if args.video_file: rdr = VideoReader(args.video_file, loop=True) @@ -205,12 +215,12 @@ if __name__ == '__main__': while True: if not args.still: frame = rdr.get_frame() - key = cp.show_frame(frame, width=args.width) + key = cp.show_frame(frame, width=args.width, manual=args.manual) if key == ord('q') or key == 27: break finally: cp.do_print() if args.output_config: - cp.save(args.output_config, args.color) + cp.save(args.output_config, args.target) if not args.still: rdr.close() diff --git a/pykick/detection_demo.py b/pykick/detection_demo.py index f8fbfc6..306a153 100644 --- a/pykick/detection_demo.py +++ b/pykick/detection_demo.py @@ -1,21 +1,115 @@ from __future__ import print_function from __future__ import division -from .utils import read_config -from .imagereaders import NaoImageReader, VideoReader -from .finders import BallFinder +import argparse + +import cv2 + +from .utils import read_config, imresize, field_mask +from .imagereaders import NaoImageReader, VideoReader, PictureReader +from .finders import BallFinder, GoalFinder if __name__ == '__main__': - video = VideoReader(0, loop=True) - cfg = read_config() - hsv_lower = cfg['red'][0] - hsv_upper = cfg['red'][1] - finder = BallFinder(hsv_lower, hsv_upper, cfg['min_radius'], None) + defaults = read_config() + parser = argparse.ArgumentParser() + # parser.add_argument( + # '-i', '--input-config', + # help='file, from which to read the initial values', + # required=True + # ) + imsource = parser.add_mutually_exclusive_group() + imsource.add_argument( + '--video-file', + help='video file to use' + ) + imsource.add_argument( + '--image-file', + help='image to use' + ) + imsource.add_argument( + '--nao-ip', + help='ip address of the nao robot, from which to capture', + default=False, + const=defaults['ip'], + nargs='?' + ) + parser.add_argument( + '--still', + help='only take one image from video stream', + action='store_true' + ) + parser.add_argument( + '--manual', + help='switch frames manually', + action='store_true', + default=False + ) + + parser.add_argument( + '--nao-cam', + choices=[0, 1], + type=int, + help='0 for top camera, 1 for bottom camera', + default=defaults['cam'] + ) + parser.add_argument( + '--nao-res', + choices=[1, 2, 3], + help='choose a nao resolution', + type=int, + default=defaults['res'] + ) + parser.add_argument( + '--width', + help='specify width of the image output', + type=int, + default=640 + ) + args = parser.parse_args() + + if args.video_file: + rdr = VideoReader(args.video_file, loop=True) + elif args.image_file: + rdr = PictureReader(args.image_file) + elif args.nao_ip: + rdr = NaoImageReader( + args.nao_ip, + cam_id=args.nao_cam, + res=args.nao_res + ) + else: + rdr = VideoReader(0) + + 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) + try: + if args.still: + frame = rdr.get_frame() + rdr.close() while True: - frame = video.get_frame() - finder.find_colored_ball(frame) - finder.visualize(frame) + 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]) + 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, + cv2.bitwise_and(frame, frame, mask=field)) + + key = cv2.waitKey(0 if args.manual else 1) + if key == ord('q') or key == 27: + break finally: - video.close() + if not args.still: + rdr.close() diff --git a/pykick/finders.py b/pykick/finders.py index b8ade36..a7e2f96 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -11,8 +11,8 @@ class GoalFinder(object): def __init__(self, hsv_lower, hsv_upper): - self.hsv_lower = hsv_lower - self.hsv_upper = hsv_upper + self.hsv_lower = tuple(hsv_lower) + self.hsv_upper = tuple(hsv_upper) def goal_similarity(self, contour): contour = contour.squeeze(axis=1) @@ -33,7 +33,7 @@ class GoalFinder(object): # Final similarity score is just the sum of both final_score = shape_sim + area_sim - print('Candidate:', shape_sim, area_sim, final_score) + print('Goal candidate:', shape_sim, area_sim, final_score) return final_score def find_goal_contour(self, frame): @@ -68,7 +68,7 @@ class GoalFinder(object): similarities = [self.goal_similarity(cnt) for cnt in good_cnts] best = min(similarities) - print('Final score:', best) + print('Final goal score:', best) print() if best > 0.35: return None @@ -83,8 +83,7 @@ class GoalFinder(object): l, r = self.left_right_post(self, contour) return (l + r) / 2 - def draw(self, frame): - goal = self.find_goal_contour(frame) + def draw(self, frame, goal): if goal is not None: cv2.drawContours(frame, (goal,), -1, (0, 255, 0), 2) @@ -93,8 +92,8 @@ class BallFinder(object): def __init__(self, hsv_lower, hsv_upper, min_radius): - self.hsv_lower = hsv_lower - self.hsv_upper = hsv_upper + self.hsv_lower = tuple(hsv_lower) + self.hsv_upper = tuple(hsv_upper) self.min_radius = min_radius self.history = deque(maxlen=64) @@ -102,55 +101,56 @@ class BallFinder(object): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) # construct a mask for the color, then perform a series of - # dilations and erosions to remove any small blobs left in the mask + # dilations and erosions to remove any small blobs left in the mask ? mask = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) - mask = cv2.erode(mask, None, iterations=2) - mask = cv2.dilate(mask, None, iterations=2) + # mask = cv2.erode(mask, None, iterations=2) + # mask = cv2.dilate(mask, None, iterations=2) - # find contours in the mask and initialize the current - # (x, y) center of the ball cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2] - # only proceed if at least one contour was found if len(cnts) == 0: + print('No red contours') + self.history.appendleft(None) return None # find the largest contour in the mask, then use it to compute # the minimum enclosing circle and centroid c = max(cnts, key=cv2.contourArea) - ((x, y), radius) = cv2.minEnclosingCircle(c) + (x, y), radius = cv2.minEnclosingCircle(c) - if radius < self.min_radius: + min_radius_abs = self.min_radius * frame.shape[0] + + if radius < min_radius_abs: + print('Radius:', radius, 'Min radius:', min_radius_abs) + self.history.appendleft(None) return None M = cv2.moments(c) - center = (int(M["m10"] / M["m00"]),int(M["m01"] // M["m00"])) + try: + center = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]) + except ZeroDivisionError: + # It's weird but happened yeah + self.history.append(None) + return None + self.history.appendleft((center, int(radius))) + print('Ball:', center, radius) return center, int(radius) - def draw(self, frame): - ball = self.find_colored_ball(frame) - self.history.appendleft(ball) - + def draw(self, frame, ball): if ball is not None: center, radius = ball cv2.circle(frame, center, radius, (255, 255, 0), 1) - cv2.circle(frame, center, 5, (0, 255, 0), -1) + # cv2.circle(frame, center, 5, (0, 255, 0), -1) # loop over the set of tracked points - for i in range(1, len(self.history)): + # for i in range(1, len(self.history)): # if either of the tracked points are None, ignore them - if self.history[i - 1] is None or self.history[i] is None: - continue + # if self.history[i - 1] is None or self.history[i] is None: + # continue # otherwise, compute the thickness of the line and # draw the connecting lines - center_now = self.history[i - 1][0] - center_prev = self.history[i][0] - thickness = int((64 / (i + 1))**0.5 * 2.5) - cv2.line(frame, center_now, center_prev, (0, 255, 0), thickness) - - # def load_hsv_config(self, filename): - # with open(filename) as f: - # hsv = json.load(f) - # self.hsv_lower = tuple(map(hsv.get, ('low_h', 'low_s', 'low_v'))) - # self.hsv_upper = tuple(map(hsv.get, ('high_h', 'high_s', 'high_v'))) + # center_now = self.history[i - 1][0] + # center_prev = self.history[i][0] + # thickness = int((64 / (i + 1))**0.5 * 2.5) + # cv2.line(frame, center_now, center_prev, (0, 255, 0), thickness) diff --git a/pykick/nao_defaults.json b/pykick/nao_defaults.json index 4809179..880f310 100644 --- a/pykick/nao_defaults.json +++ b/pykick/nao_defaults.json @@ -1,10 +1,18 @@ { - "cam": 1, - "ip": "192.168.0.11", - "min_radius": 5, - "fps": 30, + "ball": [ + [ + 0, + 175, + 100 + ], + [ + 6, + 255, + 255 + ] + ], "res": 1, - "white": [ + "goal": [ [ 0, 0, @@ -16,17 +24,21 @@ 255 ] ], - "port": 9559, - "red": [ + "fps": 30, + "ip": "192.168.0.11", + "field": [ [ - 0, - 175, - 100 + 11, + 74, + 28 ], [ - 6, + 69, 255, 255 ] - ] -} + ], + "cam": 1, + "min_radius": 0.01, + "port": 9559 +} \ No newline at end of file diff --git a/pykick/striker.py b/pykick/striker.py index 962cb70..b00fca2 100644 --- a/pykick/striker.py +++ b/pykick/striker.py @@ -12,17 +12,17 @@ from .movements import NaoMover class Striker(object): - def __init__(self, nao_ip, nao_port, res, red_hsv, white_hsv, - min_radius, run_after): + def __init__(self, nao_ip, nao_port, res, ball_hsv, goal_hsv, + ball_min_radius, run_after): self.mover = NaoMover(nao_ip=nao_ip, nao_port=nao_port) self.mover.stand_up() self.video_top = NaoImageReader(nao_ip, port=nao_port, res=res, fps=30, cam_id=0) self.video_bot = NaoImageReader(nao_ip, port=nao_port, res=res, fps=30, cam_id=1) - self.ball_finder = BallFinder(tuple(red_hsv[0]), tuple(red_hsv[1]), - min_radius) - self.goal_finder = GoalFinder(tuple(white_hsv[0]), tuple(white_hsv[1])) + self.ball_finder = BallFinder(tuple(ball_hsv[0]), tuple(ball_hsv[1]), + ball_min_radius) + self.goal_finder = GoalFinder(tuple(goal_hsv[0]), tuple(goal_hsv[1])) self.lock_counter = 0 self.loss_counter = 0 self.run_after = run_after @@ -214,9 +214,9 @@ if __name__ == '__main__': nao_ip=cfg['ip'], nao_port=cfg['port'], res=cfg['res'], - red_hsv=cfg['red'], - white_hsv=cfg['white'], - min_radius=cfg['min_radius'], + ball_hsv=cfg['ball'], + goal_hsv=cfg['goal'], + ball_min_radius=cfg['ball_min_radius'], run_after=False ) try: diff --git a/pykick/utils.py b/pykick/utils.py index 0278637..a7cc1a8 100644 --- a/pykick/utils.py +++ b/pykick/utils.py @@ -2,7 +2,8 @@ from __future__ import division import os import json -from cv2 import resize as cv2_resize + +import cv2 HERE = os.path.dirname(os.path.realpath(__file__)) @@ -26,4 +27,15 @@ def imresize(frame, width=None, height=None): if width and height: sf = 0 sz = (width, height) - return cv2_resize(frame, sz, fx=sf, fy=sf) + 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)) + + # The ususal + thr = cv2.erode(thr, None, iterations=6) + thr = cv2.dilate(thr, None, iterations=20) + return thr