guess what? - refactoring! also masking out field

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):
+    def show_frame(self, frame, width=None, manual=False):
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+        # hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-        frame_threshold = cv2.inRange(
+        # frame_threshold = cv2.inRange(
-            hsv,
+            # 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:
+        for marker in self.markers:
-            self.markers['goal'].draw(frame)
+            self.markers[marker].draw(frame)
-
-        if 'ball' in self.markers:
-            self.markers['ball'].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',
+        '--target',
-        help='specify which color is being calibrated',
+        help='specify for what target is being calibrated',
-        default='white'
+        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()

pykick/detection_demo.py

@@ -1,21 +1,115 @@
 from __future__ import print_function
 from __future__ import division
 
-from .utils import read_config
+import argparse
-from .imagereaders import NaoImageReader, VideoReader
+
-from .finders import BallFinder
+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)
+    defaults = read_config()
-    cfg = read_config()
+    parser = argparse.ArgumentParser()
-    hsv_lower = cfg['red'][0]
+    # parser.add_argument(
-    hsv_upper = cfg['red'][1]
+        # '-i', '--input-config',
-    finder = BallFinder(hsv_lower, hsv_upper, cfg['min_radius'], None)
+        # 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()
+            if not args.still:
-            finder.find_colored_ball(frame)
+                frame = rdr.get_frame()
-            finder.visualize(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()

pykick/finders.py

@@ -11,8 +11,8 @@ class GoalFinder(object):
 
     def __init__(self, hsv_lower, hsv_upper):
 
-        self.hsv_lower = hsv_lower
+        self.hsv_lower = tuple(hsv_lower)
-        self.hsv_upper = hsv_upper
+        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):
+    def draw(self, frame, goal):
-        goal = self.find_goal_contour(frame)
         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_lower = tuple(hsv_lower)
-        self.hsv_upper = hsv_upper
+        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.erode(mask, None, iterations=2)
-        mask = cv2.dilate(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):
+    def draw(self, frame, ball):
-        ball = self.find_colored_ball(frame)
-        self.history.appendleft(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:
+            # if self.history[i - 1] is None or self.history[i] is None:
-                continue
+                # continue
             # otherwise, compute the thickness of the line and
             # draw the connecting lines
-            center_now = self.history[i - 1][0]
+            # center_now = self.history[i - 1][0]
-            center_prev = self.history[i][0]
+            # center_prev = self.history[i][0]
-            thickness = int((64 / (i + 1))**0.5 * 2.5)
+            # thickness = int((64 / (i + 1))**0.5 * 2.5)
-            cv2.line(frame, center_now, center_prev, (0, 255, 0), thickness)
+            # 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')))

pykick/nao_defaults.json

@@ -1,10 +1,18 @@
 {
-    "cam": 1, 
+    "ball": [
-    "ip": "192.168.0.11", 
+        [
-    "min_radius": 5, 
+            0, 
-    "fps": 30, 
+            175, 
+            100
+        ], 
+        [
+            6, 
+            255, 
+            255
+        ]
+    ], 
     "res": 1, 
-    "white": [
+    "goal": [
         [
             0, 
             0, 
@@ -16,17 +24,21 @@
             255
         ]
     ], 
-    "port": 9559, 
+    "fps": 30, 
-    "red": [
+    "ip": "192.168.0.11", 
+    "field": [
         [
-            0, 
+            11, 
-            175, 
+            74, 
-            100
+            28
         ], 
         [
-            6, 
+            69, 
             255, 
             255
         ]
-    ]
+    ], 
+    "cam": 1, 
+    "min_radius": 0.01, 
+    "port": 9559
 }

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,
+    def __init__(self, nao_ip, nao_port, res, ball_hsv, goal_hsv,
-                 min_radius, run_after):
+                 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]),
+        self.ball_finder = BallFinder(tuple(ball_hsv[0]), tuple(ball_hsv[1]),
-                                      min_radius)
+                                      ball_min_radius)
-        self.goal_finder = GoalFinder(tuple(white_hsv[0]), tuple(white_hsv[1]))
+        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'],
+        ball_hsv=cfg['ball'],
-        white_hsv=cfg['white'],
+        goal_hsv=cfg['goal'],
-        min_radius=cfg['min_radius'],
+        ball_min_radius=cfg['ball_min_radius'],
         run_after=False
     )
     try:

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