Spaces:
Running
Running
| from pathlib import Path | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| from app.project_dataset import ( | |
| _merge_info, | |
| allocate_project_episode_id, | |
| write_project_episode_index, | |
| ) | |
| def _write_frame_data(root: Path) -> None: | |
| path = root / "data" / "chunk-000" / "episode_000000.parquet" | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| pq.write_table(pa.table({ | |
| "episode_index": pa.array([0]), | |
| "frame_index": pa.array([0]), | |
| "observation.slam_pose": pa.array([[0.0] * 7]), | |
| "observation.gripper_left_force": pa.array([[1.0, 2.0, 3.0]]), | |
| }), path) | |
| def test_new_episode_id_does_not_reuse_malformed_existing_id(tmp_path: Path): | |
| write_project_episode_index(tmp_path, [{ | |
| "episode_index": 0, | |
| "episode_id": "Uncategorized_000001", | |
| "length": 10, | |
| }]) | |
| assert allocate_project_episode_id( | |
| tmp_path, "Test9_4", "Uncategorized", | |
| ) == "Uncategorized_000002" | |
| def test_empty_project_starts_at_episode_zero(tmp_path: Path): | |
| assert allocate_project_episode_id( | |
| tmp_path, "new_upload", "UMI夹爪项目", | |
| ) == "UMI夹爪项目_000000" | |
| def test_info_is_filtered_to_real_files_and_columns(tmp_path: Path): | |
| _write_frame_data(tmp_path) | |
| video = (tmp_path / "videos" / "observation.images.gripper_rgb" | |
| / "chunk-000" / "episode_000000.mp4") | |
| video.parent.mkdir(parents=True, exist_ok=True) | |
| video.write_bytes(b"placeholder") | |
| stale_info = { | |
| "task_name": "Test9/4", | |
| "features": { | |
| "observation.slam_pose": {"dtype": "float32", "shape": [7]}, | |
| "observation.gripper_left_force": { | |
| "dtype": "float32", "shape": [3], | |
| }, | |
| "observation.images.gripper_rgb": {"dtype": "video"}, | |
| "observation.images.D435_rgb": {"dtype": "video"}, | |
| "observation.imu": {"dtype": "float32", "shape": [6]}, | |
| }, | |
| "devices": [ | |
| {"key": "d435", "name": "D435", "slots": ["D435_rgb"]}, | |
| {"key": "gripper", "name": "UMI", "slots": [ | |
| "gripper_rgb", "gripper_pose", "gripper_force_left", | |
| ]}, | |
| ], | |
| "device_names": { | |
| "D435_rgb": "D435", | |
| "gripper_rgb": "UMI", | |
| "gripper_pose": "UMI", | |
| "gripper_force_left": "UMI", | |
| }, | |
| "cameras": { | |
| "D435_rgb": {"width": 1280}, | |
| "gripper_rgb": {"width": 1280}, | |
| }, | |
| "sensors": ["right_glove"], | |
| } | |
| merged = _merge_info(tmp_path, stale_info, [{ | |
| "episode_index": 0, | |
| "length": 1, | |
| }], [{"task_index": 0, "task": "Test9/4"}]) | |
| assert set(merged["features"]) == { | |
| "observation.slam_pose", | |
| "observation.gripper_left_force", | |
| "observation.images.gripper_rgb", | |
| } | |
| assert merged["cameras"] == {"gripper_rgb": {"width": 1280}} | |
| assert [d["key"] for d in merged["devices"]] == ["gripper"] | |
| assert merged["features"]["observation.slam_pose"]["quaternion_order"] == "xyzw" | |
| assert merged["features"]["observation.gripper_left_force"]["names"] == [ | |
| "fx", "fy", "fz", | |
| ] | |