Pavel Lutskov
69b1b137a2
1. Created interface classes for reading video streams (in imagereaders.py) 2. Created a class for ball detection (for reusability) 3. Reworked colorpicker, now it is possible to choose the mode of operation from command line (available are live from Nao, from video file or from webcam). It is possible to capture only the first image of a stream and work on it. Colorpicker now can save the settings on exit or load settings on startup.
114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
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)
|
|
|
|
|
|
class BallFinder(object):
|
|
|
|
def __init__(self, hsv_lower, hsv_upper, min_radius, width):
|
|
|
|
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
|
|
|
|
cv2.namedWindow('ball_mask')
|
|
cv2.namedWindow('Frame')
|
|
|
|
def find_colored_ball(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)
|
|
cv2.imshow('ball_mask', mask)
|
|
|
|
# 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:
|
|
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)
|
|
|
|
if radius < self.min_radius:
|
|
return None
|
|
|
|
M = cv2.moments(c)
|
|
center = (int(M["m10"] / M["m00"]),int(M["m01"] // M["m00"]))
|
|
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
|
|
|
|
self.history.appendleft(self.last_center)
|
|
self.draw_ball_markers(frame)
|
|
|
|
# show the frame to screen
|
|
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)
|
|
|
|
# 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
|
|
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
|
|
|
|
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()
|