diff --git a/pykick/colorpicker.py b/pykick/colorpicker.py index eaa15ba..fcfb9c1 100644 --- a/pykick/colorpicker.py +++ b/pykick/colorpicker.py @@ -218,7 +218,11 @@ if __name__ == '__main__': rdr.close() while True: if not args.still: - frame = rdr.get_frame() + try: + frame = rdr.get_frame() + except RuntimeError as e: + print(e) + continue key = cp.show_frame(frame, width=args.width, manual=args.manual) if key == ord('q') or key == 27: break diff --git a/pykick/detection_demo.py b/pykick/detection_demo.py index d0569ee..887dec4 100644 --- a/pykick/detection_demo.py +++ b/pykick/detection_demo.py @@ -96,7 +96,11 @@ if __name__ == '__main__': rdr.close() while True: if not args.still: - frame = rdr.get_frame() + try: + frame = rdr.get_frame() + except RuntimeError as e: + print(e) + continue frame = imresize(frame, width=args.width) field = field_finder.find(frame) diff --git a/pykick/finders.py b/pykick/finders.py index bd74f08..7cbb573 100644 --- a/pykick/finders.py +++ b/pykick/finders.py @@ -133,13 +133,18 @@ class GoalFinder(object): class BallFinder(object): - def __init__(self, hsv_lower, hsv_upper, min_radius): + def __init__(self, hsv_lower, hsv_upper, min_radius=0.02): self.hsv_lower = tuple(hsv_lower) self.hsv_upper = tuple(hsv_upper) self.min_radius = min_radius self.history = deque(maxlen=64) + def primary_mask(self, frame): + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + mask = cv2.inRange(hsv, self.hsv_lower, self.hsv_upper) + return mask + def find(self, frame): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)