From 9f3fe74be6d5f33120516ae78a5d9f04e6413fef Mon Sep 17 00:00:00 2001 From: Pavel Lutskov Date: Fri, 29 Jun 2018 11:25:19 +0200 Subject: [PATCH] Better Hue overflow handling --- pykick/colorpicker.py | 2 ++ pykick/finders.py | 38 +++++++++----------------------------- pykick/utils.py | 4 ++-- 3 files changed, 13 insertions(+), 31 deletions(-) diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index 4ac2aae..a64ef95 100644 --- a/pykick/colorpicker.py +++ b/pykick/colorpicker.py @@ -76,6 +76,8 @@ class Colorpicker(object): def _hsv_updated(self, param): cv2.setTrackbarPos(param, self.WINDOW_DETECTION_NAME, self.settings[param]) + if self.marker is None: + return self.marker.hsv_lower = tuple( map(self.settings.get, ('low_h', 'low_s', 'low_v')) ) diff --git a/pykick/finders.py b/pykick/finders.py index 7cbb573..0b16ca5 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -6,6 +6,8 @@ from collections import deque import cv2 import numpy as np +from .utils import hsv_mask + class FieldFinder(object): @@ -16,7 +18,7 @@ class FieldFinder(object): 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 = hsv_mask(blurred, self.hsv_lower, self.hsv_upper) thr = cv2.erode(thr, None, iterations=6) thr = cv2.dilate(thr, None, iterations=10) return thr @@ -55,7 +57,7 @@ class GoalFinder(object): def primary_mask(self, frame): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - thr = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) + thr = hsv_mask(hsv, self.hsv_lower, self.hsv_upper) thr = cv2.erode(thr, None, iterations=2) thr = cv2.dilate(thr, None, iterations=2) return thr @@ -86,7 +88,7 @@ class GoalFinder(object): def find(self, frame): thr = self.primary_mask(frame) cnts, _ = cv2.findContours(thr, cv2.RETR_EXTERNAL, - cv2.CHAIN_APPROX_SIMPLE) + cv2.CHAIN_APPROX_SIMPLE) cnts.sort(key=cv2.contourArea, reverse=True) top_x = 6 cnts = cnts[:top_x] @@ -142,28 +144,19 @@ class BallFinder(object): def primary_mask(self, frame): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) - mask = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) + mask = hsv_mask(hsv, self.hsv_lower, self.hsv_upper) return mask def find(self, frame): - 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 ? - mask = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) - # mask = cv2.erode(mask, None, iterations=2) - # mask = cv2.dilate(mask, None, iterations=2) - - cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, - cv2.CHAIN_APPROX_SIMPLE)[-2] + mask = self.primary_mask(frame) + cnts, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, + cv2.CHAIN_APPROX_SIMPLE) 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) @@ -191,16 +184,3 @@ class BallFinder(object): 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 - # 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: - # continue - # otherwise, compute the thickness of the line and - # draw the connecting lines - # 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) diff --git a/pykick/utils.py b/pykick/utils.py index 35e951b..cbb949b 100644 --- a/pykick/utils.py +++ b/pykick/utils.py @@ -34,8 +34,8 @@ def imresize(frame, width=None, height=None): def hsv_mask(hsv, hsv_lower, hsv_upper): if hsv_lower[0] > hsv_upper[0]: mask_l = cv2.inRange(hsv, tuple(hsv_lower), - tuple([180] + hsv_upper[1:])) - mask_u = cv2.inRange(hsv, tuple([0] + hsv_lower[1:]), + (180,) + tuple(hsv_upper[1:])) + mask_u = cv2.inRange(hsv, (0,) + tuple(hsv_lower[1:]), tuple(hsv_upper)) return cv2.add(mask_l, mask_u) else: