Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,744 Bytes
246c106 |
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 |
import numpy as np
import json
import cv2
from sim.main import InteractiveDigitalWorld
from sim.simulator import GenieSimulator, ReplaySimulator
from sim.policy import ReplayPolicy
# NOTE: ad hoc
def normalize_frames(frames):
new_frames = []
for frame in frames:
H, W = frame.shape[:2]
if H < W:
Hnew, Wnew = 256, int(W * 256 / H)
else:
Hnew, Wnew = int(H * 256 / W), 256
frame = cv2.resize(frame, (Wnew, Hnew))
H, W = frame.shape[:2]
Hstart = (H - 256) // 2
Wstart = (W - 256) // 2
frame = frame[Hstart:Hstart+256, Wstart:Wstart+256]
new_frames.append(frame)
return np.stack(new_frames, axis=0)
if __name__ == '__main__':
prompt_horizon = 11
action_stride = 1
dataset_dir = "data/langtable_raw_train"
metadata = json.load(open(f"{dataset_dir}/metadata.json"))
h, w = metadata['h'], metadata['w']
action_dim = metadata['action_dim']
num_images = metadata['num_images']
actions = np.fromfile(f"{dataset_dir}/actions/actions.bin", dtype=np.float32).reshape(num_images, action_dim)
frames = np.fromfile(f"{dataset_dir}/video.bin", dtype=np.uint8).reshape(num_images, h, w, 3)
# frames = normalize_frames(frames)
segment_ids = np.fromfile(f"{dataset_dir}/segment_ids.bin", dtype=np.int32)
print(f"{actions.shape=}, {frames.shape=}, {segment_ids.shape=}")
# get chunks' start and end
chunks = [] # [start_index, end_index)
start_index = 0
end_index = 0
while end_index < len(segment_ids):
while end_index < len(segment_ids) and segment_ids[end_index] == segment_ids[start_index]:
end_index += 1
if end_index - start_index > prompt_horizon * 2:
chunks.append((start_index, end_index))
start_index = end_index
print(f"there're {len(chunks)} chunks")
for eps_idx, chunk in enumerate(chunks):
start_idx, end_idx = chunk
this_frames = frames[start_idx:end_idx]
this_actions = actions[start_idx:end_idx]
print(f"processing chunk {eps_idx} with {len(this_frames)} frames")
replay_simulator = ReplaySimulator(frames=this_frames, prompt_horizon=prompt_horizon)
replay_policy = ReplayPolicy(actions=this_actions, prompt_horizon=prompt_horizon, action_stride=action_stride)
assert len(replay_policy) == len(replay_simulator)
genie_simulator = GenieSimulator(
image_encoder_type='temporalvae',
image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
quantize=False,
backbone_type="stmar",
backbone_ckpt="data/mar_ckpt/langtable",
prompt_horizon=prompt_horizon,
action_stride=action_stride,
domain='language_table',
physics_simulator=replay_simulator,
compute_psnr=False,
compute_delta_psnr=False,
allow_external_prompt=True
)
# use whatever current state is as the initial state
image_prompt = replay_simulator.prompt()
action_prompt = replay_policy.prompt()
genie_simulator.set_initial_state((image_prompt, action_prompt))
playground = InteractiveDigitalWorld(
simulator=genie_simulator,
policy=replay_policy,
offscreen=True,
window_size=(512 * 2, 512) # [genie image | GT image] side-by-side
)
for _ in range(len(replay_policy)):
playground.step()
save_video_path = f'data/langtable_train_videos/{eps_idx}.mp4'
print(f"Saving video to {save_video_path}")
playground.save_video(save_path=save_video_path, as_gif=False)
playground.close() |