wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

support depth image when aggregating stats

This commit is contained in:
Tavish 2025-04-07 21:54:33 +08:00
parent 1c873df5c0
commit f8cfe2e1a9
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
lerobot/common/datasets/compute_stats.py
View File

@ -119,8 +119,11 @@ def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
raise ValueError("Number of dimensions must be at least 1, and is 0 instead.")
if k == "count" and v.shape != (1,):
raise ValueError(f"Shape of 'count' must be (1), but is {v.shape} instead.")
if "image" in fkey and k != "count" and v.shape != (3, 1, 1):
raise ValueError(f"Shape of '{k}' must be (3,1,1), but is {v.shape} instead.")
if "image" in fkey and k != "count":
if "depth" not in fkey and v.shape != (3, 1, 1):
raise ValueError(f"Shape of '{k}' must be (3,1,1), but is {v.shape} instead.")
if "depth" in fkey and v.shape != (1, 1, 1):
raise ValueError(f"Shape of '{k}' must be (1,1,1), but is {v.shape} instead.")
def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]: