(predicted_frames, ground_truth_frames, inferred_actions, fps, use_actions=True)
| 21 | |
| 22 | |
| 23 | def visualize_inference(predicted_frames, ground_truth_frames, inferred_actions, fps, use_actions=True): |
| 24 | # Move to CPU and convert to numpy |
| 25 | predicted_frames = predicted_frames.detach().cpu() |
| 26 | ground_truth_frames = ground_truth_frames.detach().cpu() |
| 27 | |
| 28 | # Denormalize frames from [-1, 1] to [0, 1] |
| 29 | predicted_frames = (predicted_frames + 1) / 2 |
| 30 | predicted_frames = torch.clamp(predicted_frames, 0, 1) |
| 31 | ground_truth_frames = (ground_truth_frames + 1) / 2 |
| 32 | ground_truth_frames = torch.clamp(ground_truth_frames, 0, 1) |
| 33 | |
| 34 | # Get dimensions |
| 35 | B, T, C, H, W = predicted_frames.shape |
| 36 | |
| 37 | _, num_gt_frames, _, _, _ = ground_truth_frames.shape |
| 38 | |
| 39 | # Create figure with ground truth and predictions side by side |
| 40 | fig, axes = plt.subplots(2, T, figsize=(4 * T, 8)) |
| 41 | |
| 42 | # Handle single subplot case |
| 43 | if T == 1: |
| 44 | axes = axes.reshape(2, 1) |
| 45 | |
| 46 | # Plot ground truth frames (top row) |
| 47 | for i in range(num_gt_frames): |
| 48 | frame = ground_truth_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C] |
| 49 | axes[0, i].imshow(frame) |
| 50 | axes[0, i].set_title(f'Ground Truth {i+1}', fontsize=12, color='green') |
| 51 | axes[0, i].axis('off') |
| 52 | |
| 53 | # Plot predicted frames (bottom row) |
| 54 | for i in range(T): |
| 55 | frame = predicted_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C] |
| 56 | axes[1, i].imshow(frame) |
| 57 | title = f'Predicted {i+1}' |
| 58 | if use_actions and i < len(inferred_actions): |
| 59 | title += f'\nAction {inferred_actions[i].item()}' if i < len(inferred_actions) else '' |
| 60 | axes[1, i].set_title(title, fontsize=12, color='red') |
| 61 | axes[1, i].axis('off') |
| 62 | |
| 63 | plt.suptitle('Ground Truth vs Predicted Frames', fontsize=16, fontweight='bold') |
| 64 | |
| 65 | # Save the visualization |
| 66 | timestamp = time.strftime("%Y%m%d_%H%M%S") |
| 67 | save_dir = "inference_results" |
| 68 | os.makedirs(save_dir, exist_ok=True) |
| 69 | |
| 70 | if use_actions: |
| 71 | save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_{timestamp}.png") |
| 72 | mp4_path = os.path.join(save_dir, f"inference_video_{timestamp}.mp4") |
| 73 | else: |
| 74 | save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_no_actions_{timestamp}.png") |
| 75 | mp4_path = os.path.join(save_dir, f"inference_video_no_actions_{timestamp}.mp4") |
| 76 | |
| 77 | plt.savefig(save_path, dpi=150, bbox_inches='tight') |
| 78 | plt.close() |
| 79 | |
| 80 | print(f"Visualization saved to: {save_path}") |
no test coverage detected