class VisionWrapper:
"""
Class that handles vision
"""
ENEMY_SCALE = 1.
GROUP9_SCALE = 1.
ALLY_SCALE = 1.
def __init__(self, pitch, color_settings, our_name, robot_details,
enable_gui=False, pc_name=None, robots_on_pitch=list(), goal=None):
"""
Entry point for the SDP system.
Params:
[int] pitch 0 - main pitch, 1 - secondary pitch
[int or string] color_settings [0, small, 1, big] - 0 or small for pitch color settings with small numbers (previously - pitch 0)
1 or big - pitch color settings with big numbers (previously - pitch 1)
[string] colour The colour of our teams plates
[string] our_name our name
[int] video_port port number for the camera
[boolean] enable_gui Does not draw the normal image to leave space for GUI.
Also forces the trackers to return circle contour positons in addition to robot details.
[string] pc_name Name of the PC to load the files from (BUT NOT SAVE TO). Will default to local machine if not specified.
"""
pitch = int(pitch)
assert pitch in [0, 1]
if goal: assert goal in ["left", "right"]
self.pitch = pitch
self.target_goal = goal
self.color_settings = color_settings
self.calibration = Configuration.read_calibration(
pc_name, create_if_missing=True)
self.video_settings = Configuration.read_video_config(
pc_name, create_if_missing=True)
# Set up camera for frames
self.camera = Camera(pitch)
self.camera.start_capture()
self.frame = self.camera.get_frame()
center_point = self.camera.get_adjusted_center()
# Set up vision
self.trackers = list()
self.world_objects = dict()
# Get machine-specific calibration
self.enable_gui = enable_gui
self.gui = None
# Initialize robots
self.robots = []
# Initialize ball states - which robot had the ball previously.
self.ball_median_size = 5
self.ball_index = 0
self.ball_states = [None] * self.ball_median_size
self.BALL_HOLDING_AREA_SCALE = 0.1
for r_name in robot_details.keys():
role = 'ally'
if robot_details[r_name]['main_colour'] != robot_details[our_name]['main_colour']:
role = 'enemy'
elif r_name == our_name:
role = 'group9'
print "[WRAPPER] Setting %s to %s." % (r_name, role)
self.robots.append(RobotInstance(r_name,
robot_details[r_name]['main_colour'],
robot_details[r_name]['side_colour'],
robot_details[r_name]['offset_angle'],
role,
r_name in robots_on_pitch))
# Draw various things on the image
self.draw_direction = True
self.draw_robot = True
self.draw_contours = True
self.draw_ball = True
self.vision = Vision(
pitch=pitch, frame_shape=self.frame.shape,
frame_center=center_point, calibration=self.calibration,
robots=self.robots,
return_circle_contours=enable_gui, trackers_out=self.trackers)
# Used for latency calculations
self.t0 = time()
self.delta_t = 0
if self.enable_gui:
self.gui = GUI(self.pitch, self.color_settings, self.calibration, self)
from threading import Thread
from gui.common import MainWindow
from gui.status_control import StatusUI
from gui.main import MainUI
self.status_window = None
def create_windows():
self.main_window = MainUI(self)
return [ self.main_window ]
Thread(name="Tkinter UI", target=MainWindow.create_and_show,
args=[create_windows]).start()
# Set up preprocessing and postprocessing
# self.postprocessing = Postprocessing()
self.preprocessing = Preprocessing()
self.side = our_name
self.frameQueue = []
self.video = None
def get_robots_raw(self):
# Filter robots that have no position
return [[r.name, r.visible, r.position, r.heading, RobotType.UNKNOWN]
for r in self.robots]
def get_robot_position(self, robot_name):
return filter(lambda r: r.name == robot_name, self.robots)[0].position
def get_robot_headings(self):
headings = dict()
for r in self.robots:
headings[r.name] = r.heading
return headings
def get_circle_position(self, robot_name):
for r in self.robots:
if r.name == robot_name:
return r.side_x, r.side_y
def get_all_robots(self):
robots = dict()
for r in self.robots:
robots[r.name] = self.get_robot_position(r.name)
return robots
def get_ball_position(self):
"""
:return: Actual ball position or predicted ball position if the robot was near it. Might return None.
"""
# TODO: call methods here to get robot region if new one will be used.
if not self.world_objects['ball'][2]:
r_name, _ = self._mode(self.ball_states)
if r_name:
for r in self.robots:
if r.name == r_name:
x, y = r.predicted_ball_pos
self.world_objects['ball'] = (x, y, 2)
if self.world_objects['ball'][2] and not \
(self.world_objects['ball'][0] == 42 and self.world_objects['ball'][1] == 42):
return self.world_objects['ball']
else:
return None
def get_pitch_dimensions(self):
from util.tools import get_pitch_size
return get_pitch_size()
def get_robot_direction(self, robot_name):
return filter(lambda r: r.name == robot_name, self.robots)[0]
def do_we_have_ball(self, robot_name):
return self.is_ball_in_range(robot_name, scale=self.BALL_HOLDING_AREA_SCALE)
def is_ball_in_range(self, robot_name, scale=1.):
"""
returns True if ball is in the grabbing zone.
:param robot_name: robot name
:param scale: Zone can be scaled, to accommodate for different robots
:return: boolean
"""
ball = self.get_ball_position()
if ball and ball[2]:
for r in self.robots:
if r.name == robot_name and r.visible:
if r.role == 'enemy':
if r.is_point_in_grabbing_zone(ball[0], ball[1], role=r.role, scale=self.ENEMY_SCALE, circular=False):
return True
elif r.role == 'group9':
if r.is_point_in_grabbing_zone(ball[0], ball[1], role=r.role, scale=self.GROUP9_SCALE, circular=False):
return True
else:
if r.is_point_in_grabbing_zone(ball[0], ball[1], role=r.role, scale=self.ALLY_SCALE, circular=False):
return True
break
return False
def is_ball_on_other(self, robot_name, zone, scale=1.):
"""
returns True if ball is in the side or bottom zone.
:param robot_name: robot name
:param zone: ["left", "right", "bottom"]
:param scale: Zone can be scaled, to accommodate for different robots
:return: boolean
"""
ball = self.get_ball_position()
if ball and ball[2]:
for r in self.robots:
if r.name == robot_name:
if r.is_point_in_other_zone(ball[0], ball[1], zone, role=r.role, scale=scale):
return True
break
return False
def change_drawing(self, key):
"""
Toggles drawing of contours, heading directions, robot positions and ball tracker
Usage:
* 'b' for ball tracker
* 'n' for contours
* 'j' for robot_tracker
* 'i' for heading direction
:param key: keypress
:return: nothing
"""
if key == ord('b'):
self.gui.draw_ball = not self.gui.draw_ball
elif key == ord('n'):
self.gui.draw_contours = not self.gui.draw_contours
# THIS WAS CRASHING EVERYTHING
# elif key == ord('j'):
# self.gui.draw_robot = not self.gui.draw_robot
elif key == ord('i'):
self.gui.draw_direction = not self.gui.draw_direction
def get_circle_contours(self):
"""
Careful! Does not return x and y values. Call minimum bounding circles if proper circle locations are required.
Or look at dem fish http://docs.opencv.org/2.4/doc/tutorials/imgproc/shapedescriptors/find_contours/find_contours.html
:return: [Contours]
"""
return self.world_objects.get('circles')
def get_frame_size(self):
height, width, channels = self.frame.shape
return width, height
def get_latency_seconds(self):
return self.delta_t
def _mode(self, array):
"""
:param array: some array
:return: m, c the first mode found and the number of occurences
"""
m = max(array, key = array.count)
return m, array.count(m)
def start_video(self, title):
if self.video is not None:
try:
self.video.release()
except Exception, e:
print "[WARNING] Error releasing previous video:", e
self.video = cv2.VideoWriter('recordings/' + "test" + '.mpeg', -1, int(self.camera.capture.get(cv2.CAP_PROP_FPS)),
(int(self.camera.capture.get(3)), int(self.camera.capture.get(4))))
import cv2
from vision import Camera
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("pitch", help="[0] Main pitch, [1] Secondary pitch")
args = parser.parse_args()
pitch_number = int(args.pitch)
capture = Camera(pitch=pitch_number)
while True:
frame = capture.get_frame()
cv2.imwrite("currBg_" + str(pitch_number) + ".png", frame)
#cv2.imwrite("../currBg_" + str(pitch_number) + ".png", frame)
cv2.imshow('Background view', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
print "done"
