feat(aruco): extend SVOReader with depth map support
- Add depth_mode parameter to SVOReader.__init__() - Add enable_depth property - Add depth_map field to FrameData dataclass - Add _retrieve_depth(), get_depth_at(), get_depth_window_median() methods - Update grab_all() and grab_synced() to retrieve depth when enabled
This commit is contained in:
@@ -13,16 +13,20 @@ class FrameData:
|
|||||||
timestamp_ns: int
|
timestamp_ns: int
|
||||||
frame_index: int
|
frame_index: int
|
||||||
serial_number: int
|
serial_number: int
|
||||||
|
depth_map: Optional[np.ndarray] = None
|
||||||
|
|
||||||
|
|
||||||
class SVOReader:
|
class SVOReader:
|
||||||
"""Handles synchronized playback of multiple SVO files."""
|
"""Handles synchronized playback of multiple SVO files."""
|
||||||
|
|
||||||
def __init__(self, svo_paths: List[str]):
|
def __init__(
|
||||||
|
self, svo_paths: List[str], depth_mode: sl.DEPTH_MODE = sl.DEPTH_MODE.NONE
|
||||||
|
):
|
||||||
self.svo_paths = svo_paths
|
self.svo_paths = svo_paths
|
||||||
self.cameras: List[sl.Camera] = []
|
self.cameras: List[sl.Camera] = []
|
||||||
self.camera_info: List[Dict] = []
|
self.camera_info: List[Dict] = []
|
||||||
self.runtime_params = sl.RuntimeParameters()
|
self.runtime_params = sl.RuntimeParameters()
|
||||||
|
self._depth_mode = depth_mode
|
||||||
|
|
||||||
for path in svo_paths:
|
for path in svo_paths:
|
||||||
if not os.path.exists(path):
|
if not os.path.exists(path):
|
||||||
@@ -32,7 +36,7 @@ class SVOReader:
|
|||||||
init_params = sl.InitParameters()
|
init_params = sl.InitParameters()
|
||||||
init_params.set_from_svo_file(path)
|
init_params.set_from_svo_file(path)
|
||||||
init_params.svo_real_time_mode = False
|
init_params.svo_real_time_mode = False
|
||||||
init_params.depth_mode = sl.DEPTH_MODE.NONE
|
init_params.depth_mode = depth_mode
|
||||||
|
|
||||||
cam = sl.Camera()
|
cam = sl.Camera()
|
||||||
status = cam.open(init_params)
|
status = cam.open(init_params)
|
||||||
@@ -87,6 +91,7 @@ class SVOReader:
|
|||||||
if err == sl.ERROR_CODE.SUCCESS:
|
if err == sl.ERROR_CODE.SUCCESS:
|
||||||
mat = sl.Mat()
|
mat = sl.Mat()
|
||||||
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
||||||
|
depth_map = self._retrieve_depth(cam)
|
||||||
frames.append(
|
frames.append(
|
||||||
FrameData(
|
FrameData(
|
||||||
image=mat.get_data().copy(),
|
image=mat.get_data().copy(),
|
||||||
@@ -95,6 +100,7 @@ class SVOReader:
|
|||||||
).get_nanoseconds(),
|
).get_nanoseconds(),
|
||||||
frame_index=cam.get_svo_position(),
|
frame_index=cam.get_svo_position(),
|
||||||
serial_number=self.camera_info[i]["serial"],
|
serial_number=self.camera_info[i]["serial"],
|
||||||
|
depth_map=depth_map,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
||||||
@@ -140,6 +146,7 @@ class SVOReader:
|
|||||||
if err == sl.ERROR_CODE.SUCCESS:
|
if err == sl.ERROR_CODE.SUCCESS:
|
||||||
mat = sl.Mat()
|
mat = sl.Mat()
|
||||||
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
cam.retrieve_image(mat, sl.VIEW.LEFT)
|
||||||
|
depth_map = self._retrieve_depth(cam)
|
||||||
frames[i] = FrameData(
|
frames[i] = FrameData(
|
||||||
image=mat.get_data().copy(),
|
image=mat.get_data().copy(),
|
||||||
timestamp_ns=cam.get_timestamp(
|
timestamp_ns=cam.get_timestamp(
|
||||||
@@ -147,6 +154,7 @@ class SVOReader:
|
|||||||
).get_nanoseconds(),
|
).get_nanoseconds(),
|
||||||
frame_index=cam.get_svo_position(),
|
frame_index=cam.get_svo_position(),
|
||||||
serial_number=self.camera_info[i]["serial"],
|
serial_number=self.camera_info[i]["serial"],
|
||||||
|
depth_map=depth_map,
|
||||||
)
|
)
|
||||||
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
elif err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
|
||||||
cam.set_svo_position(0)
|
cam.set_svo_position(0)
|
||||||
@@ -156,6 +164,52 @@ class SVOReader:
|
|||||||
|
|
||||||
return frames
|
return frames
|
||||||
|
|
||||||
|
@property
|
||||||
|
def enable_depth(self) -> bool:
|
||||||
|
return self._depth_mode != sl.DEPTH_MODE.NONE
|
||||||
|
|
||||||
|
def _retrieve_depth(self, cam: sl.Camera) -> Optional[np.ndarray]:
|
||||||
|
if not self.enable_depth:
|
||||||
|
return None
|
||||||
|
depth_mat = sl.Mat()
|
||||||
|
cam.retrieve_measure(depth_mat, sl.MEASURE.DEPTH)
|
||||||
|
return depth_mat.get_data().copy()
|
||||||
|
|
||||||
|
def get_depth_at(self, frame: FrameData, x: int, y: int) -> Optional[float]:
|
||||||
|
if frame.depth_map is None:
|
||||||
|
return None
|
||||||
|
h, w = frame.depth_map.shape[:2]
|
||||||
|
if x < 0 or x >= w or y < 0 or y >= h:
|
||||||
|
return None
|
||||||
|
depth = frame.depth_map[y, x]
|
||||||
|
if not np.isfinite(depth) or depth <= 0:
|
||||||
|
return None
|
||||||
|
return float(depth)
|
||||||
|
|
||||||
|
def get_depth_window_median(
|
||||||
|
self, frame: FrameData, x: int, y: int, size: int = 5
|
||||||
|
) -> Optional[float]:
|
||||||
|
if frame.depth_map is None:
|
||||||
|
return None
|
||||||
|
if size % 2 == 0:
|
||||||
|
size += 1
|
||||||
|
|
||||||
|
h, w = frame.depth_map.shape[:2]
|
||||||
|
half = size // 2
|
||||||
|
|
||||||
|
x_min = max(0, x - half)
|
||||||
|
x_max = min(w, x + half + 1)
|
||||||
|
y_min = max(0, y - half)
|
||||||
|
y_max = min(h, y + half + 1)
|
||||||
|
|
||||||
|
window = frame.depth_map[y_min:y_max, x_min:x_max]
|
||||||
|
valid_depths = window[np.isfinite(window) & (window > 0)]
|
||||||
|
|
||||||
|
if len(valid_depths) == 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return float(np.median(valid_depths))
|
||||||
|
|
||||||
def close(self):
|
def close(self):
|
||||||
"""Closes all cameras."""
|
"""Closes all cameras."""
|
||||||
for cam in self.cameras:
|
for cam in self.cameras:
|
||||||
|
|||||||
Reference in New Issue
Block a user