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