class Vision:
def __init__(self, pitch_num, stdout, reset_pitch_size, reset_thresh,
scale, colour_order, render_tlayers, file_input=None):
self.running = True
self.connected = False
self.scale = scale
self.stdout = stdout
self._logger = Logger('vision_errors.log')
if file_input is None:
self.cam = Camera(prop_set = {"width": 720, "height": 540})
else:
file_type = 'video'
if file_input.endswith(('jpg', 'png')):
file_type = 'image'
self.cam = VirtualCamera(file_input, file_type)
try:
calibration_path = os.path.join('calibration', 'pitch{0}'.format(pitch_num))
self.cam.loadCalibration(os.path.join(sys.path[0], calibration_path))
except TypeError:
error_msg = 'Calibration file not found.'
self._logger.log(error_msg)
print error_msg
self.cropper = Cropper(pitch_num=pitch_num, reset_pitch=reset_pitch_size)
self.processor = Processor(pitch_num, reset_pitch_size, reset_thresh, scale)
if self.cropper.is_ready():
self.gui = Gui(self.cropper.pitch_size)
else:
self.gui = Gui()
self.threshold_gui = ThresholdGui(self.processor, self.gui, pitch_num = pitch_num)
self.detection = Detection(self.gui, self.processor, colour_order, scale, pitch_num,
render_tlayers=render_tlayers)
self.event_handler = self.gui.get_event_handler()
self.event_handler.add_listener('q', self.quit)
while self.running:
try:
if not self.stdout:
self.connect()
else:
self.connected = True
if self.cropper.is_ready():
#self.output_pitch_size()
self.detection.set_coord_rect(self.cropper.get_coord_rect())
self.detection.set_pitch_dims(self.cropper.pitch_size)
self.processor.set_crop_rect(self.cropper.get_crop_rect())
self.gui.set_show_mouse(False)
else:
self.event_handler.set_click_listener(self.set_next_pitch_corner)
while self.running:
self.process_frame()
except socket.error:
self.connected = False
# If the rest of the system is not up yet/gets quit,
# just wait for it to come available.
time.sleep(1)
error_msg = 'Connection error, sleeping 1s...'
self._logger.log(error_msg)
print error_msg
self.process_frame()
if not self.stdout:
self.socket.close()
def process_frame(self):
"""Get frame, detect objects and display frame
"""
# This is where calibration comes in
if self.cam.getCameraMatrix is None:
frame = self.cam.getImage()
else:
frame = self.cam.getImageUndistort()
self.processor.preprocess(frame, self.scale)
if self.cropper.is_ready():
self.gui.update_layer('raw', self.processor.get_bgr_frame())
else:
self.gui.update_layer('raw', frame)
if self.cropper.is_ready():
entities = self.detection.detect_objects()
self.output_entities(entities)
self.gui.process_update()
def set_next_pitch_corner(self, where):
self.cropper.set_point(where)
if self.cropper.is_ready():
#self.output_pitch_size()
self.processor.set_crop_rect(self.cropper.get_crop_rect())
self.detection.set_pitch_dims(self.cropper.pitch_size)
self.detection.set_coord_rect(self.cropper.get_coord_rect())
self.gui.draw_crosshair(self.cropper.get_coord_rect()[0], 'corner1')
self.gui.draw_crosshair(self.cropper.get_coord_rect()[1], 'corner2')
self.cropper.get_coord_rect()[0]
self.gui.set_show_mouse(False)
self.gui.update_layer('corner', None)
else:
self.gui.draw_crosshair(where, 'corner')
def output_pitch_size(self):
self.send('{0} {1}\n'.format(PITCH_SIZE_BIT, self.processor.pitch_points_string))
def output_entities(self, entities):
if not self.connected or not self.cropper.is_ready():
return
self.send('{0} '.format(ENTITY_BIT))
for i in range(0, 4):
entity = entities[i]
x, y = entity.get_coordinates()
angle = -1 if entity.get_angle() is None else entity.get_angle()
self.send('{0} {1} {2} '.format(x, y, angle))
x, y = entities[BALL].get_coordinates()
self.send('{0} {1} '.format(x, y))
self.send(str(int(time.time() * 1000)) + "\n")
def send(self, string):
if self.stdout:
sys.stdout.write(string)
else:
self.socket.send(string)
def connect(self):
print('Attempting to connect...')
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.socket.connect((HOST, PORT))
self.connected = True
print('Successfully connected.')
def quit(self):
self.running = False
