Spaces:
Running on Zero
Running on Zero
| 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 | |