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)
