File size: 3,935 Bytes
a3f6fc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
Simplified episode loader + video writer.
"""
import argparse
import pathlib
import numpy as np
import imageio
from PIL import Image
# helpers
def _normalize_loaded_to_frames(loaded: np.ndarray):
"""
Convert a 4D numpy array (T,H,W,C) into a list of HxWxC uint8 frames.
"""
if not isinstance(loaded, np.ndarray) or loaded.ndim != 4:
raise ValueError(f"Expected 4D ndarray (T,H,W,C), got {type(loaded)} with shape {getattr(loaded, 'shape', None)}")
frames = []
for f in loaded:
f = np.asarray(f)
if f.dtype != np.uint8:
if np.issubdtype(f.dtype, np.floating):
if f.max() <= 1.0:
f = (f * 255).round().astype(np.uint8)
else:
f = f.round().astype(np.uint8)
else:
f = np.clip(f, 0, 255).astype(np.uint8)
frames.append(f)
return frames
# core loader
def load_episode_pixels(episode_path: pathlib.Path):
"""
Load an episode file (.npy, .npz, or .pth) and return a list of HxWxC uint8 frames.
"""
episode_path = pathlib.Path(episode_path)
if not episode_path.exists():
raise FileNotFoundError(f"{episode_path} does not exist")
if episode_path.suffix == '.npz':
npz = np.load(episode_path, allow_pickle=False)
loaded = npz[list(npz.files)[0]]
elif episode_path.suffix == '.npy':
loaded = np.load(episode_path)
elif episode_path.suffix in ('.pth', '.pt'):
import torch
loaded = torch.load(episode_path, map_location='cpu', weights_only=False)
if hasattr(loaded, 'numpy'):
loaded = loaded.numpy()
else:
raise ValueError(f"Unsupported file extension: {episode_path.suffix}")
frames = _normalize_loaded_to_frames(loaded)
return frames
# video writer
def write_video(frames, out_path, fps=30):
if len(frames) == 0:
raise ValueError("No frames to write")
H, W, C = frames[0].shape
# Resize frames if needed
for i, f in enumerate(frames):
if f.shape != (H, W, C):
f = Image.fromarray(f).resize((W, H), Image.BILINEAR)
frames[i] = np.asarray(f)
out_path = pathlib.Path(out_path)
out_path.parent.mkdir(parents=True, exist_ok=True)
writer = imageio.get_writer(str(out_path), fps=fps, codec='libx264',
quality=8, ffmpeg_params=['-pix_fmt', 'yuv420p'])
try:
for f in frames:
writer.append_data(f)
finally:
writer.close()
# CLI
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--save_folder", required=True,
help="Path to save folder that contains 'obses' subfolder")
parser.add_argument("--episode", type=int, required=True,
help="Episode index (integer)")
parser.add_argument("--out", required=True,
help="Output video path (e.g. ep00000.mp4)")
parser.add_argument("--fps", type=int, default=30, help="Video FPS")
args = parser.parse_args()
args.episode = args.episode - 1 # fix indexing
save_folder = pathlib.Path(args.save_folder)
obses_folder = save_folder / "obses"
candidates = [
obses_folder / f"episode_{args.episode:05d}.pth",
obses_folder / f"episode_{args.episode:05d}.pt",
obses_folder / f"episode_{args.episode:05d}.npz",
obses_folder / f"episode_{args.episode:05d}.npy",
]
epfile = next((c for c in candidates if c.exists()), None)
if epfile is None:
raise FileNotFoundError(f"No episode file found for index {args.episode} in {obses_folder}")
print(f"Loading episode from {epfile} ...")
frames = load_episode_pixels(epfile)
print(f"Loaded {len(frames)} frames. Writing video to {args.out} ...")
write_video(frames, args.out, fps=args.fps)
print("Done.")
if __name__ == "__main__":
main()
|