From 5076f87145dda4880d231e42a953186d636d1025 Mon Sep 17 00:00:00 2001 From: Pavel Lutskov Date: Wed, 13 Jun 2018 22:11:58 +0200 Subject: [PATCH] did small refactoring, probably broke a lot --- pykick/colorpicker.py | 105 ++++++++++++++--------------------------- pykick/finders.py | 54 +++++++++++---------- pykick/imagereaders.py | 3 +- 3 files changed, 67 insertions(+), 95 deletions(-) diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index 1565fff..a79f7dd 100644 --- a/pykick/colorpicker.py +++ b/pykick/colorpicker.py @@ -5,9 +5,9 @@ import json import argparse import cv2 -import numpy as np from .imagereaders import VideoReader, NaoImageReader, PictureReader +from .finders import GoalFinder from .utils import read_config, imresize class Colorpicker(object): @@ -15,7 +15,7 @@ class Colorpicker(object): WINDOW_CAPTURE_NAME = 'Video Capture' WINDOW_DETECTION_NAME = 'Object Detection' - def __init__(self): + def __init__(self, markers=()): parameters = ['low_h', 'low_s', 'low_v', 'high_h', 'high_s', 'high_v'] maxes = [180, 255, 255, 180, 255, 255] checkers = [ @@ -34,6 +34,13 @@ class Colorpicker(object): 'high_s': 255, 'high_v': 255 } + self.markers = {} + if 'goal' in markers: + self.markers['goal'] = GoalFinder( + tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))), + tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v'))) + ) + cv2.namedWindow(self.WINDOW_CAPTURE_NAME) cv2.namedWindow(self.WINDOW_DETECTION_NAME) self.trackers = [ @@ -51,79 +58,37 @@ class Colorpicker(object): def _on_trackbar(self, val, name, checker): self.settings[name] = checker(val) - cv2.setTrackbarPos(name, self.WINDOW_DETECTION_NAME, - self.settings[name]) + self._hsv_updated(name) - def goal_similarity(self, contour): - contour = contour.reshape((-1, 2)) - hull = cv2.convexHull(contour).reshape((-1, 2)) - len_h = cv2.arcLength(hull, True) + def _hsv_updated(self, param): + cv2.setTrackbarPos(param, self.WINDOW_DETECTION_NAME, + self.settings[param]) + print(param) + for marker in self.markers: + print(self.markers[marker]) + print(self.settings) + self.markers[marker].hsv_lower = tuple( + map(self.settings.get, ('low_h', 'low_s', 'low_v')) + ) + self.markers[marker].hsv_upper = tuple( + map(self.settings.get, ('high_h', 'high_s', 'high_v')) + ) - # Wild assumption that the goal should lie close to its - # enclosing convex hull - shape_sim = np.linalg.norm(contour[:,None] - hull, - axis=2).min(axis=1).sum() / len_h - - # Wild assumption that the area of the goal is rather small - # compared to its enclosing convex hull - area_c = cv2.contourArea(contour) - area_h = cv2.contourArea(hull) - - area_sim = area_c / area_h - - # Final similarity score is just the sum of both - final_score = shape_sim + area_sim - print(shape_sim, area_sim, final_score) - return final_score - - def draw_contours(self, thr): - # The ususal - thr = cv2.erode(thr, None, iterations=2) - thr = cv2.dilate(thr, None, iterations=2) - cnts, _ = 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 - - # Candidates are at most 6 biggest white areas - 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.01 * perimeters - - # Approximate resulting contours with simpler lines - cnts = [cv2.approxPolyDP(cnt, eps, True) - for cnt, eps in zip(cnts, epsilon)] - - # Goal needs normally 8 points for perfect approximation - # But with 6 can also be approximated - good_cnts = [cnt for cnt in cnts if 6 <= cnt.shape[0] <= 9 - and not cv2.isContourConvex(cnt)] - - if good_cnts: - # Find the contour with the shape closest to that of the goal - good_cnts = [min(good_cnts, key=self.goal_similarity)] - - thr = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR) - cv2.drawContours(thr, good_cnts, -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) + def show_frame(self, frame, width=None): + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) frame_threshold = cv2.inRange( - frame_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) - if draw_contours: - frame_threshold = self.draw_contours(frame_threshold) + if 'goal' in self.markers: + self.markers['goal'].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) @@ -143,8 +108,7 @@ class Colorpicker(object): with open(filename) as f: self.settings = json.load(f) for name in self.settings: - cv2.setTrackbarPos(name, self.WINDOW_DETECTION_NAME, - self.settings[name]) + self._hsv_updated(name) if __name__ == '__main__': @@ -186,6 +150,7 @@ if __name__ == '__main__': parser.add_argument( '--nao-cam', choices=[0, 1], + type=int, help='0 for top camera, 1 for bottom camera', default=defaults['cam'] ) @@ -204,7 +169,7 @@ if __name__ == '__main__': ) args = parser.parse_args() - cp = Colorpicker() + cp = Colorpicker(['goal']) if args.input_config: cp.load(args.input_config) @@ -228,7 +193,7 @@ if __name__ == '__main__': while True: if not args.still: frame = rdr.get_frame() - key = cp.show_frame(frame, width=args.width, draw_contours=True) + key = cp.show_frame(frame, width=args.width) if key == ord('q') or key == 27: break finally: diff --git a/pykick/finders.py b/pykick/finders.py index dbaa82f..815b25a 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -1,6 +1,11 @@ +from __future__ import division +from __future__ import print_function + import json from collections import deque + import cv2 +import numpy as np class GoalFinder(object): @@ -11,8 +16,8 @@ class GoalFinder(object): self.hsv_upper = hsv_upper def goal_similarity(self, contour): - contour = contour.reshape((-1, 2)) - hull = cv2.convexHull(contour).reshape((-1, 2)) + contour = contour.squeeze(axis=1) + hull = cv2.convexHull(contour).squeeze(axis=1) len_h = cv2.arcLength(hull, True) # Wild assumption that the goal should lie close to its @@ -32,8 +37,10 @@ class GoalFinder(object): print(shape_sim, area_sim, final_score) return final_score - def find_goal_contour(self, frame) - thr = + def find_goal_contour(self, frame): + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + print(self.hsv_lower, self.hsv_upper) + thr = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) # The ususal thr = cv2.erode(thr, None, iterations=2) @@ -47,8 +54,7 @@ class GoalFinder(object): 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.01 * perimeters + epsilon = [0.01 * cv2.arcLength(cnt, True) for cnt in cnts] # Approximate resulting contours with simpler lines cnts = [cv2.approxPolyDP(cnt, eps, True) @@ -64,10 +70,19 @@ class GoalFinder(object): similarities = [self.goal_similarity(cnt) for cnt in good_cnts] best = min(similarities) - if best > 0.4: - return None + # if best > 0.4: + # return None # Find the contour with the shape closest to that of the goal goal = good_cnts[similarities.index(best)] + return goal + + def left_right_post(self, contour): + return contour[:,0].min(), contour[:,0].max() + + def draw(self, frame): + goal = self.find_goal_contour(frame) + if goal is not None: + cv2.drawContours(frame, (goal,), -1, (0, 255, 0), 2) class BallFinder(object): @@ -102,7 +117,6 @@ class BallFinder(object): # only proceed if at least one contour was found if len(cnts) == 0: - self.history.appendleft(None) return None # find the largest contour in the mask, then use it to compute @@ -111,23 +125,18 @@ class BallFinder(object): ((x, y), radius) = cv2.minEnclosingCircle(c) if radius < self.min_radius: - self.history.appendleft(None) return None M = cv2.moments(c) center = (int(M["m10"] / M["m00"]),int(M["m01"] // M["m00"])) - self.history.appendleft((center, int(radius))) return center, int(radius) - def visualize(self, frame): - if not self.viz: - raise ValueError( - 'Visualization needs to be enabled when initializing' - ) + def draw(self, frame): + ball = self.find_colored_ball(frame) + self.history.appendleft(ball) - frame = frame.copy() - if self.history[0] is not None: - center, radius = self.history[0] + 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) @@ -138,13 +147,10 @@ class BallFinder(object): continue # otherwise, compute the thickness of the line and # draw the connecting lines - center_now = self.history[0][0] - center_prev = self.history[1][0] + 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) - # show the frame to screen - cv2.imshow("Frame", frame) - return cv2.waitKey(1) def load_hsv_config(self, filename): with open(filename) as f: diff --git a/pykick/imagereaders.py b/pykick/imagereaders.py index 1f67d6f..4b0ca25 100644 --- a/pykick/imagereaders.py +++ b/pykick/imagereaders.py @@ -60,7 +60,8 @@ class VideoReader(object): def get_frame(self): succ, frame = self.cap.read() if not succ: - raise ValueError('Error while reading video') + raise ValueError('Error while reading video.\n' + + 'Or video is over.') self.ctr += 1 if (self.ctr == self.cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT) and self.loop):