Saner color calibration. Finder API changed

pykick/colorpicker.py

@@ -7,15 +7,15 @@ import argparse
 import cv2
 
 from .imagereaders import VideoReader, NaoImageReader, PictureReader
-from .finders import GoalFinder
+from .finders import GoalFinder, BallFinder, FieldFinder
-from .utils import read_config, imresize, field_mask
+from .utils import read_config, imresize
 
 class Colorpicker(object):
 
-    WINDOW_CAPTURE_NAME = 'Video Capture'
+    WINDOW_CAPTURE_NAME = 'Object Detection (or not)'
-    WINDOW_DETECTION_NAME = 'Object Detection'
+    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 target == 'goal':
-        if 'goal' in markers:
+            Marker = GoalFinder
-            self.markers['goal'] = 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.marker.hsv_lower = tuple(
-            self.markers[marker].hsv_lower = tuple(
+            map(self.settings.get, ('low_h', 'low_s', 'low_v'))
-                map(self.settings.get, ('low_h', 'low_s', 'low_v'))
+        )
-            )
+        self.marker.hsv_upper = tuple(
-            self.markers[marker].hsv_upper = tuple(
+            map(self.settings.get, ('high_h', 'high_s', 'high_v'))
-                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')))
         )
 
-        for marker in self.markers:
+    def show_frame(self, frame, width=None, manual=False):
-            self.markers[marker].draw(frame)
+        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:
+            {color: [
-             [list(map(self.settings.get, ['low_h', 'low_s', 'low_v'])),
+                list(map(self.settings.get, ['low_h', 'low_s', 'low_v'])),
-              list(map(self.settings.get, ['high_h', 'high_s', 'high_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)
 

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'])
+                             defaults['ball_min_radius'])
-    window_name = 'Live detection'
+    field_finder = FieldFinder(defaults['field'][0], defaults['field'][1])
-    cv2.namedWindow(window_name)
+    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_frame = field_finder.draw(frame, field)
-            ball = ball_finder.find_colored_ball(
+            goal_frame = field_finder.draw(frame, not_field)
-                cv2.bitwise_and(frame, frame, mask=field))
+
-            goal_finder.draw(frame, goal)
+            ball = ball_finder.find(ball_frame)
-            ball_finder.draw(frame, ball)
+            goal = goal_finder.find(goal_frame)
-            cv2.imshow(window_name, 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:

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):
+    def find(self, frame):
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+        thr = self.primary_mask(frame)
-        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)
         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

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
 }

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(
+        ball = self.ball_finder.find(cam.get_frame())
-            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(
+        goal_contour = self.goal_finder.find(self.upper_camera.get_frame())
-            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)

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