Offset Corrected Gaze From Recording

gaze_with_offset_correction.py

from pyglui import ui

import file_methods as fm
import player_methods as pm
from gaze_producer.gaze_producer_base import GazeProducerBase
from methods import denormalize
from player_methods import transparent_circle


class GazeFromRecordingWithOffsetCorrection(GazeProducerBase):
    @classmethod
    def plugin_menu_label(cls) -> str:
        return "Offset Corrected Gaze Data"

    @classmethod
    def gaze_data_source_selection_order(cls) -> float:
        return 3.0

    def __init__(self, g_pool, *, x_offset=0.0, y_offset=0.0):
        super().__init__(g_pool)
        self.__x_offset = x_offset
        self.__y_offset = y_offset
        self.__active_offset = (x_offset, y_offset)
        self.__gaze_from_recording = fm.load_pldata_file(self.g_pool.rec_dir, "gaze")
        self.__publish_with_offsets()

    def get_init_dict(self):
        return {
            **super().get_init_dict(),
            "x_offset": self.__x_offset,
            "y_offset": self.__y_offset,
        }

    def recent_events(self, events):
        super().recent_events(events)
        if self.__active_offset == (self.x_offset, self.y_offset):
            return

        frame = events.get("frame")
        if not frame:
            return

        frame_width_height = frame.img.shape[:-1][::-1]

        points = []
        for datum in events.get("gaze", []):
            if datum["confidence"] < self.g_pool.min_data_confidence:
                continue
            x, y = datum["norm_pos"]
            active_x, active_y = self.__active_offset
            mapped = (x - active_x + self.x_offset, y - active_y + self.y_offset)
            points.append(denormalize(mapped, frame_width_height, flip_y=True))

        for pt in points:
            transparent_circle(
                frame.img, pt, radius=20, color=(0.0, 0.3, 1.0, 0.1), thickness=-1,
            )
            transparent_circle(
                frame.img, pt, radius=20, color=(0.0, 0.3, 1.0, 0.5), thickness=1,
            )

    def __publish_with_offsets(self):
        mutable_data = [
            data._deep_copy_dict() for data in self.__gaze_from_recording.data
        ]
        for datum in mutable_data:
            datum["norm_pos"] = (
                datum["norm_pos"][0] + self.x_offset,
                datum["norm_pos"][1] + self.y_offset,
            )
        serialized_data = [fm.Serialized_Dict(data) for data in mutable_data]
        self.g_pool.gaze_positions = pm.Bisector(
            serialized_data, self.__gaze_from_recording.timestamps
        )
        self._gaze_changed_announcer.announce_new()
        self.__active_offset = (self.x_offset, self.y_offset)

    def init_ui(self):
        super().init_ui()
        self.menu.append(
            ui.Info_Text("Using recorded gaze data with a manual offset correction")
        )
        self.menu.append(
            ui.Slider(
                "x_offset",
                self,
                min=-0.5,
                step=0.01,
                max=0.5,
                label="Manual Correction X",
            )
        )
        self.menu.append(
            ui.Slider(
                "y_offset",
                self,
                min=-0.5,
                step=0.01,
                max=0.5,
                label="Manual Correction Y",
            )
        )
        self.menu.append(ui.Button("Apply Offset", self.__publish_with_offsets))

    @property
    def x_offset(self):
        return self.__x_offset

    @x_offset.setter
    def x_offset(self, val):
        if val != self.__x_offset:
            self.__x_offset = val

    @property
    def y_offset(self):
        return self.__y_offset

    @y_offset.setter
    def y_offset(self, val):
        if val != self.__y_offset:
            self.__y_offset = val


# need to delete this, otherwise the GazeProducerBase gets loaded as user plugin, which
# fails because it is an abstract class
del GazeProducerBase