HDF5viewerPro.py

import h5py
import cv2
import numpy as np

def start_interactive_viewer(file_path):
    with h5py.File(file_path, "r") as f:
        # 1. 获取所有 Demo 列表并排序
        demo_names = sorted(list(f['data'].keys()), key=lambda x: int(x.split('_')[-1]) if '_' in x else x)
        total_demos = len(demo_names)
        
        current_demo_idx = 0
        is_paused = False
        
        print(f"\n--- 交互播放器已启动 (共 {total_demos} 个 Demo) ---")
        print("操作指南:")
        print("  [Space]: 暂停 / 播放")
        print("  [D] / [A]: 下一个 / 上一个 Demo")
        print("  [R]: 重新开始当前 Demo")
        print("  [Q]: 退出播放器")
        print("  [J]: 在终端输入数字跳转到指定 Demo")
        print("-------------------------------------------\n")

        while True:
            # 获取当前 Demo 的图像数据
            demo_name = demo_names[current_demo_idx]
            # 路径根据 LIBERO 结构，如果不对请修改此处
            img_path = f'data/{demo_name}/obs/agentview_rgb'
            images = f[img_path][:]
            
            # 数据预处理 (归一化检查)
            if images.dtype == np.float32 or images.max() <= 1.0:
                images = (images * 255).astype(np.uint8)

            frame_idx = 0
            num_frames = len(images)

            # --- 当前 Demo 的播放循环 ---
            while frame_idx < num_frames:
                img = images[frame_idx]
                # RGB 转 BGR 用于 OpenCV 显示
                img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

                # 在图片上叠加文字信息 (UI)
                status = "PAUSED" if is_paused else "PLAYING"
                info_text = f"Demo:({current_demo_idx+1}/{total_demos})"
                cv2.putText(img_bgr, info_text, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                
                cv2.imshow("LIBERO Dataset Viewer", img_bgr)

                # 等待按键输入
                # 如果暂停，等待时间设为 0 (无限等待)；如果播放，等待 30ms
                wait_time = 0 if is_paused else 30
                key = cv2.waitKey(wait_time) & 0xFF

                if key == ord('q'): # 退出
                    cv2.destroyAllWindows()
                    return
                elif key == ord(' '): # 暂停/播放
                    is_paused = not is_paused
                elif key == ord('d'): # 下一个 Demo
                    current_demo_idx = (current_demo_idx + 1) % total_demos
                    break # 跳出当前 Demo 循环
                elif key == ord('a'): # 上一个 Demo
                    current_demo_idx = (current_demo_idx - 1) % total_demos
                    break
                elif key == ord('r'): # 重置当前 Demo
                    frame_idx = 0
                    continue
                elif key == ord('j'): # 跳转
                    try:
                        new_idx = int(input(f"请输入要跳转的 Demo 索引 (0-{total_demos-1}): "))
                        if 0 <= new_idx < total_demos:
                            current_demo_idx = new_idx
                            break
                    except ValueError:
                        print("无效输入，请输入数字")

                # 如果没有暂停，索引增加
                if not is_paused:
                    frame_idx += 1
                
                # 如果播放到末尾，自动重播（或暂停在最后一帧）
                if frame_idx >= num_frames:
                    frame_idx = 0 # 这里设为 0 是循环播放，设为 num_frames-1 则是停在末尾

# 使用示例
file_path = "/home/robot/00_DYF-WS/Files/libero_spatial/pick_up_the_black_bowl_from_table_center_and_place_it_on_the_plate_demo.hdf5"
start_interactive_viewer(file_path)