diff --git a/scripts/colorpicker.py b/scripts/colorpicker.py index c822acf..e7ff280 100644 --- a/scripts/colorpicker.py +++ b/scripts/colorpicker.py @@ -67,8 +67,14 @@ class Colorpicker(object): return cv2.waitKey(1) def save(self, filename): + try: + with open(filename) as f: + conf = json.load(f) + except FileNotFoundError: + conf = {} + conf.update(self.settings) with open(filename, 'w') as f: - json.dump(self.settings, f, indent=4) + json.dump(conf, f, indent=4) def load(self, filename): with open(filename) as f: diff --git a/scripts/detection_demo.py b/scripts/detection_demo.py new file mode 100644 index 0000000..bc2c940 --- /dev/null +++ b/scripts/detection_demo.py @@ -0,0 +1,22 @@ +from __future__ import print_function +from __future__ import division + +import json +from imagereaders import NaoImageReader, VideoReader +from finders import BallFinder + + +if __name__ == '__main__': + video = VideoReader(0, loop=True) + with open('ball_hsv.json') as f: + hsv = json.load(f) + hsv_lower = tuple(map(hsv.get, ('low_h', 'low_s', 'low_v'))) + hsv_upper = tuple(map(hsv.get, ('high_h', 'high_s', 'high_v'))) + finder = BallFinder(hsv_lower, hsv_upper, 5, None) + try: + while True: + frame = video.get_frame() + finder.find_colored_ball(frame) + finder.visualize(frame) + finally: + video.close() diff --git a/scripts/live_recognition.py b/scripts/finders.py similarity index 55% rename from scripts/live_recognition.py rename to scripts/finders.py index c1f995e..7dc08dd 100644 --- a/scripts/live_recognition.py +++ b/scripts/finders.py @@ -1,30 +1,16 @@ -from __future__ import print_function -from __future__ import division - import json -import cv2 -import numpy as np -# import imutils -from imagereaders import NaoImageReader, VideoReader from collections import deque - - -red_lower = (0, 185, 170) # HSV coded red interval -red_upper = (2, 255, 255) +import cv2 class BallFinder(object): - def __init__(self, hsv_lower, hsv_upper, min_radius, width, - viz=False): + def __init__(self, hsv_lower, hsv_upper, min_radius, viz=False): self.hsv_lower = hsv_lower self.hsv_upper = hsv_upper self.min_radius = min_radius - self.width = width self.history = deque(maxlen=64) - self.last_center = None - self.last_radius = None self.viz = viz if self.viz: @@ -50,6 +36,7 @@ class BallFinder(object): # only proceed if at least one contour was found if len(cnts) == 0: print('Nothin there') + self.history.appendleft(None) return None # find the largest contour in the mask, then use it to compute @@ -58,35 +45,25 @@ 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 next_frame(self, frame): - # maybe resize the frame, maybe blur it - # if self.width is not None: - # frame = imutils.resize(frame, width=self.width) - try: - self.last_center, self.last_radius = self.find_colored_ball(frame) - except TypeError: # No red ball found and function returned None - self.last_center, self.last_radius = None, None + def visualize(self, frame): + if not self.viz: + raise ValueError( + 'Visualization needs to be enabled when initializing' + ) - self.history.appendleft(self.last_center) - self.draw_ball_markers(frame) - - # show the frame to screen - if self.viz: - cv2.imshow("Frame", frame) - return cv2.waitKey(2) - - def draw_ball_markers(self, frame): - # draw the enclosing circle and ball's centroid on the frame, - if self.last_center is not None and self.last_radius is not None: - cv2.circle(frame, self.last_center, self.last_radius, - (255, 255, 0), 1) - cv2.circle(frame, self.last_center, 5, (0, 255, 0), -1) + frame = frame.copy() + if self.history[0] is not None: + center, radius = self.history[0] + cv2.circle(frame, center, radius, (255, 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)): @@ -95,25 +72,16 @@ class BallFinder(object): continue # otherwise, compute the thickness of the line and # draw the connecting lines - thickness = int(np.sqrt(64 / float(i + 1)) * 2.5) - cv2.line(frame, self.history[i - 1], self.history[i], - (0, 255, 0), thickness) - - return frame + center_now = self.history[0][0] + center_prev = self.history[1][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: 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'))) - - -if __name__ == '__main__': - # video = NaoImageReader('192.168.0.11') - video = VideoReader(0, loop=True) - finder = BallFinder(red_lower, red_upper, 5, None) - try: - while True: - finder.next_frame(video.get_frame()) - finally: - video.close() diff --git a/scripts/nao_defaults.json b/scripts/nao_defaults.json index 32cec7e..9bfb3a4 100644 --- a/scripts/nao_defaults.json +++ b/scripts/nao_defaults.json @@ -2,5 +2,13 @@ "ip": "192.168.0.11", "res": 1, "port": 9559, - "cam": 0 + "cam": 0, + "fps": 30, + "low_s": 175, + "low_v": 100, + "high_h": 6, + "high_v": 255, + "low_h": 0, + "high_s": 255, + "min_radius": 5 }