mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 21:19:40 +00:00
Use generic names in multidimensional add_features test
This commit is contained in:
@@ -381,7 +381,7 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
|
|||||||
def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
|
def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
|
||||||
"""Test adding a multi-dimensional per-frame feature."""
|
"""Test adding a multi-dimensional per-frame feature."""
|
||||||
num_frames = sample_dataset.meta.total_frames
|
num_frames = sample_dataset.meta.total_frames
|
||||||
latent_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
|
feature_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
|
||||||
|
|
||||||
feature_info = {
|
feature_info = {
|
||||||
"dtype": "float32",
|
"dtype": "float32",
|
||||||
@@ -389,7 +389,7 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
|
|||||||
"names": None,
|
"names": None,
|
||||||
}
|
}
|
||||||
features = {
|
features = {
|
||||||
"latent_vectors": (latent_values, feature_info),
|
"multi_dim_feature": (feature_values, feature_info),
|
||||||
}
|
}
|
||||||
|
|
||||||
with (
|
with (
|
||||||
@@ -397,19 +397,19 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
|
|||||||
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
|
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
|
||||||
):
|
):
|
||||||
mock_get_safe_version.return_value = "v3.0"
|
mock_get_safe_version.return_value = "v3.0"
|
||||||
mock_snapshot_download.return_value = str(tmp_path / "with_latents")
|
mock_snapshot_download.return_value = str(tmp_path / "with_multi_dim_feature")
|
||||||
|
|
||||||
new_dataset = add_features(
|
new_dataset = add_features(
|
||||||
dataset=sample_dataset,
|
dataset=sample_dataset,
|
||||||
features=features,
|
features=features,
|
||||||
output_dir=tmp_path / "with_latents",
|
output_dir=tmp_path / "with_multi_dim_feature",
|
||||||
)
|
)
|
||||||
|
|
||||||
assert "latent_vectors" in new_dataset.meta.features
|
assert "multi_dim_feature" in new_dataset.meta.features
|
||||||
sample_item = new_dataset[0]
|
sample_item = new_dataset[0]
|
||||||
assert "latent_vectors" in sample_item
|
assert "multi_dim_feature" in sample_item
|
||||||
assert isinstance(sample_item["latent_vectors"], torch.Tensor)
|
assert isinstance(sample_item["multi_dim_feature"], torch.Tensor)
|
||||||
assert tuple(sample_item["latent_vectors"].shape) == (1, 4)
|
assert tuple(sample_item["multi_dim_feature"].shape) == (1, 4)
|
||||||
|
|
||||||
|
|
||||||
def test_add_existing_feature(sample_dataset, tmp_path):
|
def test_add_existing_feature(sample_dataset, tmp_path):
|
||||||
|
|||||||
Reference in New Issue
Block a user