3_eyetracker.py

import cv2
import mss
import time
import numpy as np

from datetime import datetime
from collections import deque
from Gaze import Detector, Predictor
from Models import FullModel
from utils import get_config, clamp_value


# Read config.ini file
SETTINGS, COLOURS, EYETRACKER, TF = get_config("config.ini")

with mss.mss() as sct:
    mon = sct.monitors[EYETRACKER["monitor_num"]]
    w, h = mon["width"], mon["height"]
    monitor = {
        "top": mon["top"],
        "left": mon["left"],
        "width": w,
        "height": h,
        "mon": EYETRACKER["monitor_num"],
    }

    # Load trained model
    detector = Detector(output_size=SETTINGS["image_size"])
    predictor = Predictor(
        FullModel,
        model_data="trained_models/eyetracking_model.pt",
        config_file="trained_models/eyetracking_config.json",
    )
    screen_errors = region_map = np.load("trained_models/eyetracking_errors.npy")

    track_x = deque(
        [0] * SETTINGS["avg_window_length"], maxlen=SETTINGS["avg_window_length"]
    )
    track_y = deque(
        [0] * SETTINGS["avg_window_length"], maxlen=SETTINGS["avg_window_length"]
    )
    track_error = deque(
        [0] * (SETTINGS["avg_window_length"] * 2),
        maxlen=SETTINGS["avg_window_length"] * 2,
    )

    videoWriter = None
    if EYETRACKER["write_to_disk"]:
        # FIXME: fix video codec on windows
        date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        fourcc = cv2.VideoWriter_fourcc(*"avc1")
        videoWriter = cv2.VideoWriter(
            "media/recordings/{}.mp4".format(date_time),
            fourcc,
            EYETRACKER["tracker_frame_rate"],
            (
                int(w * EYETRACKER["screen_capture_scale"]),
                int(h * EYETRACKER["screen_capture_scale"]),
            ),
        )

    last_time = time.time()
    while True:
        if cv2.waitKey(1) & 0xFF == 27:  # wait for escape key
            detector.close()
            cv2.destroyAllWindows()
            break

        cur_time = time.time()

        if (cur_time - last_time) >= 1 / EYETRACKER["tracker_frame_rate"]:
            fps = 1 / (cur_time - last_time)
            last_time = cur_time

            # Get camera data
            l_eye, r_eye, face, face_align, head_pos, angle = detector.get_frame()

            # Get screenshot
            screenshot = np.array(sct.grab(monitor))
            overlay = screenshot.copy()

            # Overlays
            x_hat, y_hat = predictor.predict(
                face, l_eye, r_eye, head_pos, head_angle=angle
            )

            track_x.append(x_hat)
            track_y.append(y_hat)

            x_hat_clamp = clamp_value(x_hat, w)
            y_hat_clamp = clamp_value(y_hat, h)
            error = screen_errors[int(x_hat_clamp) - 1][int(y_hat_clamp) - 1]
            track_error.append(error * 0.75)

            weights = np.arange(1, SETTINGS["avg_window_length"] + 1)
            weights_error = np.arange(1, (SETTINGS["avg_window_length"] * 2) + 1)

            cv2.circle(
                overlay,
                (
                    int(np.average(track_x, weights=weights)),
                    int(np.average(track_y, weights=weights)),
                ),
                int(np.average(track_error, weights=weights_error)),
                (255, 255, 255, 50),
                -1,
            )

            cv2.circle(
                screenshot,
                (
                    int(np.average(track_x, weights=weights)),
                    int(np.average(track_y, weights=weights)),
                ),
                int(np.average(track_error, weights=weights_error)),
                COLOURS["green"],
                5,
            )

            screenshot = cv2.addWeighted(overlay, 0.3, screenshot, 0.7, 0)

            cv2.putText(
                screenshot,
                "fps: {}".format(round(fps, 2)),
                (0, h),
                cv2.FONT_HERSHEY_PLAIN,
                2,
                COLOURS["green"],
            )

            if EYETRACKER["show_webcam"]:
                large_size = SETTINGS["image_size"] * 2

                screenshot[0:large_size, 0:large_size, 0:3] = cv2.resize(
                    face_align, (large_size, large_size)
                )

                screenshot[
                    0:large_size,
                    large_size : large_size * 2,
                    0:3,
                ] = cv2.resize(
                    np.repeat(head_pos[:, :, np.newaxis], 3, axis=2),
                    (large_size, large_size),
                )

                screenshot[
                    large_size : large_size * 2,
                    0:large_size,
                    0:3,
                ] = cv2.resize(l_eye, (large_size, large_size))

                screenshot[
                    large_size : large_size * 2,
                    large_size : large_size * 2,
                    0:3,
                ] = cv2.resize(r_eye, (large_size, large_size))

            # Resize and write frame
            screenshot = cv2.resize(
                screenshot,
                (
                    int(w * EYETRACKER["screen_capture_scale"]),
                    int(h * EYETRACKER["screen_capture_scale"]),
                ),
            )

            cv2.imshow("Eyetracker", screenshot)
            # cv2.imshow("Webcam", face)

            if EYETRACKER["write_to_disk"]:
                videoWriter.write(screenshot)

    if EYETRACKER["write_to_disk"]:
        videoWriter.release()