240 lines
7.6 KiB
Python
240 lines
7.6 KiB
Python
from __future__ import print_function
|
|
from __future__ import division
|
|
|
|
import json
|
|
import argparse
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .imagereaders import VideoReader, NaoImageReader, PictureReader
|
|
from .utils import read_config, imresize
|
|
|
|
class Colorpicker(object):
|
|
|
|
WINDOW_CAPTURE_NAME = 'Video Capture'
|
|
WINDOW_DETECTION_NAME = 'Object Detection'
|
|
|
|
def __init__(self):
|
|
parameters = ['low_h', 'low_s', 'low_v', 'high_h', 'high_s', 'high_v']
|
|
maxes = [180, 255, 255, 180, 255, 255]
|
|
checkers = [
|
|
lambda x: min(x, self.settings['high_h'] - 1), # LOW H
|
|
lambda x: min(x, self.settings['high_s'] - 1), # LOW S
|
|
lambda x: min(x, self.settings['high_v'] - 1), # LOW V
|
|
lambda x: max(x, self.settings['low_h'] + 1), # HIGH H
|
|
lambda x: max(x, self.settings['low_s'] + 1), # HIGH S
|
|
lambda x: max(x, self.settings['low_v'] + 1), # HIGH V
|
|
]
|
|
self.settings = {
|
|
'low_h': 0,
|
|
'low_s': 0,
|
|
'low_v': 0,
|
|
'high_h': 180,
|
|
'high_s': 255,
|
|
'high_v': 255
|
|
}
|
|
cv2.namedWindow(self.WINDOW_CAPTURE_NAME)
|
|
cv2.namedWindow(self.WINDOW_DETECTION_NAME)
|
|
self.trackers = [
|
|
cv2.createTrackbar(
|
|
name, self.WINDOW_DETECTION_NAME, self.settings[name], max_v,
|
|
lambda val, name=name, checker=checker: self._on_trackbar(
|
|
val, name, checker
|
|
)
|
|
)
|
|
for name, max_v, checker in zip(parameters, maxes, checkers)
|
|
]
|
|
|
|
def do_print(self):
|
|
print(self.settings)
|
|
|
|
def _on_trackbar(self, val, name, checker):
|
|
self.settings[name] = checker(val)
|
|
cv2.setTrackbarPos(name, self.WINDOW_DETECTION_NAME,
|
|
self.settings[name])
|
|
|
|
def goal_similarity(self, contour):
|
|
contour = contour.reshape((-1, 2))
|
|
hull = cv2.convexHull(contour).reshape((-1, 2))
|
|
len_h = cv2.arcLength(hull, True)
|
|
|
|
# Wild assumption that the goal should lie close to its
|
|
# enclosing convex hull
|
|
shape_sim = np.linalg.norm(contour[:,None] - hull,
|
|
axis=2).min(axis=1).sum() / len_h
|
|
|
|
# Wild assumption that the area of the goal is rather small
|
|
# compared to its enclosing convex hull
|
|
area_c = cv2.contourArea(contour)
|
|
area_h = cv2.contourArea(hull)
|
|
|
|
area_sim = area_c / area_h
|
|
|
|
# Final similarity score is just the sum of both
|
|
final_score = shape_sim + area_sim
|
|
print(shape_sim, area_sim, final_score)
|
|
return final_score
|
|
|
|
def draw_contours(self, thr):
|
|
# The ususal
|
|
thr = cv2.erode(thr, None, iterations=2)
|
|
thr = cv2.dilate(thr, None, iterations=2)
|
|
cnts, _ = cv2.findContours(thr.copy(), cv2.RETR_EXTERNAL,
|
|
cv2.CHAIN_APPROX_SIMPLE)
|
|
areas = np.array([cv2.contourArea(cnt) for cnt in cnts])
|
|
perimeters = np.array([cv2.arcLength(cnt, True) for cnt in cnts])
|
|
epsilon = 0.04 * perimeters
|
|
|
|
# Candidates are at most 6 biggest white areas
|
|
top_x = 6
|
|
if len(areas) > top_x:
|
|
cnt_ind = np.argpartition(areas, -top_x)[-top_x:]
|
|
cnts = [cnts[i] for i in cnt_ind]
|
|
|
|
perimeters = np.array([cv2.arcLength(cnt, True) for cnt in cnts])
|
|
epsilon = 0.01 * perimeters
|
|
|
|
# Approximate resulting contours with simpler lines
|
|
cnts = [cv2.approxPolyDP(cnt, eps, True)
|
|
for cnt, eps in zip(cnts, epsilon)]
|
|
|
|
# Goal needs normally 8 points for perfect approximation
|
|
# But with 6 can also be approximated
|
|
good_cnts = [cnt for cnt in cnts if 6 <= cnt.shape[0] <= 9
|
|
and not cv2.isContourConvex(cnt)]
|
|
|
|
if good_cnts:
|
|
# Find the contour with the shape closest to that of the goal
|
|
good_cnts = [min(good_cnts, key=self.goal_similarity)]
|
|
|
|
thr = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR)
|
|
cv2.drawContours(thr, good_cnts, -1, (0, 255, 0), 2)
|
|
return thr
|
|
|
|
def show_frame(self, frame, width=None, draw_contours=False):
|
|
frame_HSV = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
|
|
frame_threshold = cv2.inRange(
|
|
frame_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)
|
|
|
|
if draw_contours:
|
|
frame_threshold = self.draw_contours(frame_threshold)
|
|
|
|
cv2.imshow(self.WINDOW_CAPTURE_NAME, frame)
|
|
cv2.imshow(self.WINDOW_DETECTION_NAME, frame_threshold)
|
|
return cv2.waitKey(1)
|
|
|
|
def save(self, filename):
|
|
try:
|
|
with open(filename) as f:
|
|
conf = json.load(f)
|
|
except IOError:
|
|
conf = {}
|
|
conf.update(self.settings)
|
|
with open(filename, 'w') as f:
|
|
json.dump(conf, f, indent=4)
|
|
|
|
def load(self, filename):
|
|
with open(filename) as f:
|
|
self.settings = json.load(f)
|
|
for name in self.settings:
|
|
cv2.setTrackbarPos(name, self.WINDOW_DETECTION_NAME,
|
|
self.settings[name])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
defaults = read_config()
|
|
parser = argparse.ArgumentParser(
|
|
epilog='When called without arguments specifying the video source, ' +
|
|
'will try to use the webcam'
|
|
)
|
|
parser.add_argument(
|
|
'-o', '--output-config',
|
|
help='file, to which the settings will be saved (if given)'
|
|
)
|
|
parser.add_argument(
|
|
'-i', '--input-config',
|
|
help='file, from which to read the initial values'
|
|
)
|
|
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(
|
|
'--nao-cam',
|
|
choices=[0, 1],
|
|
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()
|
|
|
|
cp = Colorpicker()
|
|
if args.input_config:
|
|
cp.load(args.input_config)
|
|
|
|
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)
|
|
|
|
try:
|
|
if args.still:
|
|
frame = rdr.get_frame()
|
|
rdr.close()
|
|
while True:
|
|
if not args.still:
|
|
frame = rdr.get_frame()
|
|
key = cp.show_frame(frame, width=args.width, draw_contours=True)
|
|
if key == ord('q') or key == 27:
|
|
break
|
|
finally:
|
|
cp.do_print()
|
|
if args.output_config:
|
|
cp.save(args.output_config)
|
|
if not args.still:
|
|
rdr.close()
|