adapt to pythonic style. merge searchsorted and nearest neighbor samp…

…ling. adapt examples.

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

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()

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:")

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)

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"]

src/pupil_labs/neon_recording/calib.py

@@ -1,4 +1,5 @@
 import pathlib
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -7,19 +8,20 @@
 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")
 
     calib_raw_data: bytes = b""
     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