imu_timeline.py

import pathlib

import numpy as np
import OpenGL.GL as gl
from pyglui import ui, pyfontstash
from pyglui.cygl import utils as cygl_utils

import gl_utils
from plugin import Plugin
from pupil_recording import PupilRecording


class IMURecording:
    DTYPE = np.dtype(
        [
            ("gyro_x", "<f4"),
            ("gyro_y", "<f4"),
            ("gyro_z", "<f4"),
            ("accel_x", "<f4"),
            ("accel_y", "<f4"),
            ("accel_z", "<f4"),
        ]
    )

    def __init__(self, path_to_imu_raw: pathlib.Path):
        stem = path_to_imu_raw.stem
        self.path_raw = path_to_imu_raw
        self.path_ts = path_to_imu_raw.with_name(stem + "_timestamps.npy")
        self.load()

    def load(self):
        if not self.path_raw.exists() and self.path_ts.exists():
            self.ts = np.empty(0, dtype=np.float64)
            self.raw = np.empty(0, dtype=self.DTYPE).view(np.recarray)
            return

        self.ts = np.load(str(self.path_ts))
        self.raw = np.fromfile(str(self.path_raw), dtype=self.DTYPE).view(np.recarray)
        num_ts_during_init = self.ts.size - self.raw.size
        if num_ts_during_init > 0:
            self.ts = self.ts[num_ts_during_init:]


class IMUTimeline(Plugin):
    IMU_PATTERN_RAW = r"^extimu ps(\d+).raw"
    CMAP = {
        "gyro_x": cygl_utils.RGBA(0.12156, 0.46666, 0.70588, 1.0),
        "gyro_y": cygl_utils.RGBA(1.0, 0.49803, 0.05490, 1.0),
        "gyro_z": cygl_utils.RGBA(0.17254, 0.62745, 0.1725, 1.0),
        "accel_x": cygl_utils.RGBA(0.83921, 0.15294, 0.15686, 1.0),
        "accel_y": cygl_utils.RGBA(0.58039, 0.40392, 0.74117, 1.0),
        "accel_z": cygl_utils.RGBA(0.54901, 0.33725, 0.29411, 1.0),
    }
    NUMBER_SAMPLES_TIMELINE = 4000
    TIMELINE_LINE_HEIGHT = 16

    @classmethod
    def parse_pretty_class_name(cls) -> str:
        return "IMU Timeline"

    def __init__(self, g_pool):
        super().__init__(g_pool)
        rec = PupilRecording(g_pool.rec_dir)
        imu_files = sorted(rec.files().filter_patterns(self.IMU_PATTERN_RAW))
        imu_recs = [IMURecording(imu_file) for imu_file in imu_files]
        self.data_ts = np.concatenate([rec.ts for rec in imu_recs])
        self.data_raw = np.concatenate([rec.raw for rec in imu_recs])

    def init_ui(self):
        self.gyro_timeline = ui.Timeline(
            "gyro",
            self.draw_raw_gyro,
            self.draw_legend_gyro,
            self.TIMELINE_LINE_HEIGHT * (len(self.CMAP) // 2),
        )
        self.accel_timeline = ui.Timeline(
            "accel",
            self.draw_raw_accel,
            self.draw_legend_accel,
            self.TIMELINE_LINE_HEIGHT * (len(self.CMAP) // 2),
        )
        self.g_pool.user_timelines.append(self.gyro_timeline)
        self.g_pool.user_timelines.append(self.accel_timeline)
        self.glfont = pyfontstash.fontstash.Context()
        self.glfont.add_font("opensans", ui.get_opensans_font_path())
        self.glfont.set_color_float((1.0, 1.0, 1.0, 0.8))
        self.glfont.set_align_string(v_align="right", h_align="top")

    def deinit_ui(self):
        self.g_pool.user_timelines.remove(self.gyro_timeline)
        self.g_pool.user_timelines.remove(self.accel_timeline)
        del self.gyro_timeline
        del self.accel_timeline
        del self.glfont

    def draw_raw_gyro(self, width, height, scale):
        keys = ["gyro_x", "gyro_y", "gyro_z"]
        y_limits = [-10, 10]  # TODO: ADJUST
        self._draw_raw_grouped(keys, y_limits, width, height, scale)

    def draw_raw_accel(self, width, height, scale):
        keys = ["accel_x", "accel_y", "accel_z"]
        y_limits = [-1, 1]  # TODO: ADJUST
        self._draw_raw_grouped(keys, y_limits, width, height, scale)

    def _draw_raw_grouped(self, keys, y_limits, width, height, scale):
        ts_min = self.g_pool.timestamps[0]
        ts_max = self.g_pool.timestamps[-1]
        data = self.data_raw[keys]
        with gl_utils.Coord_System(ts_min, ts_max, *y_limits):
            for key in keys:
                data_keyed = data[key]
                points = list(zip(self.data_ts, data_keyed))
                cygl_utils.draw_points(points, size=2.0 * scale, color=self.CMAP[key])

    def draw_legend_gyro(self, width, height, scale):
        self._draw_legend_grouped(["gyro_x", "gyro_y", "gyro_z"], width, height, scale)

    def draw_legend_accel(self, width, height, scale):
        self._draw_legend_grouped(
            ["accel_x", "accel_y", "accel_z"], width, height, scale
        )

    def _draw_legend_grouped(self, labels, width, height, scale):
        self.glfont.set_size(self.TIMELINE_LINE_HEIGHT * 0.8 * scale)
        pad = width * 2 / 3
        for label in labels:
            color = self.CMAP[label]
            self.glfont.draw_text(width, 0, label)

            cygl_utils.draw_polyline(
                [
                    (pad, self.TIMELINE_LINE_HEIGHT / 2),
                    (width / 4, self.TIMELINE_LINE_HEIGHT / 2),
                ],
                color=color,
                line_type=gl.GL_LINES,
                thickness=4.0 * scale,
            )
            gl.glTranslatef(0, self.TIMELINE_LINE_HEIGHT * scale, 0)