Implemented goal detection in colorpicker

pykick/colorpicker.py

@@ -3,10 +3,12 @@ from __future__ import division
 
 import json
 import argparse
+
 import cv2
+import numpy as np
 
 from .imagereaders import VideoReader, NaoImageReader, PictureReader
-from .utils import read_config
+from .utils import read_config, imresize
 
 class Colorpicker(object):
 
@@ -19,7 +21,7 @@ class Colorpicker(object):
         checkers = [
             lambda x: min(x, self.settings['high_h'] - 1),  # LOW H
             lambda x: min(x, self.settings['high_s'] - 1),  # LOW S
-            lambda x: min(x, self.settings['high_v'] - 1),  # LOW H
+            lambda x: min(x, self.settings['high_v'] - 1),  # LOW V
             lambda x: max(x, self.settings['low_h'] + 1),  # HIGH H
             lambda x: max(x, self.settings['low_s'] + 1),  # HIGH S
             lambda x: max(x, self.settings['low_v'] + 1),  # HIGH V
@@ -52,17 +54,75 @@ class Colorpicker(object):
         cv2.setTrackbarPos(name, self.WINDOW_DETECTION_NAME,
                            self.settings[name])
 
-    def show_frame(self, frame, width=None):
+    def goal_similarity(self, contour):
+        contour = contour.reshape((-1, 2))
+        left, right = contour[:,0].min(), contour[:,0].max()
+        top, bottom = contour[:,1].min(), contour[:,1].max()
+        approx_line = np.array([
+            [left, bottom],
+            [left, top],
+            [right, top],
+            [right, bottom]
+        ])
+        shape_sim = np.array([
+            (np.abs(contour - al)).sum(axis=1) for al in approx_line
+        ])
+        len_a = cv2.arcLength(approx_line, False)
+        shape_sim = shape_sim.min(axis=1) / len_a
+
+        shape_sim = shape_sim.sum()
+
+        # len_c = cv2.arcLength(contour, True)
+        area_c = cv2.contourArea(contour)
+
+        # len_similarity = ((len_c / 2 - len_a) / len_a)**2
+        area_sim = area_c / ((right - left) * (bottom - top))
+        print(shape_sim, area_sim)
+        return shape_sim * area_sim
+
+    def draw_contours(self, thr):
+        thr = cv2.erode(thr, None, iterations=2)
+        thr = cv2.dilate(thr, None, iterations=2)
+        cnts, hier = cv2.findContours(thr.copy(), cv2.RETR_EXTERNAL,
+                                            cv2.CHAIN_APPROX_SIMPLE)
+        areas = np.array([cv2.contourArea(cnt) for cnt in cnts])
+        perimeters = np.array([cv2.arcLength(cnt, True) for cnt in cnts])
+        epsilon = 0.04 * perimeters
+
+        top_x = 6
+        if len(areas) > top_x:
+            cnt_ind = np.argpartition(areas, -top_x)[-top_x:]
+            cnts = [cnts[i] for i in cnt_ind]
+
+        perimeters = np.array([cv2.arcLength(cnt, True) for cnt in cnts])
+        epsilon = 0.005 * perimeters
+
+        cnts = [cv2.approxPolyDP(cnt, eps, True)
+                for cnt, eps in zip(cnts, epsilon)]
+
+        good_cnt = [cnt for cnt in cnts if 6 <= cnt.shape[0] <= 9
+                    and not cv2.isContourConvex(cnt)]
+        if good_cnt:
+            good_cnt = [min(good_cnt, key=self.goal_similarity)]
+            # print(good_cnt[0])
+
+        thr = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR)
+        cv2.drawContours(thr, good_cnt, -1, (0, 255, 0), 2)
+        return thr
+
+    def show_frame(self, frame, width=None, draw_contours=False):
         frame_HSV = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
         frame_threshold = cv2.inRange(
             frame_HSV,
             tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))),
             tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v')))
         )
-        if width:
-            sf = width / frame.shape[1]
-            frame = cv2.resize(frame, (0, 0), fx=sf, fy=sf)
-            frame_threshold = cv2.resize(frame_threshold, (0, 0), fx=sf, fy=sf)
+        frame = imresize(frame, width=width)
+        frame_threshold = imresize(frame_threshold, width=width)
+
+        if draw_contours:
+            frame_threshold = self.draw_contours(frame_threshold)
+
         cv2.imshow(self.WINDOW_CAPTURE_NAME, frame)
         cv2.imshow(self.WINDOW_DETECTION_NAME, frame_threshold)
         return cv2.waitKey(1)
@@ -71,7 +131,7 @@ class Colorpicker(object):
         try:
             with open(filename) as f:
                 conf = json.load(f)
-        except FileNotFoundError:
+        except IOError:
             conf = {}
         conf.update(self.settings)
         with open(filename, 'w') as f:
@@ -166,7 +226,7 @@ if __name__ == '__main__':
         while True:
             if not args.still:
                 frame = rdr.get_frame()
-            key = cp.show_frame(frame, args.width)
+            key = cp.show_frame(frame, width=args.width, draw_contours=True)
             if key == ord('q') or key == 27:
                 break
     finally:

pykick/goal_hsv.json

@@ -0,0 +1,8 @@
+{
+    "low_s": 0, 
+    "low_v": 159, 
+    "high_h": 180, 
+    "high_v": 255, 
+    "low_h": 0, 
+    "high_s": 62
+}

pykick/imagereaders.py

@@ -62,9 +62,10 @@ class VideoReader(object):
         if not succ:
             raise ValueError('Error while reading video')
         self.ctr += 1
-        if self.ctr == self.cap.get(cv2.CAP_PROP_FRAME_COUNT) and self.loop:
+        if (self.ctr == self.cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT) and
+            self.loop):
             self.ctr = 0
-            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+            self.cap.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, 0)
         return frame
 
     def close(self):

pykick/utils.py

@@ -1,5 +1,8 @@
+from __future__ import division
+
 import os
 import json
+from cv2 import resize as cv2_resize
 
 
 HERE = os.path.dirname(os.path.realpath(__file__))
@@ -9,3 +12,18 @@ def read_config(cfg_file=os.path.join(HERE, 'nao_defaults.json')):
     with open(cfg_file) as f:
         cfg = json.load(f)
     return cfg
+
+
+def imresize(frame, width=None, height=None):
+    if not width and not height:
+        return frame
+    if not height:
+        sf = width / frame.shape[1]
+        sz = (0, 0)
+    if not width:
+        sf = height / frame.shape[0]
+        sz = (0, 0)
+    if width and height:
+        sf = 0
+        sz = (width, height)
+    return cv2_resize(frame, sz, fx=sf, fy=sf)