Exemple #1
0
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