From abc69fb69aacf3d4fd6958af54410cda66907243 Mon Sep 17 00:00:00 2001 From: Pavel Lutskov Date: Sun, 10 Jun 2018 21:04:20 +0200 Subject: [PATCH] Implemented goal detection in colorpicker --- pykick/colorpicker.py | 78 +++++++++++++++++++++++++++++++++++++----- pykick/goal_hsv.json | 8 +++++ pykick/imagereaders.py | 5 +-- pykick/utils.py | 18 ++++++++++ 4 files changed, 98 insertions(+), 11 deletions(-) create mode 100644 pykick/goal_hsv.json diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index 7bb6de3..a300815 100644 --- a/pykick/colorpicker.py +++ b/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: diff --git a/pykick/goal_hsv.json b/pykick/goal_hsv.json new file mode 100644 index 0000000..5614b77 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/pykick/imagereaders.py b/pykick/imagereaders.py index 8ac0f2a..1f67d6f 100644 --- a/pykick/imagereaders.py +++ b/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): diff --git a/pykick/utils.py b/pykick/utils.py index 0da75a9..0278637 100644 --- a/pykick/utils.py +++ b/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)