131 lines
3.6 KiB
Python
131 lines
3.6 KiB
Python
from __future__ import print_function
|
|
from __future__ import division
|
|
|
|
import argparse
|
|
|
|
import cv2
|
|
|
|
from .utils import read_config, imresize
|
|
from .imagereaders import NaoImageReader, VideoReader, PictureReader
|
|
from .finders import BallFinder, GoalFinder, FieldFinder
|
|
|
|
|
|
if __name__ == '__main__':
|
|
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['ball_min_radius'])
|
|
field_finder = FieldFinder(defaults['field'][0], defaults['field'][1])
|
|
ball_window = 'Ball detection'
|
|
goal_window = 'Goal detection'
|
|
cv2.namedWindow(ball_window)
|
|
cv2.namedWindow(goal_window)
|
|
|
|
try:
|
|
if args.still:
|
|
frame = rdr.get_frame()
|
|
rdr.close()
|
|
while True:
|
|
if not args.still:
|
|
try:
|
|
frame = rdr.get_frame()
|
|
except RuntimeError as e:
|
|
print(e)
|
|
continue
|
|
frame = imresize(frame, width=args.width)
|
|
|
|
field = field_finder.find(frame)
|
|
|
|
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)
|
|
|
|
if goal is not None:
|
|
print(goal)
|
|
print('LR post', goal_finder.left_right_post(goal))
|
|
print('----')
|
|
|
|
ball_frame = ball_finder.draw(ball_frame, ball)
|
|
goal_frame = goal_finder.draw(goal_frame, goal)
|
|
|
|
cv2.imshow(ball_window, ball_frame)
|
|
cv2.imshow(goal_window, goal_frame)
|
|
|
|
key = cv2.waitKey(0 if args.manual else 1)
|
|
if key == ord('q') or key == 27:
|
|
break
|
|
finally:
|
|
if not args.still:
|
|
rdr.close()
|