diff --git a/examples/load_video_as_numpy_array.py b/examples/load_video_as_numpy_array.py index f49f112..d3611ce 100644 --- a/examples/load_video_as_numpy_array.py +++ b/examples/load_video_as_numpy_array.py @@ -20,5 +20,5 @@ # all_frames_np = np.dstack(all_frames) # print('all_frames', all_frames_np.shape) -# need to implement nicely +# can this be implemented nicely? # all_frames_ts = scene.sample(between_two_events).ts diff --git a/examples/make_gaze_overlay_video.py b/examples/make_gaze_overlay_video.py index 3eecaac..eef2f2f 100644 --- a/examples/make_gaze_overlay_video.py +++ b/examples/make_gaze_overlay_video.py @@ -1,14 +1,24 @@ import sys +from fractions import Fraction import cv2 import numpy as np import pupil_labs.neon_recording as nr +import pupil_labs.video as plv rec = nr.load(sys.argv[1]) gaze = rec.gaze eye = rec.eye scene = rec.scene +imu = rec.imu + + +def transparent_rect(img, x, y, w, h): + sub_img = img[y : y + h, x : x + w] + white_rect = np.ones(sub_img.shape, dtype=np.uint8) * 85 + res = cv2.addWeighted(sub_img, 0.5, white_rect, 0.5, 1.0) + img[y : y + h, x : x + w] = res def overlay_image(img, img_overlay, x, y): @@ -30,35 +40,144 @@ def overlay_image(img, img_overlay, x, y): img_crop[:] = img_overlay_crop -video = cv2.VideoWriter( - "video.mp4", cv2.VideoWriter.fourcc("M", "P", "4", "V"), 30, (1600, 1200) +def convert_neon_pts_to_video_pts(neon_pts, neon_time_base, video_time_base): + return int(float(neon_pts * neon_time_base) / video_time_base) + + +fps = 65535 +container = plv.open("video.mp4", mode="w") + +stream = container.add_stream("mpeg4", rate=fps) +stream.width = scene.width +stream.height = scene.height +stream.pix_fmt = "yuv420p" + +neon_time_base = scene.data[0].time_base +video_time_base = Fraction(1, fps) + +avg_neon_pts_size = int(np.mean(np.diff([f.pts for f in scene.data if f is not None]))) +avg_video_pts_size = convert_neon_pts_to_video_pts( + avg_neon_pts_size, neon_time_base, video_time_base ) -try: - start_ts = rec.unique_events["recording.begin"] - end_ts = rec.unique_events["recording.end"] - - my_ts = np.arange(start_ts, end_ts, np.mean(np.diff(scene.ts))) - - for gaze_datum, eye_frame, scene_frame in zip( - gaze.sample(my_ts), eye.sample(my_ts), scene.sample(my_ts) - ): - scn_img = ( - scene_frame.cv2 - if scene_frame is not None - else np.ones((1200, 1600, 3), dtype="uint8") * 128 # gray frames - ) - ey_img = ( - eye_frame.cv2 - if eye_frame is not None - else np.zeros((192, 384, 3), dtype="uint8") # black frames + +start_ts = rec.unique_events["recording.begin"] +end_ts = rec.unique_events["recording.end"] + +avg_frame_dur = np.mean(np.diff(scene.ts)) +pre_ts = np.arange(start_ts, scene.ts[0] - avg_frame_dur, avg_frame_dur) +post_ts = np.arange(scene.ts[-1] + avg_frame_dur, end_ts, avg_frame_dur) + +my_ts = np.concatenate((pre_ts, scene.ts, post_ts)) + +fields = [ + "gyro_x", + "gyro_y", + "gyro_z", + "pitch", + "yaw", + "roll", + "accel_x", + "accel_y", + "accel_z", +] +colors = [ + (208, 203, 228), + (135, 157, 115), + (179, 133, 124), + (101, 118, 223), + (189, 201, 138), + (235, 167, 124), + (93, 197, 128), + (188, 181, 0), + (24, 50, 170), +] +imu_maxes = {} +for field in fields: + imu_maxes[field] = np.max(np.abs(imu[field])) + +ts_rel_max = np.max(imu.ts_rel) + +# gyro_data = dict.fromkeys(fields, []) +gyro_data = {} +for field in fields: + gyro_data[field] = [] + +pts_offset = 0 +video_pts = 0 +reached_video_start = False +for gaze_datum, eye_frame, scene_frame, imu_datum in zip( + gaze.sample(my_ts), eye.sample(my_ts), scene.sample(my_ts), imu.sample(my_ts) +): + scene_image = ( + scene_frame.cv2 + if scene_frame is not None + else np.ones((scene.height, scene.width, 3), dtype="uint8") * 128 # gray frames + ) + eye_image = ( + eye_frame.cv2 + if eye_frame is not None + else np.zeros((eye.height, eye.width, 3), dtype="uint8") # black frames + ) + + overlay_image(scene_image, eye_image, 0, 0) + if gaze_datum: + cv2.circle( + scene_image, (int(gaze_datum.x), int(gaze_datum.y)), 50, (0, 0, 255), 10 ) - overlay_image(scn_img, ey_img, 0, 0) - if gaze_datum: - cv2.circle( - scn_img, (int(gaze_datum.x), int(gaze_datum.y)), 50, (0, 0, 255), 10 + border_cols = [(245, 201, 176), (225, 181, 156), (225, 181, 156)] + transparent_rect(scene_image, 0, 950, scene.width, scene.height - 950) + cv2.line(scene_image, (0, 950), (scene.width, 950), border_cols[0], 2) + cv2.line(scene_image, (0, 950 + 80), (scene.width, 950 + 80), border_cols[1], 2) + cv2.line(scene_image, (0, 950 + 160), (scene.width, 950 + 160), border_cols[2], 2) + if imu_datum: + sep = 0 + for i, field in enumerate(imu_maxes): + if i > 0 and i % 3 == 0: + sep += 80 + + datum_mx = imu_maxes[field] + gyro_data[field].append( + [ + scene.width * imu_datum.ts_rel / ts_rel_max, + -1.0 * imu_datum[field] / datum_mx * 20 + 1000 + sep, + ] + ) + + cv2.polylines( + scene_image, + np.array([gyro_data[field]], dtype=np.int32), + isClosed=False, + color=colors[i], + thickness=2, ) - video.write(scn_img) + frame = plv.VideoFrame.from_ndarray(scene_image, format="bgr24") + if scene_frame is not None: + reached_video_start = True + video_pts = convert_neon_pts_to_video_pts( + scene_frame.pts, neon_time_base, video_time_base + ) + elif reached_video_start and scene_frame is None: + video_pts += avg_video_pts_size + + frame.pts = pts_offset + video_pts + frame.time_base = video_time_base + for packet in stream.encode(frame): + container.mux(packet) + + if scene_frame is not None: + cv2.imshow("frame", scene_image) + if cv2.waitKey(1) & 0xFF == ord("q"): + break + + if not reached_video_start and scene_frame is None: + pts_offset += avg_video_pts_size + +try: + # Flush stream + for packet in stream.encode(): + container.mux(packet) finally: - video.release() + # Close the file + container.close() diff --git a/examples/print_diagnostic_info.py b/examples/print_diagnostic_info.py index 7045c0f..c8499b6 100644 --- a/examples/print_diagnostic_info.py +++ b/examples/print_diagnostic_info.py @@ -13,16 +13,16 @@ pprint.pprint(rec.info) print() -print("scene camera info:") -pprint.pprint(rec.scene_camera) +print("scene camera calibration values:") +pprint.pprint(rec.scene_camera_calibration) print() -print("eye 1 camera info:") -pprint.pprint(rec.eye1_camera) +print("right eye camera calibration values:") +pprint.pprint(rec.right_eye_camera_calibration) print() -print("eye 2 camera info:") -pprint.pprint(rec.eye2_camera) +print("left eye camera calibration values:") +pprint.pprint(rec.left_eye_camera_calibration) print() print("available data streams:") diff --git a/examples/sample_and_iterate_streams.py b/examples/sample_and_iterate_streams.py index cd275ff..708e925 100644 --- a/examples/sample_and_iterate_streams.py +++ b/examples/sample_and_iterate_streams.py @@ -34,8 +34,7 @@ img = scene_frame.rgb gray = scene_frame.gray img_index = scene_frame.index # frame index in the stream - # img_ts = scene_frame.ts # TODO(rob) - fix this: same as rec.streams['scene'].ts[world.index] - img_ts = rec.streams["scene"].ts[scene_frame.index] + img_ts = scene_frame.ts time_into_the_scene_stream = img_ts - rec.start_ts print("scene_frame", img_ts) @@ -49,6 +48,7 @@ # gets me the closest sample within -+0.01s gaze_single_sample = gaze.sample_one(gaze.ts[-20], dt=0.01) +print() print(gaze_single_sample) print() @@ -57,13 +57,14 @@ print() print(gaze[42:45]) +print() # get all samples in a list gaze_samples_list = list(gaze.sample(gaze.ts[:15])) # get all samples as a numpy recarray (gaze/imu) or ndarray of frames (video) -gaze_samples_np = nr.subsampled_to_numpy(gaze.sample(gaze.ts[:15])) +gaze_samples_np = nr.sampled_to_numpy(gaze.sample(gaze.ts[:15])) # NOTE: the following is quite intense on the RAM. -scene_samples_np = nr.subsampled_to_numpy(rec.scene.sample(rec.scene.ts[:15])) +scene_samples_np = nr.sampled_to_numpy(rec.scene.sample(rec.scene.ts[:15])) print(scene_samples_np.shape) diff --git a/src/pupil_labs/neon_recording/__init__.py b/src/pupil_labs/neon_recording/__init__.py index 6e65998..d34265e 100644 --- a/src/pupil_labs/neon_recording/__init__.py +++ b/src/pupil_labs/neon_recording/__init__.py @@ -36,6 +36,6 @@ log.info("NeonRecording: package loaded.") from .neon_recording import load -from .stream.stream import subsampled_to_numpy +from .stream.stream import sampled_to_numpy -__all__ = ["__version__", "load", "subsampled_to_numpy"] +__all__ = ["__version__", "load", "sampled_to_numpy"] diff --git a/src/pupil_labs/neon_recording/calib.py b/src/pupil_labs/neon_recording/calib.py index 93645ce..cd528b2 100644 --- a/src/pupil_labs/neon_recording/calib.py +++ b/src/pupil_labs/neon_recording/calib.py @@ -1,4 +1,5 @@ import pathlib +from dataclasses import dataclass import numpy as np @@ -7,6 +8,13 @@ log = structlog.get_logger(__name__) +@dataclass +class Calibration: + camera_matrix: np.ndarray + distortion_coefficients: np.ndarray + extrinsics_affine_matrix: np.ndarray + + def parse_calib_bin(rec_dir: pathlib.Path): log.debug("NeonRecording: Loading calibration.bin data") @@ -14,12 +22,6 @@ def parse_calib_bin(rec_dir: pathlib.Path): try: with open(rec_dir / "calibration.bin", "rb") as f: calib_raw_data = f.read() - except FileNotFoundError: - raise FileNotFoundError( - f"File not found: {rec_dir / 'calibration.bin'}. Please double check the recording download." - ) - except OSError: - raise OSError(f"Error opening file: {rec_dir / 'calibration.bin'}") except Exception as e: log.exception(f"Unexpected error loading calibration.bin: {e}") raise diff --git a/src/pupil_labs/neon_recording/data_utils.py b/src/pupil_labs/neon_recording/data_utils.py deleted file mode 100644 index f188be1..0000000 --- a/src/pupil_labs/neon_recording/data_utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import json -import pathlib -import typing - -import numpy as np - -from . import structlog - -log = structlog.get_logger(__name__) - - -def load_with_error_check( - func: typing.Callable, fpath: pathlib.Path, err_msg_supp: str -): - log.debug(f"NeonRecording: Loading {fpath.name}") - - try: - res = func(fpath) - if res is not None: - return res - except FileNotFoundError: - raise FileNotFoundError(f"File not found: {fpath.name}. {err_msg_supp}") - except OSError: - raise OSError(f"Error opening file: {fpath.name}.") - except Exception as e: - log.exception(f"Unexpected error loading {fpath.name}: {e}") - raise - - -def load_json_with_error_check(fpath: pathlib.Path): - log.debug(f"NeonRecording: Loading {fpath.name}") - - try: - with open(fpath) as f: - data = json.load(f) - except FileNotFoundError: - raise FileNotFoundError( - f"File not found: {fpath}. Please double check the recording download." - ) - except OSError: - raise OSError(f"Error opening file: {fpath}") - except Exception as e: - log.exception(f"Unexpected error loading info.json: {e}") - raise - else: - return data - - -# obtained from @dom: -# https://github.com/pupil-labs/neon-player/blob/master/pupil_src/shared_modules/pupil_recording/update/neon.py -# -# can return none, as worn data is always 1 for Neon -def load_worn_data(path: pathlib.Path): - log.debug("NeonRecording: Loading worn data.") - - if not (path and path.exists()): - return None - - confidences = np.fromfile(str(path), " float: + return self._start_ts - self.events = [] - self.unique_events = {} + @property + def streams(self): + return self._streams - self._calib = [] - self._version = "" - self._serial = 0 + @property + def info(self): + return self._info - self._gaze_ps1_raw_time_ns = [] - self._gaze_200hz_raw_time_ns = [] - self._gaze_right_ps1_raw_time_ns = [] - self._gaze_ps1_ts = [] - self._gaze_ps1_raw = [] - self._gaze_right_ps1_ts = [] - self._gaze_right_ps1_raw = [] - # self._worn_ps1_raw = [] - self._events_ts_ns = [] + @property + def wearer(self): + return self._wearer - # TODO: save for the end of development - def check(self): - pass + @property + def gaze(self) -> Stream: + return self._streams["gaze"] @property - def gaze(self) -> GazeStream: - return self.streams["gaze"] + def imu(self) -> Stream: + return self._streams["imu"] @property - def imu(self) -> IMUStream: - return self.streams["imu"] + def scene(self) -> Stream: + return self._streams["scene"] @property - def scene(self) -> VideoStream: - return self.streams["scene"] + def eye(self) -> Stream: + return self._streams["eye"] @property - def eye(self) -> VideoStream: - return self.streams["eye"] + def events(self): + return self._events + @property + def unique_events(self): + if self._unique_events is None: + log.info("NeonRecording: Parsing unique events") -def load(rec_dir_in: pathlib.Path | str) -> NeonRecording: - log.info(f"NeonRecording: Loading recording from: {rec_dir_in}") - if isinstance(rec_dir_in, str): - rec_dir = pathlib.Path(rec_dir_in) - else: - rec_dir = rec_dir_in - - if not rec_dir.is_dir(): - raise NotADirectoryError( - f"Please provide the directory with the Native Recording Data: {rec_dir}" + # when converting a list of tuples to dict, if elements repeat, then the last one + # is what ends up in the dict. + # but mpk would prefer that the first occurence of each repeated event is what + # makes it into the unique_events dict, so flippy-floppy + self._events.reverse() + self._unique_events = dict(self.events) + self._events.reverse() + + return self._unique_events + + @property + def calib_version(self): + if not self._calib_bin_loaded: + self._load_calib_bin() + self._calib_bin_loaded = True + + return self._calib_version + + @property + def serial(self): + if not self._calib_bin_loaded: + self._load_calib_bin() + self._calib_bin_loaded = True + + return self._serial + + @property + def scene_camera_calibration(self): + if not self._calib_bin_loaded: + self._load_calib_bin() + self._calib_bin_loaded = True + + return self._scene_camera_calibration + + @property + def right_eye_camera_calibration(self): + if not self._calib_bin_loaded: + self._load_calib_bin() + self._calib_bin_loaded = True + + return self._right_eye_camera_calibration + + @property + def left_eye_camera_calibration(self): + if not self._calib_bin_loaded: + self._load_calib_bin() + self._calib_bin_loaded = True + + return self._left_eye_camera_calibration + + def _load_calib_bin(self): + log.info("NeonRecording: Loading calibration data") + self._calib = parse_calib_bin(self._rec_dir) + + self._calib_version = str(self._calib["version"]) + self._serial = int(self._calib["serial"][0]) + self._scene_camera_calibration = Calibration( + self._calib["scene_camera_matrix"], + self._calib["scene_distortion_coefficients"], + self._calib["scene_extrinsics_affine_matrix"], + ) + self._right_eye_camera_calibration = Calibration( + self._calib["right_camera_matrix"], + self._calib["right_distortion_coefficients"], + self._calib["right_extrinsics_affine_matrix"], + ) + self._left_eye_camera_calibration = Calibration( + self._calib["left_camera_matrix"], + self._calib["left_distortion_coefficients"], + self._calib["left_extrinsics_affine_matrix"], + ) + + def __init__(self, rec_dir_in: pathlib.Path | str): + self._streams = { + "gaze": GazeStream("gaze", self), + "imu": IMUStream("imu", self), + "scene": VideoStream("scene", self), + "eye": VideoStream("eye", self), + } + + self._calib_bin_loaded = False + + log.info(f"NeonRecording: Loading recording from: {rec_dir_in}") + if isinstance(rec_dir_in, str): + self._rec_dir = pathlib.Path(rec_dir_in) + else: + self._rec_dir = rec_dir_in + + if not self._rec_dir.is_dir(): + raise NotADirectoryError( + f"Please provide the directory with the Native Recording Data: {self._rec_dir}" + ) + + if not self._rec_dir.exists(): + raise FileNotFoundError(f"Directory not found: {self._rec_dir}") + + log.info("NeonRecording: Loading recording info") + try: + with open(self._rec_dir / "info.json") as f: + self._info = json.load(f) + except Exception as e: + log.exception(f"Unexpected error loading 'info.json': {e}") + raise + + self._start_ts_ns = self._info["start_time"] + self._start_ts = ns_to_s(self._info["start_time"]) + + log.info("NeonRecording: Loading wearer") + self._wearer = {"uuid": "", "name": ""} + try: + with open(self._rec_dir / "wearer.json") as f: + wearer_data = json.load(f) + except Exception as e: + log.exception(f"Unexpected error loading 'info.json': {e}") + raise + + self._wearer["uuid"] = wearer_data["uuid"] + self._wearer["name"] = wearer_data["name"] + + # load up raw times, in case useful at some point + log.info("NeonRecording: Loading raw time (ns) files") + self._gaze_ps1_raw_time_ns = np.fromfile( + str(self._rec_dir / "gaze ps1.time"), dtype=" NeonRecording: + return NeonRecording(rec_dir_in) diff --git a/src/pupil_labs/neon_recording/stream/__init__.py b/src/pupil_labs/neon_recording/stream/__init__.py index a810b72..01ef138 100644 --- a/src/pupil_labs/neon_recording/stream/__init__.py +++ b/src/pupil_labs/neon_recording/stream/__init__.py @@ -1,3 +1,3 @@ from .stream import Stream # noqa: F401 -all = ["Stream"] +__all__ = ["Stream"] diff --git a/src/pupil_labs/neon_recording/stream/gaze_stream.py b/src/pupil_labs/neon_recording/stream/gaze_stream.py index 4200f39..43a53a8 100644 --- a/src/pupil_labs/neon_recording/stream/gaze_stream.py +++ b/src/pupil_labs/neon_recording/stream/gaze_stream.py @@ -4,7 +4,6 @@ import numpy as np from .. import structlog -from ..data_utils import load_with_error_check from ..time_utils import load_and_convert_tstamps from .stream import Stream @@ -34,20 +33,19 @@ def _convert_gaze_data_to_recarray(gaze_data, ts, ts_rel): class GazeStream(Stream): - def linear_interp(self, sorted_ts): - xs = self.data.x - ys = self.data.y - ts_rel = self.data.ts_rel + def _sample_linear_interp(self, sorted_ts): + xs = self._data.x + ys = self._data.y interp_data = np.zeros( len(sorted_ts), dtype=[("x", " None: + def _load(self, file_name: Optional[str] = None) -> None: # we use gaze_200hz from cloud for the rec gaze stream # ts, raw = self._load_ts_and_data(rec_dir, 'gaze ps1') - gaze_200hz_ts, gaze_200hz_raw = self._load_ts_and_data(rec_dir, "gaze_200hz") - gaze_200hz_ts_rel = gaze_200hz_ts - start_ts + gaze_200hz_ts, gaze_200hz_raw = self._load_ts_and_data("gaze_200hz") + gaze_200hz_ts_rel = gaze_200hz_ts - self._recording._start_ts data = _convert_gaze_data_to_recarray( gaze_200hz_raw, gaze_200hz_ts, gaze_200hz_ts_rel ) - self._backing_data = data - self.data = self._backing_data[:] - self.ts = self._backing_data[:].ts - self.ts_rel = self._backing_data[:].ts_rel + self._data = data + self._ts = self._data[:].ts - def _load_ts_and_data(self, rec_dir: pathlib.Path, stream_name: str): + def _load_ts_and_data(self, stream_filename: str): log.debug("NeonRecording: Loading gaze data and timestamps.") - time_path = rec_dir / (stream_name + ".time") - raw_path = rec_dir / (stream_name + ".raw") - - ts = load_with_error_check( - load_and_convert_tstamps, - time_path, - "Possible error when converting timestamps.", - ) - raw = load_with_error_check( - self._load_raw_data, raw_path, "Please double check the recording download." - ) + try: + ts = load_and_convert_tstamps( + self._recording._rec_dir / (stream_filename + ".time") + ) + except Exception as e: + log.exception(f"Error loading timestamps: {e}") + raise + + try: + raw = self._load_gaze_raw_data( + self._recording._rec_dir / (stream_filename + ".raw") + ) + except Exception as e: + log.exception(f"Error loading raw data: {e}") + raise return ts, raw - # adapted from @dom: - # https://github.com/pupil-labs/neon-player/blob/master/pupil_src/shared_modules/pupil_recording/update/neon.py - def _load_raw_data(self, path: pathlib.Path): + def _load_gaze_raw_data(self, path: pathlib.Path): log.debug("NeonRecording: Loading gaze raw data.") raw_data = np.fromfile(str(path), " None: - imu_rec = IMURecording(rec_dir / "extimu ps1.raw", start_ts) + def _load(self, file_name: Optional[str] = None) -> None: + imu_rec = IMURecording( + self._recording._rec_dir / "extimu ps1.raw", self._recording._start_ts + ) - self._backing_data = imu_rec.raw - self.data = self._backing_data[:] - self.ts = self._backing_data[:].ts - self.ts_rel = self._backing_data[:].ts_rel + self._data = imu_rec.raw + self._ts = self._data[:].ts diff --git a/src/pupil_labs/neon_recording/stream/stream.py b/src/pupil_labs/neon_recording/stream/stream.py index 397a54d..c997582 100644 --- a/src/pupil_labs/neon_recording/stream/stream.py +++ b/src/pupil_labs/neon_recording/stream/stream.py @@ -1,5 +1,6 @@ import abc -import pathlib +import math +from enum import Enum from typing import Optional import numpy as np @@ -9,69 +10,81 @@ log = structlog.get_logger(__name__) +class InterpolationMethod(Enum): + NEAREST = "nearest" + LINEAR = "linear" + + # this implements mpk's idea for a stream of data # from a neon sensor that can be sampled via ts values # see here: # https://www.notion.so/pupillabs/Neon-Recording-Python-Lib-5b247c33e1c74f638af2964fa78018ff?pvs=4 class Stream(abc.ABC): - def __init__(self, name): + def __init__(self, name, recording): self.name = name - self._backing_data = [] - self.data = [] - self.ts = [] - self.ts_rel = [] + self._recording = recording + self._backing_data = None + self._data = [] + self._ts = [] + self._ts_rel = None + + @property + def recording(self): + return self._recording + + @property + def data(self): + return self._data + + @property + def ts(self): + return self._ts + + @property + def ts_rel(self): + if self._ts_rel is None: + self._ts_rel = self._ts - self.recording.start_ts + + return self._ts_rel @abc.abstractmethod - def load( - self, rec_dir: pathlib.Path, start_ts: float, file_name: Optional[str] = None - ) -> None: + def _load(self, file_name: Optional[str] = None) -> None: pass @abc.abstractmethod - def linear_interp(self, sorted_ts): + def _sample_linear_interp(self, sorted_ts): pass def __getitem__(self, idxs): - return self.data[idxs] + return self._data[idxs] def _ts_oob(self, ts: float): - return ts < self.ts[0] or ts > self.ts[-1] + return ts < self._ts[0] or ts > self._ts[-1] - def sample_one(self, ts_wanted: float, dt: float = 0.01, method="nearest"): + def sample_one( + self, ts_wanted: float, dt: float = 0.01, method=InterpolationMethod.NEAREST + ): log.debug("NeonRecording: Sampling one timestamp.") - # in case they pass multiple timestamps - # see note at https://numpy.org/doc/stable/reference/generated/numpy.isscalar.html - if not np.ndim(ts_wanted) == 0: - raise ValueError( - "This function can only sample a single timestamp. Use 'sample' for multiple timestamps." - ) - if self._ts_oob(ts_wanted): return None - if method == "insert_order": - datum = self.data[np.searchsorted(self.ts, ts_wanted)] - if np.abs(datum.ts - ts_wanted) < dt: - return datum - else: - return None - elif method == "nearest": - diffs = np.abs(self.ts - ts_wanted) + if method == InterpolationMethod.NEAREST: + diffs = np.abs(self._ts - ts_wanted) if np.any(diffs < dt): idx = int(np.argmin(diffs)) - return self.data[idx] + return self._data[idx] else: return None - elif method == "linear": - datum = self.linear_interp([ts_wanted]) + elif method == InterpolationMethod.LINEAR: + datum = self._sample_linear_interp([ts_wanted]) if np.abs(datum.ts - ts_wanted) < dt: return datum else: return None - def sample(self, tstamps, method="nearest"): + def sample(self, tstamps, method=InterpolationMethod.NEAREST): log.debug("NeonRecording: Sampling timestamps.") # in case they pass one float @@ -83,51 +96,53 @@ def sample(self, tstamps, method="nearest"): if self._ts_oob(tstamps[0]): return None - sorted_ts = np.sort(tstamps) + sorted_tses = np.sort(tstamps) - if method == "linear": - return self.linear_interp(sorted_ts) - elif method == "nearest": - return self.sample_nearest(sorted_ts) - elif method == "insert_order": - return self.sample_sorted(sorted_ts) + if method == InterpolationMethod.NEAREST: + return self._sample_nearest(sorted_tses) + elif method == InterpolationMethod.LINEAR: + return self._sample_linear_interp(sorted_tses) else: return ValueError( - "Only 'linear', 'nearest', and 'insert_order' methods are supported." + "Only LINEAR, NEAREST, and INSERT_ORDER methods are supported." ) - # this works for all the different streams, so define it here, rather than in multiple different places - def sample_nearest(self, sorted_tses): + def _sample_nearest_rob(self, sorted_tses): log.debug("NeonRecording: Sampling nearest timestamps.") closest_idxs = [ - np.argmin(np.abs(self.ts - curr_ts)) if not self._ts_oob(curr_ts) else None + np.argmin(np.abs(self._ts - curr_ts)) if not self._ts_oob(curr_ts) else None for curr_ts in sorted_tses ] for idx in closest_idxs: if idx is not None and not np.isnan(idx): - yield self.data[int(idx)] + yield self._data[int(idx)] else: yield None - def sample_sorted(self, sorted_tses): + # from stack overflow: + # https://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array + def _sample_nearest(self, sorted_tses): # uggcf://jjj.lbhghor.pbz/jngpu?i=FRKKRF5i59b - log.debug("NeonRecording: Sampling sorted timestamps.") + log.debug("NeonRecording: Sampling nearest timestamps.") - closest_idxs = np.searchsorted(self.ts, sorted_tses) + closest_idxs = np.searchsorted(self._ts, sorted_tses, side="right") for i, ts in enumerate(sorted_tses): if self._ts_oob(ts): - closest_idxs[i] = np.nan - - for idx in closest_idxs: - if idx is not None and not np.isnan(idx): - yield self.data[int(idx)] - else: yield None + else: + idx = closest_idxs[i] + if idx > 0 and ( + idx == len(self._ts) + or math.fabs(ts - self._ts[idx - 1]) < math.fabs(ts - self._ts[idx]) + ): + yield self._data[int(idx - 1)] + else: + yield self._data[int(idx)] -def subsampled_to_numpy(sample_generator): +def sampled_to_numpy(sample_generator): fst = next(sample_generator) if isinstance(fst, np.record): diff --git a/src/pupil_labs/neon_recording/stream/video_stream.py b/src/pupil_labs/neon_recording/stream/video_stream.py index 81e5ee8..4215c45 100644 --- a/src/pupil_labs/neon_recording/stream/video_stream.py +++ b/src/pupil_labs/neon_recording/stream/video_stream.py @@ -1,64 +1,113 @@ -import pathlib -from enum import Enum -from typing import Optional, Tuple +import math +from typing import Optional +import numpy as np import pupil_labs.video as plv from .. import structlog -from ..data_utils import load_with_error_check from ..time_utils import load_and_convert_tstamps from .stream import Stream log = structlog.get_logger(__name__) -def _load_video(rec_dir: pathlib.Path, start_ts: float, video_name: str) -> Tuple: - log.debug(f"NeonRecording: Loading video and associated timestamps: {video_name}.") - - if not (rec_dir / (video_name + ".mp4")).exists(): - raise FileNotFoundError( - f"File not found: {rec_dir / (video_name + '.mp4')}. Please double check the recording download." - ) - - container = plv.open(rec_dir / (video_name + ".mp4")) - - # use hardware ts - ts = load_with_error_check( - load_and_convert_tstamps, - rec_dir / (video_name + ".time"), - "Possible error when converting timestamps.", - ) - # ts = load_with_error_check(load_and_convert_tstamps, rec_dir / (video_name + '.time_aux'), "Possible error when converting timestamps.") - - ts_rel = ts - start_ts - - return container, ts, ts_rel - +class VideoStream(Stream): + @property + def width(self): + return self._width -class VideoType(Enum): - EYE = 1 - SCENE = 2 + @property + def height(self): + return self._height + @property + def ts_rel(self): + # if self._ts_rel is None: + # self._ts_rel = self._ts - self._recording._start_ts + # setattr(self._data, "ts_rel", self._ts_rel) -class VideoStream(Stream): - def __init__(self, name, type): - super().__init__(name) - self.type = type + return self._ts_rel - def linear_interp(self, sorted_ts): + def _sample_linear_interp(self, sorted_ts): raise NotImplementedError( "NeonRecording: Video streams only support nearest neighbor interpolation." ) - def load( - self, rec_dir: pathlib.Path, start_ts: float, file_name: Optional[str] = None - ) -> None: + def _load(self, file_name: Optional[str] = None) -> None: if file_name is not None: - container, ts, ts_rel = _load_video(rec_dir, start_ts, file_name) + container, ts = self._load_video(file_name) self._backing_data = container.streams.video[0] - self.data = self._backing_data.frames - self.ts = ts - self.ts_rel = ts_rel + self._data = self._backing_data.frames + self._ts = ts + setattr(self._data, "ts", self._ts) + self._ts_rel = self._ts - self._recording._start_ts + + self._width = self._data[0].width + self._height = self._data[0].height else: raise ValueError("Filename must be provided when loading a VideoStream.") + + def _load_video(self, video_name: str): + log.debug( + f"NeonRecording: Loading video and associated timestamps: {video_name}." + ) + + if not (self._recording._rec_dir / (video_name + ".mp4")).exists(): + raise FileNotFoundError( + f"File not found: {self._recording._rec_dir / (video_name + '.mp4')}. Please double check the recording download." + ) + + container = plv.open(self._recording._rec_dir / (video_name + ".mp4")) + + # use hardware ts + # ts = load_and_convert_tstamps(self._recording._rec_dir / (video_name + '.time_aux')) + try: + ts = load_and_convert_tstamps( + self._recording._rec_dir / (video_name + ".time") + ) + except Exception as e: + log.exception(f"Error loading timestamps: {e}") + raise + + return container, ts + + def _sample_nearest_rob(self, sorted_tses): + log.debug("NeonRecording: Sampling nearest timestamps.") + + closest_idxs = [ + np.argmin(np.abs(self._ts - curr_ts)) if not self._ts_oob(curr_ts) else None + for curr_ts in sorted_tses + ] + + for idx in closest_idxs: + if idx is not None and not np.isnan(idx): + d = self._data[int(idx)] + setattr(d, "ts", self._ts[int(idx)]) + setattr(d, "ts_rel", self._ts_rel[int(idx)]) + yield d + else: + yield None + + # from stack overflow: + # https://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array + def _sample_nearest(self, sorted_tses): + # uggcf://jjj.lbhghor.pbz/jngpu?i=FRKKRF5i59b + log.debug("NeonRecording: Sampling nearest timestamps.") + + closest_idxs = np.searchsorted(self._ts, sorted_tses, side="right") + for i, ts in enumerate(sorted_tses): + if self._ts_oob(ts): + yield None + else: + idx = closest_idxs[i] + if idx > 0 and ( + idx == len(self._ts) + or math.fabs(ts - self._ts[idx - 1]) < math.fabs(ts - self._ts[idx]) + ): + idx = idx - 1 + + d = self._data[int(idx)] + setattr(d, "ts", self._ts[int(idx)]) + setattr(d, "ts_rel", self._ts_rel[int(idx)]) + yield d diff --git a/src/pupil_labs/neon_recording/time_utils.py b/src/pupil_labs/neon_recording/time_utils.py index 2c31c86..df74e61 100644 --- a/src/pupil_labs/neon_recording/time_utils.py +++ b/src/pupil_labs/neon_recording/time_utils.py @@ -12,10 +12,6 @@ def ns_to_s(ns): return ns * SECONDS_PER_NANOSECOND -# adapted from @dom and @pfaion: -# https://github.com/pupil-labs/neon-player/blob/master/pupil_src/shared_modules/pupil_recording/update/update_utils.py -# -# modified according to mpk's note def load_and_convert_tstamps(path: pathlib.Path): log.debug(f"NeonRecording: Loading timestamps from: {path.name}") diff --git a/todo.txt b/todo.txt index 706f230..ba3e039 100644 --- a/todo.txt +++ b/todo.txt @@ -3,15 +3,15 @@ - get from cloud or run offline on file load? - fixation stream - - eyestate/pupillometry stream + - eyestate/pupillometry stream - coming soon to an android near you - clarify how to deal with TsNs in imu stream - - file to notify that it has been checked - -- for video, concatenate raw then convert to np array +- for multi-part video, concatenate raw then convert to np array specify default option for agg, and provide stream specific aggs see here: https://www.notion.so/pupillabs/Neon-Recording-Python-Lib-5b247c33e1c74f638af2964fa78018ff?pvs=4#6bf3848ff3ad435aa6ea8fe37e72ee9c + +- check if data is from cloud or direct from phone