path: root/src/argaze/TobiiGlassesPro2/TobiiData.py

#!/usr/bin/env python

from dataclasses import dataclass
import threading
import uuid
import gzip
import json
import time
import queue

from argaze import DataStructures
from argaze.TobiiGlassesPro2 import TobiiNetworkInterface

@dataclass
class DirSig():
    """Define dir sig data (dir sig)."""

    dir: int # meaning ?
    sig: int # meaning ?

@dataclass
class PresentationTimeStamp():
    """Define presentation time stamp data (pts)."""

    value: int

@dataclass
class EventSynch():
    """Define event synch data (evts)."""

    value: int # meaning ?

@dataclass
class Event():
    """Define event data (ets type tag)."""

    ets: int # meaning ?
    type: str
    tag: str # dict ?

@dataclass
class Accelerometer():
    """Define accelerometer data (ac)."""

    value: tuple((float, float, float))

@dataclass
class Gyroscope():
    """Define gyroscope data (gy)."""

    value: tuple((float, float, float))

@dataclass
class PupilCenter():
    """Define pupil center data (gidx pc eye)."""

    index: int
    value: tuple((float, float, float))
    eye: str # 'right' or 'left'

@dataclass
class PupilDiameter():
    """Define pupil diameter data (gidx pd eye)."""

    index: int
    value: float
    eye: str # 'right' or 'left'

@dataclass
class GazeDirection():
    """Define gaze direction data (gidx gd eye)."""

    index: int
    value: tuple((float, float, float))
    eye: str # 'right' or 'left'

@dataclass
class GazePosition():
    """Define gaze position data (gidx l gp)."""

    index: int
    l: str # ?
    value: tuple((float, float))

@dataclass
class GazePosition3D():
    """Define gaze position 3D data (gidx gp3)."""

    index: int
    value: tuple((float, float))

@dataclass
class MarkerPosition():
    """Define marker data (marker3d marker2d)."""

    value_3d: tuple((float, float, float))
    value_2d: tuple((float, float))

class TobiiJsonDataParser():

    def parse_dir_sig(self, json_data):

        return DirSig(json_data['dir'], json_data['sig'])

    def parse_pts(self, json_data):

        return PresentationTimeStamp(json_data['pts'])

    def parse_event_synch(self, json_data):

        return EventSynch(json_data['evts'])

    def parse_event(self, json_data):

        return Event(json_data['ets'], json_data['type'], json_data['tag'])

    def parse_accelerometer(self, json_data):

        return Accelerometer(json_data['ac'])

    def parse_gyroscope(self, json_data):

        return Gyroscope(json_data['gy'])

    def parse_pupil_center(self, gaze_index, json_data):

        return PupilCenter(gaze_index, json_data['pc'], json_data['eye'])

    def parse_pupil_diameter(self, gaze_index, json_data):

        return PupilDiameter(gaze_index, json_data['pd'], json_data['eye'])

    def parse_gaze_direction(self, gaze_index, json_data):

        return GazeDirection(gaze_index, json_data['gd'], json_data['eye'])

    def parse_gaze_position(self, gaze_index, json_data):

        return GazePosition(gaze_index, json_data['l'], json_data['gp'])

    def parse_gaze_position_3d(self, gaze_index, json_data):

        return GazePosition3D(gaze_index, json_data['gp3'])

    def parse_marker_position(self, json_data):

        return MarkerPosition(json_data['marker3d'], json_data['marker2d'])

    def parse_pupil_or_gaze(self, json_data):

        gaze_index = json_data.pop('gidx')

        # parse pupil or gaze data depending second json key
        second_key = next(iter(json_data))

        parse_map = {
            'pc': self.parse_pupil_center,
            'pd': self.parse_pupil_diameter,
            'gd': self.parse_gaze_direction,
            'l': self.parse_gaze_position,
            'gp3': self.parse_gaze_position_3d
        }

        return parse_map[second_key](gaze_index, json_data)

    def parse_data(self, json_data):

        # parse data depending first json key
        first_key = next(iter(json_data))

        parse_map = {
            'dir': self.parse_dir_sig,
            'pts': self.parse_pts,
            'evts': self.parse_event_synch,
            'ets': self.parse_event,
            'ac': self.parse_accelerometer,
            'gy': self.parse_gyroscope,
            'gidx': self.parse_pupil_or_gaze,
            'marker3d': self.parse_marker_position
        }

        return parse_map[first_key](json_data)

class TobiiDataSegment():
    """Handle Tobii Glasses Pro 2 segment data file."""

    def __init__(self, segment_data_path, start_timestamp = 0, end_timestamp = None):
        """Load segment data from segment directory then parse and register each recorded dataflow as a TimeStampedBuffer member of the TobiiSegmentData instance."""

        self.__segment_data_path = segment_data_path

        self.__vts_offset = 0
        self.__vts_ts = -1

        self.__json_data_parser = TobiiJsonDataParser()

        self.__ts_data_buffer_dict = {
            'DirSig': DataStructures.TimeStampedBuffer(),
            'PresentationTimeStamp': DataStructures.TimeStampedBuffer(),
            'EventSynch': DataStructures.TimeStampedBuffer(),
            'Event': DataStructures.TimeStampedBuffer(),
            'Accelerometer': DataStructures.TimeStampedBuffer(),
            'Gyroscope': DataStructures.TimeStampedBuffer(),
            'PupilCenter': DataStructures.TimeStampedBuffer(),
            'PupilDiameter': DataStructures.TimeStampedBuffer(),
            'GazeDirection': DataStructures.TimeStampedBuffer(),
            'GazePosition': DataStructures.TimeStampedBuffer(),
            'GazePosition3D': DataStructures.TimeStampedBuffer(),
            'MarkerPosition': DataStructures.TimeStampedBuffer()
        }

        # define a decoder function
        def decode(json_data):

            # accept only valid data (e.g. with status value equal to 0)
            if json_data.pop('s', -1) == 0:

                # convert timestamp
                ts = json_data.pop('ts')

                # watch for vts data to offset timestamps
                try:
                    self.__vts_offset = json_data['vts']
                    self.__vts_ts = ts

                    return True # continue

                except KeyError:
                    pass

                # ignore data before first vts entry
                if self.__vts_ts == -1:
                    return True # continue

                ts -= self.__vts_ts
                ts += self.__vts_offset

                # ignore timestamps out of the given time range
                if ts < start_timestamp:
                    return True # continue

                if ts >= end_timestamp:
                    return False # stop

                # convert json data into data object
                data_object = self.__json_data_parser.parse_data(json_data)
                data_object_type = type(data_object).__name__

                # store data object into dedicated timestamped buffer
                self.__ts_data_buffer_dict[data_object_type][ts] = data_object

            return True # continue

        # start loading
        with gzip.open(self.__segment_data_path) as f:

            for item in f:
                if not json.loads(item.decode('utf-8'), object_hook=decode):
                    break

    def __getitem__(self, key):
        return self.__ts_data_buffer_dict[key]

    def keys(self):
        """Get all registered data keys"""
        return list(self.__ts_data_buffer_dict.keys())

    def get_path(self):
        return self.__segment_data_path

class TobiiDataStream(threading.Thread):
    """Capture Tobii Glasses Pro 2 data stream in separate thread."""

    def __init__(self, network_interface: TobiiNetworkInterface.TobiiNetworkInterface):
        """Initialise thread super class as a deamon dedicated to data reception."""

        threading.Thread.__init__(self)
        threading.Thread.daemon = True
        
        self.__network = network_interface
        self.__data_socket = self.__network.make_socket()

        self.__data_queue = queue.Queue()
        
        self.__stop_event = threading.Event()
        self.__read_lock = threading.Lock() 

        # prepare keep alive message
        self.__keep_alive_msg = "{\"type\": \"live.data.unicast\", \"key\": \""+ str(uuid.uuid4()) +"\", \"op\": \"start\"}"
        self.__keep_alive_thread = threading.Thread(target = self.__keep_alive)
        self.__keep_alive_thread.daemon = True

    def __del__(self):
        """Stop data reception before destruction."""

        self.close()

    def __keep_alive(self):
        """Maintain connection."""

        while not self.__stop_event.isSet():

            self.__network.send_keep_alive_msg(self.__data_socket, self.__keep_alive_msg)

            time.sleep(1)

    def open(self):
        """Start data reception."""

        self.__first_ts = 0
        self.__keep_alive_thread.start()
        threading.Thread.start(self)

    def close(self):
        """Stop data reception definitively."""

        self.__stop_event.set()

        threading.Thread.join(self.__keep_alive_thread)
        threading.Thread.join(self)

        self.__data_socket.close()

    def run(self):
        """Store received data into a queue for further reading."""

        while not self.__stop_event.isSet():

            # lock data queue access
            self.__read_lock.acquire()

            # write in data queue
            data = self.__network.grab_data(self.__data_socket)
            json_data = json.loads(data.decode('utf-8'))
            self.__data_queue.put(json_data)

            # unlock data queue access
            self.__read_lock.release()

    def read(self):

        json_data_parser = TobiiJsonDataParser()

        # create a dictionary of timestamped data buffers
        ts_data_buffer_dict = {
            'DirSig': DataStructures.TimeStampedBuffer(),
            'PresentationTimeStamp': DataStructures.TimeStampedBuffer(),
            'EventSynch': DataStructures.TimeStampedBuffer(),
            'Event': DataStructures.TimeStampedBuffer(),
            'Accelerometer': DataStructures.TimeStampedBuffer(),
            'Gyroscope': DataStructures.TimeStampedBuffer(),
            'PupilCenter': DataStructures.TimeStampedBuffer(),
            'PupilDiameter': DataStructures.TimeStampedBuffer(),
            'GazeDirection': DataStructures.TimeStampedBuffer(),
            'GazePosition': DataStructures.TimeStampedBuffer(),
            'GazePosition3D': DataStructures.TimeStampedBuffer(),
            'MarkerPosition': DataStructures.TimeStampedBuffer()
        }

        # if the data acquisition thread is not running
        if self.__stop_event.isSet():
            return ts_data_buffer_dict

        # lock data queue access
        self.__read_lock.acquire()

        # read data queue
        while not self.__data_queue.empty():

            json_data = self.__data_queue.get()

            # accept only valid data (e.g. with status value equal to 0)
            if json_data.pop('s', -1) == 0:

                # convert timestamp
                ts = json_data.pop('ts')

                # keep first timestamp to offset all timestamps
                if self.__first_ts == 0:
                    self.__first_ts = ts
                
                ts -= self.__first_ts
                
                # ignore negative timestamp
                if ts < 0:
                    break

                # convert json data into data object
                data_object = json_data_parser.parse_data(json_data)
                data_object_type = type(data_object).__name__

                # store data object into dedicated timestamped buffer
                ts_data_buffer_dict[data_object_type][ts] = data_object
                
        # unlock data queue access
        self.__read_lock.release()

        return ts_data_buffer_dict