Spaces:
Running on Zero
Running on Zero
File size: 8,397 Bytes
dcc4594 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | import torch
import time
import os
import matplotlib.pyplot as plt
import cv2
import numpy as np
from utils.utils import load_videotokenizer_from_checkpoint, load_latent_actions_from_checkpoint, load_dynamics_from_checkpoint
from einops import repeat
def load_models(video_tokenizer_path, latent_actions_path, dynamics_path, device, use_actions=True):
# Load tokenizer and dynamics, and Latent Actions if using actions
video_tokenizer, _vt_ckpt = load_videotokenizer_from_checkpoint(video_tokenizer_path, device)
video_tokenizer.eval()
latent_action_model = None
if use_actions:
latent_action_model, _latent_action_ckpt = load_latent_actions_from_checkpoint(latent_actions_path, device)
latent_action_model.eval()
dynamics_model, _dyn_ckpt = load_dynamics_from_checkpoint(dynamics_path, device)
dynamics_model.eval()
return video_tokenizer, latent_action_model, dynamics_model
def visualize_inference(predicted_frames, ground_truth_frames, inferred_actions, fps, use_actions=True):
# Move to CPU and convert to numpy
predicted_frames = predicted_frames.detach().cpu()
ground_truth_frames = ground_truth_frames.detach().cpu()
# Denormalize frames from [-1, 1] to [0, 1]
predicted_frames = (predicted_frames + 1) / 2
predicted_frames = torch.clamp(predicted_frames, 0, 1)
ground_truth_frames = (ground_truth_frames + 1) / 2
ground_truth_frames = torch.clamp(ground_truth_frames, 0, 1)
# Get dimensions
B, T, C, H, W = predicted_frames.shape
_, num_gt_frames, _, _, _ = ground_truth_frames.shape
# Create figure with ground truth and predictions side by side
fig, axes = plt.subplots(2, T, figsize=(4 * T, 8))
# Handle single subplot case
if T == 1:
axes = axes.reshape(2, 1)
# Plot ground truth frames (top row)
for i in range(num_gt_frames):
frame = ground_truth_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C]
axes[0, i].imshow(frame)
axes[0, i].set_title(f'Ground Truth {i+1}', fontsize=12, color='green')
axes[0, i].axis('off')
# Plot predicted frames (bottom row)
for i in range(T):
frame = predicted_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C]
axes[1, i].imshow(frame)
title = f'Predicted {i+1}'
if use_actions and i < len(inferred_actions):
title += f'\nAction {inferred_actions[i].item()}' if i < len(inferred_actions) else ''
axes[1, i].set_title(title, fontsize=12, color='red')
axes[1, i].axis('off')
plt.suptitle('Ground Truth vs Predicted Frames', fontsize=16, fontweight='bold')
# Save the visualization
timestamp = time.strftime("%Y%m%d_%H%M%S")
save_dir = "inference_results"
os.makedirs(save_dir, exist_ok=True)
if use_actions:
save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_{timestamp}.png")
mp4_path = os.path.join(save_dir, f"inference_video_{timestamp}.mp4")
else:
save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_no_actions_{timestamp}.png")
mp4_path = os.path.join(save_dir, f"inference_video_no_actions_{timestamp}.mp4")
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
print(f"Visualization saved to: {save_path}")
all_frames = torch.cat([ground_truth_frames, predicted_frames], dim=1)
save_frames_as_mp4(all_frames, mp4_path, fps)
# Calculate and display reconstruction error
mse_error = torch.mean((predicted_frames - ground_truth_frames) ** 2).item()
print(f"\nInference stats:")
print(f"Total frames generated: {T}")
print(f"Mean Squared Error (GT vs Pred): {mse_error:.6f}")
if use_actions:
print(f"Actions used: {[action.item() for action in inferred_actions]}")
else:
print(f"No actions used.")
# TODO: get working mp4
def save_frames_as_mp4(frames, output_path, fps=2):
B, T, C, H, W = frames.shape
# OpenCV expects (W, H)
fourcc = cv2.VideoWriter_fourcc(*'avc1')
out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
for i in range(T):
frame = frames[0, i].detach().cpu().permute(1, 2, 0).numpy() # [H, W, C]
# Ensure float32
frame = frame.astype(np.float32)
# Clamp and scale
frame = np.clip(frame, 0, 1)
frame = (frame * 255).astype(np.uint8)
# If grayscale, convert to 3 channels
if frame.shape[2] == 1:
frame = np.repeat(frame, 3, axis=2)
# Convert RGB to BGR for OpenCV
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
out.write(frame_bgr)
out.release()
print(f"MP4 video saved to: {output_path}")
def sample_random_action(n_actions):
random_action = torch.randint(0, n_actions, (1,))
return random_action
def get_action_latent(args, inferred_actions, n_actions, context_frames, latent_action_model, step):
if args.use_interactive_mode: # let user input actions
print("using interactive mode")
user_input = input(f"Enter action id [0..{n_actions-1}] for step {step+1}: ").strip()
assert user_input.isdigit() and 0 <= int(user_input) < n_actions, f"Invalid input. Please enter an integer in [0,{n_actions-1}]"
val = int(user_input)
sampled_action_index = torch.tensor([val], device=args.device)
inferred_actions.append(sampled_action_index)
recent = inferred_actions[-args.context_window:] if len(inferred_actions) > args.context_window else inferred_actions
recent_tensor = repeat(torch.tensor(recent, device=args.device), 'i -> 1 i') # [1, i or T_ctx]
action_latent = latent_action_model.quantizer.get_latents_from_indices(recent_tensor)
if args.prediction_horizon > 1:
action_latent = repeat(action_latent, 'b 1 a -> b ph a', ph=args.prediction_horizon)
if len(recent) < args.context_window:
gt_pad_actions = latent_action_model.encode(context_frames[:, :args.context_window - len(recent) + 1])
quantized_gt_pad_actions = latent_action_model.quantizer(gt_pad_actions)
action_latent = torch.cat([quantized_gt_pad_actions, action_latent], dim=1)
elif args.use_gt_actions: # use action tokenizer actions
print("using gt actions")
gt_action_latents = latent_action_model.encode(context_frames) # [1, T - 1, A]
sampled_action_index = sample_random_action(n_actions) # [1]
inferred_actions.append(sampled_action_index) # [i]
sampled_action_index_tensor = repeat(torch.tensor(sampled_action_index, device=args.device), 'i -> 1 i') # [1, i]
sampled_action_latent = latent_action_model.quantizer.get_latents_from_indices(sampled_action_index_tensor) # [1, i, A]
action_latent = torch.cat([gt_action_latents, sampled_action_latent], dim=1) # [1, T, A]
elif args.use_actions: # use random actions
print(f"using random actions")
sampled_action_index = sample_random_action(n_actions) # [1]
inferred_actions.append(sampled_action_index) # [i]
recent_inferred_actions = inferred_actions[-args.context_window:] if len(inferred_actions) > args.context_window else inferred_actions # [i or T_ctx]
recent_inferred_actions_tensor = repeat(torch.tensor(recent_inferred_actions, device=args.device), 'i -> 1 i') # [1, i or T_ctx]
action_latent = latent_action_model.quantizer.get_latents_from_indices(recent_inferred_actions_tensor) # [1, i or T_ctx, A]
if args.prediction_horizon > 1:
action_latent = repeat(action_latent, 'b 1 a -> b ph a', ph=args.prediction_horizon)
if len(recent_inferred_actions) < args.context_window:
# if we dont have enough inferred actions (in the beginning) add enough gt to fill the sequence
gt_pad_actions = latent_action_model.encode(context_frames[:, :args.context_window - len(recent_inferred_actions) + 1]) # [1, context_window - len(inferred_actions), A]
quantized_gt_pad_actions = latent_action_model.quantizer(gt_pad_actions) # [1, context_window - len(inferred_actions), A]
action_latent = torch.cat([quantized_gt_pad_actions, action_latent], dim=1) # [1, S, A]
else:
sampled_action_index = None
action_latent = None
return sampled_action_index, action_latent
|