pavel/kick-it
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

guess what? - refactoring! also masking out field

Browse Source
This commit is contained in:
Pavel Lutskov
2018-06-20 20:38:11 +02:00
parent a5e6059d3c
commit 0b53d11861
6 changed files with 217 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
pykick/colorpicker.py
View File

@@ -8,7 +8,7 @@ import cv2
from .imagereaders import VideoReader, NaoImageReader, PictureReader
from .finders import GoalFinder
from .utils import read_config, imresize
from .utils import read_config, imresize, field_mask
class Colorpicker(object):
@@ -72,25 +72,28 @@ class Colorpicker(object):
map(self.settings.get, ('high_h', 'high_s', 'high_v'))
)
def show_frame(self, frame, width=None):
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
frame_threshold = cv2.inRange(
hsv,
def show_frame(self, frame, width=None, manual=False):
# hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# frame_threshold = cv2.inRange(
# hsv,
# tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))),
# tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v')))
# )
frame = imresize(frame, width=width)
# frame_threshold = imresize(frame_threshold, width=width)
frame_threshold = field_mask(
frame,
tuple(map(self.settings.get, ('low_h', 'low_s', 'low_v'))),
tuple(map(self.settings.get, ('high_h', 'high_s', 'high_v')))
)
frame = imresize(frame, width=width)
frame_threshold = imresize(frame_threshold, width=width)
if 'goal' in self.markers:
self.markers['goal'].draw(frame)
if 'ball' in self.markers:
self.markers['ball'].draw(frame)
for marker in self.markers:
self.markers[marker].draw(frame)
cv2.imshow(self.WINDOW_CAPTURE_NAME, frame)
cv2.imshow(self.WINDOW_DETECTION_NAME, frame_threshold)
return cv2.waitKey(1)
return cv2.waitKey(0 if manual else 1)
def save(self, filename, color):
try:
@@ -154,6 +157,13 @@ if __name__ == '__main__':
help='only take one image from video stream',
action='store_true'
)
parser.add_argument(
'--manual',
help='switch frames manually',
action='store_true',
default=False
)
parser.add_argument(
'--nao-cam',
choices=[0, 1],
@@ -175,15 +185,15 @@ if __name__ == '__main__':
default=640
)
parser.add_argument(
'--color',
help='specify which color is being calibrated',
default='white'
'--target',
help='specify for what target is being calibrated',
default='field'
)
args = parser.parse_args()
cp = Colorpicker()
if args.input_config:
cp.load(args.input_config, args.color)
cp.load(args.input_config, args.target)
if args.video_file:
rdr = VideoReader(args.video_file, loop=True)
@@ -205,12 +215,12 @@ if __name__ == '__main__':
while True:
if not args.still:
frame = rdr.get_frame()
key = cp.show_frame(frame, width=args.width)
key = cp.show_frame(frame, width=args.width, manual=args.manual)
if key == ord('q') or key == 27:
break
finally:
cp.do_print()
if args.output_config:
cp.save(args.output_config, args.color)
cp.save(args.output_config, args.target)
if not args.still:
rdr.close()

118
pykick/detection_demo.py
View File

@@ -1,21 +1,115 @@
from __future__ import print_function
from __future__ import division
from .utils import read_config
from .imagereaders import NaoImageReader, VideoReader
from .finders import BallFinder
import argparse
import cv2
from .utils import read_config, imresize, field_mask
from .imagereaders import NaoImageReader, VideoReader, PictureReader
from .finders import BallFinder, GoalFinder
if __name__ == '__main__':
video = VideoReader(0, loop=True)
cfg = read_config()
hsv_lower = cfg['red'][0]
hsv_upper = cfg['red'][1]
finder = BallFinder(hsv_lower, hsv_upper, cfg['min_radius'], None)
defaults = read_config()
parser = argparse.ArgumentParser()
# parser.add_argument(
# '-i', '--input-config',
# help='file, from which to read the initial values',
# required=True
# )
imsource = parser.add_mutually_exclusive_group()
imsource.add_argument(
'--video-file',
help='video file to use'
)
imsource.add_argument(
'--image-file',
help='image to use'
)
imsource.add_argument(
'--nao-ip',
help='ip address of the nao robot, from which to capture',
default=False,
const=defaults['ip'],
nargs='?'
)
parser.add_argument(
'--still',
help='only take one image from video stream',
action='store_true'
)
parser.add_argument(
'--manual',
help='switch frames manually',
action='store_true',
default=False
)
parser.add_argument(
'--nao-cam',
choices=[0, 1],
type=int,
help='0 for top camera, 1 for bottom camera',
default=defaults['cam']
)
parser.add_argument(
'--nao-res',
choices=[1, 2, 3],
help='choose a nao resolution',
type=int,
default=defaults['res']
)
parser.add_argument(
'--width',
help='specify width of the image output',
type=int,
default=640
)
args = parser.parse_args()
if args.video_file:
rdr = VideoReader(args.video_file, loop=True)
elif args.image_file:
rdr = PictureReader(args.image_file)
elif args.nao_ip:
rdr = NaoImageReader(
args.nao_ip,
cam_id=args.nao_cam,
res=args.nao_res
)
else:
rdr = VideoReader(0)
goal_finder = GoalFinder(defaults['goal'][0], defaults['goal'][1])
ball_finder = BallFinder(defaults['ball'][0], defaults['ball'][1],
defaults['min_radius'])
window_name = 'Live detection'
cv2.namedWindow(window_name)
try:
if args.still:
frame = rdr.get_frame()
rdr.close()
while True:
frame = video.get_frame()
finder.find_colored_ball(frame)
finder.visualize(frame)
if not args.still:
frame = rdr.get_frame()
frame = imresize(frame, width=args.width)
field_hsv = defaults['field']
field = field_mask(frame, field_hsv[0], field_hsv[1])
not_field = cv2.bitwise_not(field)
goal = goal_finder.find_goal_contour(
cv2.bitwise_and(frame, frame, mask=not_field))
ball = ball_finder.find_colored_ball(
cv2.bitwise_and(frame, frame, mask=field))
goal_finder.draw(frame, goal)
ball_finder.draw(frame, ball)
cv2.imshow(window_name,
cv2.bitwise_and(frame, frame, mask=field))
key = cv2.waitKey(0 if args.manual else 1)
if key == ord('q') or key == 27:
break
finally:
video.close()
if not args.still:
rdr.close()

70
pykick/finders.py
View File

@@ -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)

36
pykick/nao_defaults.json
View File

@@ -1,10 +1,18 @@
{
"cam": 1,
"ip": "192.168.0.11",
"min_radius": 5,
"fps": 30,
"ball": [
[
0,
175,
100
],
[
6,
255,
255
]
],
"res": 1,
"white": [
"goal": [
[
0,
0,
@@ -16,17 +24,21 @@
255
]
],
"port": 9559,
"red": [
"fps": 30,
"ip": "192.168.0.11",
"field": [
[
0,
175,
100
11,
74,
28
],
[
6,
69,
255,
255
]
]
],
"cam": 1,
"min_radius": 0.01,
"port": 9559
}

16
pykick/striker.py
View File

@@ -12,17 +12,17 @@ from .movements import NaoMover
class Striker(object):
def __init__(self, nao_ip, nao_port, res, red_hsv, white_hsv,
min_radius, run_after):
def __init__(self, nao_ip, nao_port, res, ball_hsv, goal_hsv,
ball_min_radius, run_after):
self.mover = NaoMover(nao_ip=nao_ip, nao_port=nao_port)
self.mover.stand_up()
self.video_top = NaoImageReader(nao_ip, port=nao_port, res=res,
fps=30, cam_id=0)
self.video_bot = NaoImageReader(nao_ip, port=nao_port, res=res,
fps=30, cam_id=1)
self.ball_finder = BallFinder(tuple(red_hsv[0]), tuple(red_hsv[1]),
min_radius)
self.goal_finder = GoalFinder(tuple(white_hsv[0]), tuple(white_hsv[1]))
self.ball_finder = BallFinder(tuple(ball_hsv[0]), tuple(ball_hsv[1]),
ball_min_radius)
self.goal_finder = GoalFinder(tuple(goal_hsv[0]), tuple(goal_hsv[1]))
self.lock_counter = 0
self.loss_counter = 0
self.run_after = run_after
@@ -214,9 +214,9 @@ if __name__ == '__main__':
nao_ip=cfg['ip'],
nao_port=cfg['port'],
res=cfg['res'],
red_hsv=cfg['red'],
white_hsv=cfg['white'],
min_radius=cfg['min_radius'],
ball_hsv=cfg['ball'],
goal_hsv=cfg['goal'],
ball_min_radius=cfg['ball_min_radius'],
run_after=False
)
try:

16
pykick/utils.py
View File

@@ -2,7 +2,8 @@ from __future__ import division
import os
import json
from cv2 import resize as cv2_resize
import cv2
HERE = os.path.dirname(os.path.realpath(__file__))
@@ -26,4 +27,15 @@ def imresize(frame, width=None, height=None):
if width and height:
sf = 0
sz = (width, height)
return cv2_resize(frame, sz, fx=sf, fy=sf)
return cv2.resize(frame, sz, fx=sf, fy=sf)
def field_mask(frame, hsv_lower, hsv_upper):
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
blurred = cv2.GaussianBlur(hsv, (25, 25), 20)
thr = cv2.inRange(blurred, tuple(hsv_lower), tuple(hsv_upper))
# The ususal
thr = cv2.erode(thr, None, iterations=6)
thr = cv2.dilate(thr, None, iterations=20)
return thr