import freehead as fh
import numpy as np
import pygame
pygame.init()
pygame.display.set_mode((300, 300))
othread = fh.OptotrakThread()
othread.start()
othread.started_running.wait()
athread = fh.ArduinoThread()
athread.start()
athread.started_running.wait()
probe = fh.FourMarkerProbe()
n_samples = 5
led_positions = np.full((n_samples, 3), np.nan)
rig_calibration_indices = np.round(np.linspace(0, 254, n_samples)).astype(np.int)
for i, ci in enumerate(rig_calibration_indices):
athread.write_uint8(ci, 100, 0, 0)
while True:
fh.wait_for_keypress(pygame.K_SPACE)
probe_rotation, probe_tip = probe.solve(othread.current_sample.copy()[15:27].reshape(4,3))
if fh.anynan(probe_rotation):
continue
led_positions[i, :] = probe_tip[:]
break
def create_helmet(self):
head_measurement_points = [
'head straight', 'nasion', 'inion', 'right eye'
]
for i, measurement_point in enumerate(head_measurement_points):
print('Press space to measure: ' + measurement_point)
# light up an led to signal which measurement is going on
signal_led = 255
signal_length = 1
while True:
self.athread.write_uint8(
signal_led, 255, 255,
255) # bright light to start and see something
fh.wait_for_keypress(pygame.K_SPACE)
current_sample = self.othread.current_sample.copy()
helmet_leds = current_sample[HELMET].reshape((4, 3))
if np.any(np.isnan(helmet_leds)):
print('Helmet LEDs not all visible. Try again.')
self.athread.write_uint8(signal_led, 255, 0,
0) # red light for failure
time.sleep(signal_length)
continue
if i == I_BARY:
helmet = fh.Rigidbody(helmet_leds)
self.athread.write_uint8(signal_led, 0, 255,
0) # green light for success
time.sleep(signal_length)
break
else:
_, probe_tip = fh.FourMarkerProbe().solve(
current_sample[PROBE].reshape((4, 3)))
if np.any(np.isnan(probe_tip)):
print('Probe not visible. Try again.')
self.athread.write_uint8(signal_led, 255, 0,
0) # red light for failure
time.sleep(signal_length)
continue
helmet.add_reference_points(helmet_leds, probe_tip)
# replace eye measurement
if i == I_EYE:
nasion_to_inion = fh.to_unit(
helmet.ref_points[I_INION, :] -
helmet.ref_points[I_NASION, :])
# estimate eye at 15 mm inwards from probe in nasion inion direction
estimated_eye_position = helmet.ref_points[
I_EYE, :] + 15 * nasion_to_inion
# replace measured value with estimation
helmet.ref_points[I_EYE, :] = estimated_eye_position
self.athread.write_uint8(signal_led, 0, 255,
0) # green light for success
time.sleep(signal_length)
break
self.helmet = helmet
print('Helmet creation done.')
self.athread.write_uint8(255, 0, 0, 0) # turn off leds