Add ArUco detection and pose estimation modules
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -0,0 +1,163 @@
|
||||
import pyzed.sl as sl
|
||||
import numpy as np
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Dict
|
||||
import os
|
||||
|
||||
|
||||
@dataclass
|
||||
class FrameData:
|
||||
"""Data structure for a single frame from an SVO."""
|
||||
|
||||
image: np.ndarray
|
||||
timestamp_ns: int
|
||||
frame_index: int
|
||||
serial_number: int
|
||||
|
||||
|
||||
class SVOReader:
|
||||
"""Handles synchronized playback of multiple SVO files."""
|
||||
|
||||
def __init__(self, svo_paths: List[str]):
|
||||
self.svo_paths = svo_paths
|
||||
self.cameras: List[sl.Camera] = []
|
||||
self.camera_info: List[Dict] = []
|
||||
self.runtime_params = sl.RuntimeParameters()
|
||||
|
||||
for path in svo_paths:
|
||||
if not os.path.exists(path):
|
||||
print(f"Warning: SVO file not found: {path}")
|
||||
continue
|
||||
|
||||
init_params = sl.InitParameters()
|
||||
init_params.set_from_svo_file(path)
|
||||
init_params.svo_real_time_mode = False
|
||||
init_params.depth_mode = sl.DEPTH_MODE.NONE
|
||||
|
||||
cam = sl.Camera()
|
||||
status = cam.open(init_params)
|
||||
if status != sl.ERROR_CODE.SUCCESS:
|
||||
print(f"Error: Could not open {path}: {status}")
|
||||
continue
|
||||
|
||||
info = cam.get_camera_information()
|
||||
self.cameras.append(cam)
|
||||
self.camera_info.append(
|
||||
{
|
||||
"serial": info.serial_number,
|
||||
"fps": info.camera_configuration.fps,
|
||||
"total_frames": cam.get_svo_number_of_frames(),
|
||||
}
|
||||
)
|
||||
|
||||
def sync_to_latest_start(self):
|
||||
"""Aligns all SVOs to the timestamp of the latest starting SVO."""
|
||||
if not self.cameras:
|
||||
return
|
||||
|
||||
start_timestamps = []
|
||||
for cam in self.cameras:
|
||||
err = cam.grab(self.runtime_params)
|
||||
if err == sl.ERROR_CODE.SUCCESS:
|
||||
ts = cam.get_timestamp(sl.TIME_REFERENCE.IMAGE).get_nanoseconds()
|
||||
start_timestamps.append(ts)
|
||||
else:
|
||||
start_timestamps.append(0)
|
||||
|
||||
if not start_timestamps:
|
||||
return
|
||||
|
||||
max_start_ts = max(start_timestamps)
|
||||
|
||||
for i, cam in enumerate(self.cameras):
|
||||
ts = start_timestamps[i]
|
||||
if ts < max_start_ts:
|
||||
diff_ns = max_start_ts - ts
|
||||
fps = self.camera_info[i]["fps"]
|
||||
frames_to_skip = int((diff_ns / 1_000_000_000) * fps)
|
||||
cam.set_svo_position(frames_to_skip)
|
||||
else:
|
||||
cam.set_svo_position(0)
|
||||
|
||||
def grab_all(self) -> List[Optional[FrameData]]:
|
||||
"""Grabs a frame from all cameras without strict synchronization."""
|
||||
frames = []
|
||||
for i, cam in enumerate(self.cameras):
|
||||
err = cam.grab(self.runtime_params)
|
||||
if err == sl.ERROR_CODE.SUCCESS:
|
||||
mat = sl.Mat()
|
||||
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
||||
frames.append(
|
||||
FrameData(
|
||||
image=mat.get_data().copy(),
|
||||
timestamp_ns=cam.get_timestamp(
|
||||
sl.TIME_REFERENCE.IMAGE
|
||||
).get_nanoseconds(),
|
||||
frame_index=cam.get_svo_position(),
|
||||
serial_number=self.camera_info[i]["serial"],
|
||||
)
|
||||
)
|
||||
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
||||
cam.set_svo_position(0)
|
||||
frames.append(None)
|
||||
else:
|
||||
frames.append(None)
|
||||
return frames
|
||||
|
||||
def grab_synced(self, tolerance_ms: int = 33) -> List[Optional[FrameData]]:
|
||||
"""
|
||||
Grabs frames from all cameras, attempting to keep them within tolerance_ms.
|
||||
If a camera falls behind, it skips frames.
|
||||
"""
|
||||
frames = self.grab_all()
|
||||
if not any(frames):
|
||||
return frames
|
||||
|
||||
# Find latest timestamp
|
||||
valid_timestamps = [f.timestamp_ns for f in frames if f is not None]
|
||||
if not valid_timestamps:
|
||||
return frames
|
||||
|
||||
max_ts = max(valid_timestamps)
|
||||
tolerance_ns = tolerance_ms * 1_000_000
|
||||
|
||||
for i, frame in enumerate(frames):
|
||||
if frame is None:
|
||||
continue
|
||||
|
||||
if max_ts - frame.timestamp_ns > tolerance_ns:
|
||||
cam = self.cameras[i]
|
||||
fps = self.camera_info[i]["fps"]
|
||||
diff_ns = max_ts - frame.timestamp_ns
|
||||
frames_to_skip = int((diff_ns / 1_000_000_000) * fps)
|
||||
|
||||
if frames_to_skip > 0:
|
||||
current_pos = cam.get_svo_position()
|
||||
cam.set_svo_position(current_pos + frames_to_skip)
|
||||
|
||||
err = cam.grab(self.runtime_params)
|
||||
|
||||
if err == sl.ERROR_CODE.SUCCESS:
|
||||
mat = sl.Mat()
|
||||
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
||||
frames[i] = FrameData(
|
||||
image=mat.get_data().copy(),
|
||||
timestamp_ns=cam.get_timestamp(
|
||||
sl.TIME_REFERENCE.IMAGE
|
||||
).get_nanoseconds(),
|
||||
frame_index=cam.get_svo_position(),
|
||||
serial_number=self.camera_info[i]["serial"],
|
||||
)
|
||||
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
||||
cam.set_svo_position(0)
|
||||
frames[i] = None
|
||||
else:
|
||||
frames[i] = None
|
||||
|
||||
return frames
|
||||
|
||||
def close(self):
|
||||
"""Closes all cameras."""
|
||||
for cam in self.cameras:
|
||||
cam.close()
|
||||
self.cameras = []
|
||||
Reference in New Issue
Block a user