mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
feat(robots/so_follower): emit + populate depth keys when use_depth
When an SO follower has a camera configured with use_depth=True (e.g. a RealSense), the robot now exposes a paired depth feature so the dataset records both modalities: - _cameras_ft adds a 2D "<cam>_depth" entry alongside the 3-channel color shape; hw_to_dataset_features turns this into observation.depth.<cam> with the depth-map marker. - get_observation reads cam.read_latest_depth() (float32 metric meters from the RealSense async depth API) into <cam>_depth so build_dataset_frame can route it. Detection is duck-typed via getattr(..., "use_depth", False) so other cameras without that attribute keep their RGB-only behaviour unchanged. Made-with: Cursor
This commit is contained in:
@@ -68,9 +68,16 @@ class SOFollower(Robot):
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def _cameras_ft(self) -> dict[str, tuple]:
|
def _cameras_ft(self) -> dict[str, tuple]:
|
||||||
return {
|
features: dict[str, tuple] = {}
|
||||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
for cam in self.cameras:
|
||||||
}
|
cam_cfg = self.config.cameras[cam]
|
||||||
|
features[cam] = (cam_cfg.height, cam_cfg.width, 3)
|
||||||
|
# Cameras with a depth stream (e.g. RealSense with use_depth=True) also
|
||||||
|
# emit a 2D depth feature; hw_to_dataset_features routes 2D shapes to
|
||||||
|
# ``observation.depth.<bare>`` with the depth-map marker.
|
||||||
|
if getattr(cam_cfg, "use_depth", False):
|
||||||
|
features[f"{cam}_depth"] = (cam_cfg.height, cam_cfg.width)
|
||||||
|
return features
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def observation_features(self) -> dict[str, type | tuple]:
|
def observation_features(self) -> dict[str, type | tuple]:
|
||||||
@@ -190,6 +197,14 @@ class SOFollower(Robot):
|
|||||||
dt_ms = (time.perf_counter() - start) * 1e3
|
dt_ms = (time.perf_counter() - start) * 1e3
|
||||||
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
|
||||||
|
|
||||||
|
# Cameras with a depth stream populate a sibling ``<cam>_depth`` key
|
||||||
|
# (consumed by hw_to_dataset_features / build_dataset_frame).
|
||||||
|
if getattr(self.config.cameras[cam_key], "use_depth", False):
|
||||||
|
start = time.perf_counter()
|
||||||
|
obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth()
|
||||||
|
dt_ms = (time.perf_counter() - start) * 1e3
|
||||||
|
logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms")
|
||||||
|
|
||||||
return obs_dict
|
return obs_dict
|
||||||
|
|
||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
|
|||||||
Reference in New Issue
Block a user