#!/usr/bin/env python3 """ 将 HDF5 数据集转换为视频，用于可视化检查 功能： 1. 将单个 episode 转换为视频 2. 对比多个 episode 的视频 3. 放慢播放速度便于观察 """ import os import h5py import glob import cv2 import numpy as np def episode_to_video(episode_file, output_path, camera='top', fps=30, slow_factor=1): """ 将单个 episode 转换为视频 Args: episode_file: HDF5 文件路径 output_path: 输出视频路径 camera: 要使用的相机名称 fps: 帧率 slow_factor: 慢放倍数（1=正常，2=半速） """ try: with h5py.File(episode_file, 'r') as f: # 读取图像序列 img_path = f'/observations/images/{camera}' if img_path not in f: print(f" ❌ 相机 {camera} 不存在") return False images = f[img_path][:] # shape: (T, H, W, C) qpos = f['/observations/qpos'][:] actions = f['/action'][:] total_frames = len(images) height, width = images.shape[1], images.shape[2] # 创建视频写入器 fourcc = cv2.VideoWriter_fourcc(*'mp4v') actual_fps = fps // slow_factor out = cv2.VideoWriter(output_path, fourcc, actual_fps, (width, height)) # 逐帧写入 for i in range(total_frames): frame = images[i].astype(np.uint8) # 在图像上添加信息 info_text = [ f"Episode: {os.path.basename(episode_file).replace('.hdf5', '')}", f"Frame: {i}/{total_frames}", f"qpos[0:3]: [{qpos[i, 0]:.2f}, {qpos[i, 1]:.2f}, {qpos[i, 2]:.2f}]", ] for j, text in enumerate(info_text): cv2.putText(frame, text, (10, 30 + j*30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) out.write(frame) out.release() print(f" ✅ 保存: {output_path}") print(f" 帧数: {total_frames}, 尺寸: {width}x{height}, FPS: {actual_fps}") return True except Exception as e: print(f" ❌ 错误: {e}") return False def generate_all_videos(camera='top', num_episodes=5, slow_factor=1): """生成前 N 个 episode 的视频""" dataset_dir = "roboimi/demos/dataset/sim_transfer" episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) if len(episode_files) == 0: print(f"❌ 没有找到数据文件: {dataset_dir}") return # 创建输出目录 output_dir = '/tmp/dataset_videos' os.makedirs(output_dir, exist_ok=True) print(f"找到 {len(episode_files)} 个 episode 文件") print(f"将生成前 {min(num_episodes, len(episode_files))} 个 episode 的视频\n") # 生成视频 for i in range(min(num_episodes, len(episode_files))): ep_file = episode_files[i] ep_name = os.path.basename(ep_file).replace('.hdf5', '') output_path = f"{output_dir}/{ep_name}_{camera}.mp4" print(f"[{i+1}/{min(num_episodes, len(episode_files))}] {ep_name}") episode_to_video(ep_file, output_path, camera=camera, slow_factor=slow_factor) print() print(f"✅ 所有视频已保存到: {output_dir}") print(f"\n播放方法:") print(f" # 播放单个视频") print(f" vlc {output_dir}/*.mp4") print(f" ") print(f" # 或用文件管理器") print(f" nautilus {output_dir}") def generate_multi_camera_video(episode_idx=0, slow_factor=1): """生成包含多个相机的视频（分屏显示）""" dataset_dir = "roboimi/demos/dataset/sim_transfer" episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) if episode_idx >= len(episode_files): print(f"❌ Episode {episode_idx} 不存在") return ep_file = episode_files[episode_idx] try: with h5py.File(ep_file, 'r') as f: # 获取所有相机 cameras = [] for key in f.keys(): if 'images' in key: for cam_name in f[key].keys(): if cam_name not in cameras: cameras.append(cam_name) print(f"Episode {episode_idx} 的相机: {cameras}") # 读取所有相机的图像 all_images = {} for cam in cameras: img_path = f'/observations/images/{cam}' if img_path in f: all_images[cam] = f[img_path][:] if not all_images: print("❌ 没有找到图像数据") return # 获取第一个相机的尺寸 first_cam = list(all_images.keys())[0] total_frames = len(all_images[first_cam]) height, width = all_images[first_cam].shape[1], all_images[first_cam].shape[2] # 创建多相机布局 num_cams = len(all_images) cols = min(2, num_cams) rows = (num_cams + cols - 1) // cols canvas_width = width * cols canvas_height = height * rows # 创建视频写入器 output_path = f'/tmp/dataset_videos/episode_{episode_idx}_all_cameras.mp4' fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, 30 // slow_factor, (canvas_width, canvas_height)) # 逐帧合成 for i in range(total_frames): canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) for cam_idx, cam_name in enumerate(all_images.keys()): img = all_images[cam_name][i] # 计算在画布上的位置 row = cam_idx // cols col = cam_idx % cols y_start = row * height y_end = y_start + height x_start = col * width x_end = x_start + width # 调整大小（如果需要） if img.shape[:2] != (height, width): img = cv2.resize(img, (width, height)) # 放到画布上 canvas[y_start:y_end, x_start:x_end] = img # 添加相机名称 cv2.putText(canvas, cam_name, (x_start + 10, y_start + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2) # 添加帧信息 cv2.putText(canvas, f"Frame: {i}/{total_frames}", (10, canvas_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) out.write(canvas) out.release() print(f"✅ 保存多相机视频: {output_path}") except Exception as e: print(f"❌ 错误: {e}") def compare_episodes(camera='top', slow_factor=2): """并排对比多个 episode 的视频""" dataset_dir = "roboimi/demos/dataset/sim_transfer" episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) # 选择要对比的 episode episodes_to_compare = [0, 1, 2, 3, 4] # 对比前 5 个 print(f"对比 Episodes: {episodes_to_compare}") # 读取所有 episode 的数据 all_data = [] for ep_idx in episodes_to_compare: if ep_idx >= len(episode_files): continue try: with h5py.File(episode_files[ep_idx], 'r') as f: img_path = f'/observations/images/{camera}' if img_path in f: all_data.append({ 'idx': ep_idx, 'images': f[img_path][:], 'qpos': f['/observations/qpos'][:] }) except: pass if len(all_data) == 0: print("❌ 没有数据") return # 获取参数 first_data = all_data[0] height, width = first_data['images'].shape[1], first_data['images'].shape[2] total_frames = min([d['images'].shape[0] for d in all_data]) # 创建并排布局 num_compare = len(all_data) canvas_width = width * num_compare canvas_height = height # 创建视频 output_path = f'/tmp/dataset_videos/compare_{camera}.mp4' fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, 30 // slow_factor, (canvas_width, canvas_height)) print(f"生成对比视频，共 {total_frames} 帧...") # 逐帧对比 for i in range(total_frames): canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) for j, data in enumerate(all_data): img = data['images'][i] qpos = data['qpos'][i] # 调整大小（如果需要） if img.shape[:2] != (height, width): img = cv2.resize(img, (width, height)) # 放到画布上 x_start = j * width x_end = x_start + width canvas[:, x_start:x_end] = img # 添加信息 ep_name = f"Ep {data['idx']}" cv2.putText(canvas, ep_name, (x_start + 10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2) cv2.putText(canvas, f"qpos[0:3]: [{qpos[0]:.2f}, {qpos[1]:.2f}, {qpos[2]:.2f}]", (x_start + 10, height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) # 添加帧号 cv2.putText(canvas, f"Frame: {i}/{total_frames}", (10, canvas_height - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) out.write(canvas) if i % 100 == 0: print(f" 进度: {i}/{total_frames}") out.release() print(f"✅ 保存对比视频: {output_path}") if __name__ == "__main__": import sys print("="*60) print("数据集视频生成工具") print("="*60) if len(sys.argv) > 1: command = sys.argv[1] if command == 'compare': # 对比多个 episode camera = sys.argv[2] if len(sys.argv) > 2 else 'top' compare_episodes(camera=camera, slow_factor=2) elif command == 'multi': # 多相机视频 ep_idx = int(sys.argv[2]) if len(sys.argv) > 2 else 0 generate_multi_camera_video(episode_idx=ep_idx, slow_factor=1) else: print("未知命令") else: # 默认：生成前 5 个 episode 的视频 print("\n生成前 5 个 episode 的视频（top 相机，慢放 2x）...") print("="*60 + "\n") generate_all_videos(camera='top', num_episodes=5, slow_factor=2) print("\n" + "="*60) print("其他用法:") print(" python generate_dataset_videos.py compare top # 对比多个 episode") print(" python generate_dataset_videos.py multi 0 # 多相机视频") print("="*60)