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
class Vision: def __init__(self, pitchnum, stdout, sourcefile, resetPitchSize): self.running = True self.connected = False self.stdout = stdout if sourcefile is None: self.cap = Camera() else: filetype = 'video' if sourcefile.endswith(('jpg', 'png')): filetype = 'image' self.cap = VirtualCamera(sourcefile, filetype) calibrationPath = os.path.join('calibration', 'pitch{0}'.format(pitchnum)) self.cap.loadCalibration(os.path.join(sys.path[0], calibrationPath)) self.gui = Gui() self.threshold = Threshold(pitchnum) self.thresholdGui = ThresholdGui(self.threshold, self.gui) self.preprocessor = Preprocessor(resetPitchSize) self.features = Features(self.gui, self.threshold) eventHandler = self.gui.getEventHandler() eventHandler.addListener('q', self.quit) while self.running: try: if not self.stdout: self.connect() else: self.connected = True if self.preprocessor.hasPitchSize: self.outputPitchSize() self.gui.setShowMouse(False) else: eventHandler.setClickListener(self.setNextPitchCorner) while self.running: self.doStuff() 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) # Strange things seem to happen to X sometimes if the # display isn't updated for a while self.doStuff() if not self.stdout: self.socket.close() 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 def quit(self): self.running = False def doStuff(self): if self.cap.getCameraMatrix is None: frame = self.cap.getImage() else: frame = self.cap.getImageUndistort() frame = self.preprocessor.preprocess(frame) self.gui.updateLayer('raw', frame) ents = self.features.extractFeatures(frame) self.outputEnts(ents) self.gui.loop() def setNextPitchCorner(self, where): self.preprocessor.setNextPitchCorner(where) if self.preprocessor.hasPitchSize: print("Pitch size: {0!r}".format(self.preprocessor.pitch_size)) self.outputPitchSize() self.gui.setShowMouse(False) self.gui.updateLayer('corner', None) else: self.gui.drawCrosshair(where, 'corner') def outputPitchSize(self): print(self.preprocessor.pitch_size) self.send('{0} {1} {2} \n'.format( PITCH_SIZE_BIT, self.preprocessor.pitch_size[0], self.preprocessor.pitch_size[1])) def outputEnts(self, ents): # Messyyy if not self.connected or not self.preprocessor.hasPitchSize: return self.send("{0} ".format(ENTITY_BIT)) for name in ['yellow', 'blue', 'ball']: entity = ents[name] x, y = entity.coordinates() # The rest of the system needs (0, 0) at the bottom left if y != -1: y = self.preprocessor.pitch_size[1] - y if name == 'ball': self.send('{0} {1} '.format(x, y)) else: angle = 360 - (((entity.angle() * (180/math.pi)) - 360) % 360) self.send('{0} {1} {2} '.format(x, y, angle)) self.send(str(int(time.time() * 1000)) + " \n") def send(self, string): if self.stdout: sys.stdout.write(string) else: self.socket.send(string)