feat(is_depth): simplifying is_depth nested name + legacy support

src/lerobot/datasets/dataset_metadata.py

@@ -341,14 +341,20 @@ class LeRobotDatasetMetadata:
     def depth_keys(self) -> list[str]:
         """Keys to access depth-map modalities stored as videos or images.
 
-        A depth key is a feature whose ``info`` dict carries ``"<dtype>.is_depth_map": True``.
+        A depth key is a feature whose ``info`` dict carries ``"is_depth_map": True``
+        (or the legacy ``"video.is_depth_map"`` inside ``info`` or ``video_info``).
         """
-        return [
-            key
-            for key, ft in self.features.items()
-            # TODO(CarolinePascal): Make sure the legacy video_info works here as well.
-            if (ft.get("info") or {}).get(ft["dtype"] + ".is_depth_map", False)
-        ]
+
+        def _is_depth(ft: dict) -> bool:
+            info = ft.get("info") or {}
+            video_info = ft.get("video_info") or {}
+            return (
+                info.get("is_depth_map", False)
+                or info.get("video.is_depth_map", False)
+                or video_info.get("video.is_depth_map", False)
+            )
+
+        return [key for key, ft in self.features.items() if _is_depth(ft)]
 
     @property
     def camera_keys(self) -> list[str]:

src/lerobot/datasets/dataset_writer.py

@@ -155,7 +155,7 @@ class DatasetWriter:
         return ep_buffer
 
     def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
-        path_template = DEFAULT_DEPTH_PATH if self.image_key in self._meta.depth_keys else DEFAULT_IMAGE_PATH
+        path_template = DEFAULT_DEPTH_PATH if image_key in self._meta.depth_keys else DEFAULT_IMAGE_PATH
         fpath = path_template.format(
             image_key=image_key, episode_index=episode_index, frame_index=frame_index
         )

src/lerobot/datasets/video_utils.py

@@ -1112,7 +1112,6 @@ def get_video_info(
         video_info["video.width"] = video_stream.width
         video_info["video.codec"] = video_stream.codec.canonical_name
         video_info["video.pix_fmt"] = video_stream.pix_fmt
-        video_info["video.is_depth_map"] = False
 
         # Calculate fps from r_frame_rate
         video_info["video.fps"] = int(video_stream.base_rate)
@@ -1133,7 +1132,7 @@ def get_video_info(
             if field_name == "vcodec":
                 continue
             video_info.setdefault(f"video.{field_name}", field_value)
-        video_info["video.is_depth_map"] = isinstance(video_encoder, DepthEncoderConfig)
+        video_info["is_depth_map"] = isinstance(video_encoder, DepthEncoderConfig)
 
     return video_info
 

src/lerobot/utils/feature_utils.py

@@ -88,21 +88,18 @@ def hw_to_dataset_features(
 
     for key, shape in cam_fts.items():
         dtype = "video" if use_video else "image"
-        if len(shape) == 2 or shape[2] == 1:
-            if len(shape) == 2:
-                shape = (shape[0], shape[1], 1)
-            features[f"{prefix}.depth_maps.{key}"] = {
-                "dtype": dtype,
-                "shape": shape,
-                "names": ["height", "width", "channels"],
-                "info": {dtype + ".is_depth_map": True},
-            }
-        else:
+        if len(shape) == 3 and shape[2] in (1, 3):
             features[f"{prefix}.images.{key}"] = {
                 "dtype": dtype,
                 "shape": shape,
                 "names": ["height", "width", "channels"],
+                "info": {"is_depth_map": shape[2] == 1},
             }
+        else:
+            raise ValueError(
+                f"Camera feature '{key}' has shape {shape}. "
+                f"Expected a 3-tuple (H, W, C), e.g. (480, 640, 3) for RGB or (480, 640, 1) for depth."
+            )
 
     _validate_feature_names(features)
     return features
@@ -132,10 +129,7 @@ def build_dataset_frame(
         elif ft["dtype"] == "float32" and len(ft["shape"]) == 1:
             frame[key] = np.array([values[name] for name in ft["names"]], dtype=np.float32)
         elif ft["dtype"] in ["image", "video"]:
-            if ft["info"].get(ft["dtype"] + ".is_depth_map"):
-                frame[key] = values[key.removeprefix(f"{prefix}.depth_maps.")]
-            else:
-                frame[key] = values[key.removeprefix(f"{prefix}.images.")]
+            frame[key] = values[key.removeprefix(f"{prefix}.images.")]
 
     return frame
 
@@ -164,11 +158,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
             type = FeatureType.VISUAL
             if len(shape) != 3:
                 raise ValueError(f"Number of dimensions of {key} != 3 (shape={shape})")
-
-            names = ft["names"]
-            # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
-            if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
-                shape = (shape[2], shape[0], shape[1])
+            else:
+                names = ft["names"]
+                # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
+                if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
+                    shape = (shape[2], shape[0], shape[1])
         elif key == OBS_ENV_STATE:
             type = FeatureType.ENV
         elif key.startswith(OBS_STR):