Spaces:
Running
Running
Aditya Shankar
feat: added dataset recording; hf uploader, s3 uploader; runpod trainer (#6)
4384839
unverified
| export interface VideoInfo { | |
| height: number; | |
| width: number; | |
| channels: number; | |
| codec: string; | |
| pix_fmt: string; | |
| is_depth_map: boolean; | |
| has_audio: boolean; | |
| } | |
| /** | |
| * Metadata parameters interface | |
| */ | |
| interface MetadataParams { | |
| total_episodes: number; | |
| total_frames: number; | |
| total_tasks: number; | |
| chunks_size: number; | |
| fps: number; | |
| splits: { [key: string]: string }; | |
| features: { [key: string]: any }; | |
| videos_info: VideoInfo[]; | |
| data_files_size_in_mb: number; | |
| video_files_size_in_mb: number; | |
| } | |
| /** | |
| * Generates and returns a metadata information dictionary | |
| * Needs some named parameters passed as parameters | |
| */ | |
| function getMetadataInfo(params: MetadataParams) { | |
| return { | |
| "codebase_version": "v2.1", | |
| "robot_type": "so100", | |
| "total_episodes": params.total_episodes, | |
| "total_frames": params.total_frames, | |
| "total_tasks": params.total_tasks, | |
| "total_videos": params.videos_info.length, | |
| "total_chunks": 1, | |
| "chunks_size": params.chunks_size, | |
| "fps": params.fps, | |
| "splits": { | |
| "train": `0:${params.total_episodes}` | |
| }, | |
| "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", | |
| "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", | |
| "features": { | |
| "action": { | |
| "dtype": "float32", | |
| "shape": [ | |
| 6 | |
| ], | |
| "names": [ | |
| "main_shoulder_pan", | |
| "main_shoulder_lift", | |
| "main_elbow_flex", | |
| "main_wrist_flex", | |
| "main_wrist_roll", | |
| "main_gripper" | |
| ], | |
| "fps": params.fps | |
| }, | |
| "observation.state": { | |
| "dtype": "float32", | |
| "shape": [ | |
| 6 | |
| ], | |
| "names": [ | |
| "main_shoulder_pan", | |
| "main_shoulder_lift", | |
| "main_elbow_flex", | |
| "main_wrist_flex", | |
| "main_wrist_roll", | |
| "main_gripper" | |
| ], | |
| "fps": params.fps | |
| }, | |
| "observation.images.front": { | |
| "dtype": "video", | |
| "shape": [ | |
| 480, | |
| 640, | |
| 3 | |
| ], | |
| "names": [ | |
| "height", | |
| "width", | |
| "channels" | |
| ], | |
| "info": { | |
| "video.fps": params.fps, | |
| "video.height": 480, | |
| "video.width": 640, | |
| "video.channels": 3, | |
| "video.codec": "av1", | |
| "video.pix_fmt": "yuv420p", | |
| "video.is_depth_map": false, | |
| "has_audio": false | |
| } | |
| }, | |
| "timestamp": { | |
| "dtype": "float32", | |
| "shape": [ | |
| 1 | |
| ], | |
| "names": null, | |
| "fps": params.fps | |
| }, | |
| "frame_index": { | |
| "dtype": "int64", | |
| "shape": [ | |
| 1 | |
| ], | |
| "names": null, | |
| "fps": params.fps | |
| }, | |
| "episode_index": { | |
| "dtype": "int64", | |
| "shape": [ | |
| 1 | |
| ], | |
| "names": null, | |
| "fps": params.fps | |
| }, | |
| "index": { | |
| "dtype": "int64", | |
| "shape": [ | |
| 1 | |
| ], | |
| "names": null, | |
| "fps": params.fps | |
| }, | |
| "task_index": { | |
| "dtype": "int64", | |
| "shape": [ | |
| 1 | |
| ], | |
| "names": null, | |
| "fps": params.fps | |
| } | |
| }, | |
| "data_files_size_in_mb": 100, | |
| "video_files_size_in_mb": 500 | |
| } | |
| } | |
| export function getVideoInfo(width: number, height: number): VideoInfo { | |
| return { | |
| height, | |
| width, | |
| channels: 3, | |
| codec: "h264", | |
| pix_fmt: "yuv420p", | |
| is_depth_map: false, | |
| has_audio: false | |
| }; | |
| } | |
| export default getMetadataInfo; | |