diff --git a/pykick/detection_demo.py b/pykick/detection_demo.py index e4220ce..d0569ee 100644 --- a/pykick/detection_demo.py +++ b/pykick/detection_demo.py @@ -5,7 +5,7 @@ import argparse import cv2 -from .utils import read_config, imresize, field_mask +from .utils import read_config, imresize from .imagereaders import NaoImageReader, VideoReader, PictureReader from .finders import BallFinder, GoalFinder, FieldFinder @@ -100,10 +100,9 @@ if __name__ == '__main__': frame = imresize(frame, width=args.width) field = field_finder.find(frame) - not_field = cv2.bitwise_not(field) - ball_frame = field_finder.draw(frame, field) - goal_frame = field_finder.draw(frame, not_field) + ball_frame = field_finder.mask_it(frame, field) + goal_frame = field_finder.mask_it(frame, field, inverse=True) ball = ball_finder.find(ball_frame) goal = goal_finder.find(goal_frame) diff --git a/pykick/finders.py b/pykick/finders.py index 08e7424..2db0b5f 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -29,13 +29,22 @@ class FieldFinder(object): return None field = max(cnts, key=cv2.contourArea) field = cv2.convexHull(field) - mask = np.zeros(thr.shape, dtype=np.uint8) - cv2.drawContours(mask, (field,), -1, 255, -1) - return mask + return field def draw(self, frame, field): if field is not None: - frame = cv2.bitwise_and(frame, frame, mask=field) + frame = frame.copy() + cv2.drawContours(frame, (field,), -1, (0, 0, 255), 2) + return frame + + def mask_it(self, frame, field, inverse=False): + if field is not None: + print(frame.shape) + mask = np.zeros(frame.shape[:2], dtype=np.uint8) + cv2.drawContours(mask, (field,), -1, 255, -1) + if inverse: + mask = cv2.bitwise_not(mask) + frame = cv2.bitwise_and(frame, frame, mask=mask) return frame diff --git a/pykick/nao_defaults.json b/pykick/nao_defaults.json index a0ec935..d5150fe 100644 --- a/pykick/nao_defaults.json +++ b/pykick/nao_defaults.json @@ -29,11 +29,11 @@ "field": [ [ 19, - 57, - 84 + 51, + 60 ], [ - 65, + 67, 255, 255 ] diff --git a/pykick/striker.py b/pykick/striker.py index f76c5f9..924b186 100644 --- a/pykick/striker.py +++ b/pykick/striker.py @@ -1,6 +1,8 @@ from __future__ import print_function from __future__ import division +import argparse +from threading import Thread from math import pi from time import sleep, time @@ -8,9 +10,7 @@ from .utils import read_config from .imagereaders import NaoImageReader from .finders import BallFinder, GoalFinder from .movements import NaoMover -import argparse from naoqi import ALProxy -from threading import Thread class Striker(object): @@ -72,7 +72,7 @@ class Striker(object): """Detect the ball and return its angles in camera coordinates.""" try: ball = self.ball_finder.find(cam.get_frame()) - except RuntimeError as e: + except RuntimeError as e: # Sometimes camera doesn't return an image print(e) return None