Transform eyeball centers from scene camera coodinates to head pose coordinates

transform_eye_centers.py

import argparse
import csv
import logging
import os

import cv2
import numpy as np
import pandas as pd


logger = logging.getLogger(__name__)


def main(export_directory, marker_size, overwrite=False):
    """Process a given recording's exported gaze and head pose data
    handles cases where no gaze_positions.csv or
    head_pose_tacker_poses.csv files could be found
    """
    try:
        logger.info(f"Extracting {export_directory}...")
        process_recording(export_directory, marker_size, overwrite=overwrite)
    except FileNotFoundError as err:
        logger.warning(f"{err}")
        return


def process_recording(export_directory, marker_size, overwrite=False):
    """Parses gaze and head pose data from the export,
    transforms eyeball centers to head pose coordinates,
    writes the results to csv file
    """
    data = load_and_parse_data(export_directory)
    transformed = transform_and_yield(data, marker_size)
    write_results(transformed, export_directory, overwrite)


def load_and_parse_data(export_directory):
    gaze_data = pd.read_csv(os.path.join(export_directory, "gaze_positions.csv"))
    head_pose_data = pd.read_csv(
        os.path.join(export_directory, "head_pose_tacker_poses.csv")
    )

    target_rows = find_closest(gaze_data["gaze_timestamp"], head_pose_data["timestamp"])
    target_axes = ("x", "y", "z")
    target_variables = ("rotation", "translation")  # head pose variables

    eye0_centers = (
        gaze_data[["eye_center0_3d_" + ax for ax in target_axes]]
        .iloc[target_rows]
        .to_numpy()
    )
    eye1_centers = (
        gaze_data[["eye_center1_3d_" + ax for ax in target_axes]]
        .iloc[target_rows]
        .to_numpy()
    )
    camera_poses = head_pose_data[
        [f"{var}_{ax}" for var in target_variables for ax in target_axes]
    ].to_numpy()

    timestamps = gaze_data["gaze_timestamp"].iloc[target_rows].to_numpy()

    return timestamps, eye0_centers, eye1_centers, camera_poses


def find_closest(target, source):
    """Find indices of closest `target` elements for elements in `source`.
    `target` is assumed to be sorted. Result has same shape as `source`.
    Implementation taken from:
    https://stackoverflow.com/questions/8914491/finding-the-nearest-value-and-return-the-index-of-array-in-python/8929827#8929827
    helper function to find world indices
    """
    target = np.asarray(target)  # fixes https://github.com/pupil-labs/pupil/issues/1439
    idx = np.searchsorted(target, source)
    idx = np.clip(idx, 1, len(target) - 1)
    left = target[idx - 1]
    right = target[idx]
    idx -= source - left < right - source
    return idx


def transform_and_yield(data, marker_size):
    timestamps, eye0_centers, eye1_centers, camera_poses = data
    for ts, eye0, eye1, camera in zip(
        timestamps, eye0_centers, eye1_centers, camera_poses
    ):
        x0, y0, z0 = transform_centers_by_pose(eye0, camera, marker_size)
        x1, y1, z1 = transform_centers_by_pose(eye1, camera, marker_size)
        yield ts, x0, y0, z0, x1, y1, z1


def camera_pose_is_nan(camera_pose):
    return camera_pose is None or (np.isnan(camera_pose)).any()


def split_pose(pose):
    pose = np.array(pose, dtype=np.float32)
    assert pose.size == 6
    rotation = pose.ravel()[0:3]
    translation = pose.ravel()[3:6]
    return rotation, translation


def transform_centers_by_pose(centers_scene_cam, head_pose, marker_size):
    """
    Transform eyeball center estimates from scene camera coordinates to head pose coordinates
    :param centers_scene_cam: eyeball centers, shape: (3,)
    :param head_pose: pose of scene cam in head pose coordinates, shape: (6,)
    :param marker_size: physical size of the marker (float)
    :return: eyeball centers in head pose coordinates, shape: (3,)
    """
    centers_scene_cam = (
        np.asarray(centers_scene_cam, dtype=np.float64) / marker_size
    )  # scale to head pose units

    if camera_pose_is_nan(head_pose):
        return np.full((len(centers_scene_cam), 3), np.nan)

    rot_eye_to_head_pose, trans_eye_to_head_pose = split_pose(head_pose)
    rotmat_eye_to_head_pose = cv2.Rodrigues(rot_eye_to_head_pose)[0]

    points_3d_cam2 = (
        np.dot(rotmat_eye_to_head_pose, centers_scene_cam.T).T + trans_eye_to_head_pose
    )
    return points_3d_cam2 * marker_size  # scale to marker size (mm)


def csv_header():
    """CSV header fields"""
    return (
        "timestamp",
        "eye_center0_x",
        "eye_center0_y",
        "eye_center0_z",
        "eye_center1_x",
        "eye_center1_y",
        "eye_center1_z",
    )


def write_results(data, export_directory, overwrite):
    csv_out_path = os.path.join(export_directory, "eyeball_centers_in_head_pose.csv")
    if os.path.exists(csv_out_path):
        if not overwrite:
            logger.warning(f"{csv_out_path} exists already! Not overwriting.")
            return
        else:
            logger.warning(f"{csv_out_path} exists already! Overwriting.")
    else:
        logger.info(f"Writing to csv {csv_out_path}...")

    with open(csv_out_path, "w") as csv_file:
        writer = csv.writer(csv_file, dialect=csv.unix_dialect)
        writer.writerow(csv_header())
        writer.writerows(data)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(
        description=(
            "Transform eyeball center coordinates to head pose tracker "
            "coordinates using a given recording's gaze and head pose export. "
            "The resulting csv file will be saved in the corresponding "
            "export folder."
        )
    )
    # Positional arguments
    parser.add_argument(
        "export_directory",
        help="One export directory. Must contain exported gaze and head pose export",
    )
    parser.add_argument(
        "marker_size",
        type=float,
        help="Marker size (mm). Required for accurate transformation of eyeball "
        "centers from scene camera to head pose coordinates",
    )
    # Optional argument
    parser.add_argument(
        "-f",
        "--overwrite",
        action="store_true",
        help=(
            "Usually, the command refuses to overwrite existing csv files. "
            "This flag disables these checks."
        ),
    )
    args = parser.parse_args()
    main(
        export_directory=args.export_directory,
        marker_size=args.marker_size,
        overwrite=args.overwrite,
    )