detect.py

import platform

import av
import numpy as np
import pandas as pd
from pupil_detectors import Detector2D, Detector3D, __version__


def luma_component(av_frame) -> np.ndarray:
    YUV = av_frame.planes
    luma_plane = YUV[0]
    luma = np.frombuffer(luma_plane, dtype=np.uint8)
    try:
        luma.shape = av_frame.height, av_frame.width
    except ValueError:
        luma.shape = -1, luma_plane.line_size
        luma = np.ascontiguousarray(luma[:, : av_frame.width])
    return luma

def main():
    print("Pupil Detector version:", __version__)

    detect_2d = Detector2D()
    detect_3d = Detector3D()

    container = av.open("pupil_right.mp4")
    time_base = container.streams.video[0].time_base

    results = []
    for frame in container.decode(video=0):
        ts = float(frame.pts * time_base)
        img = luma_component(frame)

        result_2d = detect_2d.detect(img)
        result_3d = detect_3d.detect(img, ts)

        results.append(
            {
                "ts": ts,
                "kind": "2d",
                "confidence": result_2d["confidence"],
            }
        )
        results.append(
            {
                "ts": ts,
                "kind": "3d",
                "confidence": result_3d["confidence"],
            }
        )

    result_file = f"results_{platform.platform()}.csv"
    pd.DataFrame(results).to_csv(result_file)


if __name__ == "__main__":
    main()

vis.py

from pathlib import Path

import pandas as pd
import seaborn as sns


sns.set()

files = {}
for file in Path().glob("results_*.csv"):
    plat = file.stem.replace("results_", "")
    data = pd.read_csv(file, index_col=0)
    files[plat] = data

data = pd.concat(files.values(), keys=files.keys(), names=["platform"])
data = data.reset_index(level=0)

g = sns.FacetGrid(data=data, row="kind", aspect=3)
g.map_dataframe(sns.lineplot, x="ts", y="confidence", hue="platform")
g.add_legend()
g.set_axis_labels("time since start", "confidence")
g.savefig("results.png")