guess what? - refactoring! also masking out field

pykick/finders.py

@@ -11,8 +11,8 @@ class GoalFinder(object):
 
     def __init__(self, hsv_lower, hsv_upper):
 
-        self.hsv_lower = hsv_lower
-        self.hsv_upper = hsv_upper
+        self.hsv_lower = tuple(hsv_lower)
+        self.hsv_upper = tuple(hsv_upper)
 
     def goal_similarity(self, contour):
         contour = contour.squeeze(axis=1)
@@ -33,7 +33,7 @@ class GoalFinder(object):
 
         # Final similarity score is just the sum of both
         final_score = shape_sim + area_sim
-        print('Candidate:', shape_sim, area_sim, final_score)
+        print('Goal candidate:', shape_sim, area_sim, final_score)
         return final_score
 
     def find_goal_contour(self, frame):
@@ -68,7 +68,7 @@ class GoalFinder(object):
 
         similarities = [self.goal_similarity(cnt) for cnt in good_cnts]
         best = min(similarities)
-        print('Final score:', best)
+        print('Final goal score:', best)
         print()
         if best > 0.35:
             return None
@@ -83,8 +83,7 @@ class GoalFinder(object):
         l, r = self.left_right_post(self, contour)
         return (l + r) / 2
 
-    def draw(self, frame):
-        goal = self.find_goal_contour(frame)
+    def draw(self, frame, goal):
         if goal is not None:
             cv2.drawContours(frame, (goal,), -1, (0, 255, 0), 2)
 
@@ -93,8 +92,8 @@ class BallFinder(object):
 
     def __init__(self, hsv_lower, hsv_upper, min_radius):
 
-        self.hsv_lower = hsv_lower
-        self.hsv_upper = hsv_upper
+        self.hsv_lower = tuple(hsv_lower)
+        self.hsv_upper = tuple(hsv_upper)
         self.min_radius = min_radius
         self.history = deque(maxlen=64)
 
@@ -102,55 +101,56 @@ class BallFinder(object):
         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
+        # 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)
+        # mask = cv2.erode(mask, None, iterations=2)
+        # mask = cv2.dilate(mask, None, iterations=2)
 
-        # find contours in the mask and initialize the current
-        # (x, y) center of the ball
         cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
                                 cv2.CHAIN_APPROX_SIMPLE)[-2]
 
-        # only proceed if at least one contour was found
         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)
+        (x, y), radius = cv2.minEnclosingCircle(c)
 
-        if radius < self.min_radius:
+        min_radius_abs = self.min_radius * frame.shape[0]
+
+        if radius < min_radius_abs:
+            print('Radius:', radius, 'Min radius:', min_radius_abs)
+            self.history.appendleft(None)
             return None
 
         M = cv2.moments(c)
-        center = (int(M["m10"] / M["m00"]),int(M["m01"] // M["m00"]))
+        try:
+            center = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])
+        except ZeroDivisionError:
+            # It's weird but happened yeah
+            self.history.append(None)
+            return None
+        self.history.appendleft((center, int(radius)))
+        print('Ball:', center, radius)
         return center, int(radius)
 
-    def draw(self, frame):
-        ball = self.find_colored_ball(frame)
-        self.history.appendleft(ball)
-
+    def draw(self, frame, ball):
         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)
+            # cv2.circle(frame, center, 5, (0, 255, 0), -1)
 
         # loop over the set of tracked points
-        for i in range(1, len(self.history)):
+        # 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
+            # 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)
-
-    # def load_hsv_config(self, filename):
-        # with open(filename) as f:
-            # hsv = json.load(f)
-        # self.hsv_lower = tuple(map(hsv.get, ('low_h', 'low_s', 'low_v')))
-        # self.hsv_upper = tuple(map(hsv.get, ('high_h', 'high_s', 'high_v')))
+            # 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)