Function visualize_batch

scripts/visualize_batch.py:9–48

(frames, save_path=None, title="Video Sequences Batch", max_batch_size=8, max_seq_length=8)

from datasets.data_utils import load_data_and_data_loaders

def visualize_batch(frames, save_path=None, title="Video Sequences Batch", max_batch_size=8, max_seq_length=8):
    # move to CPU and get dimensions
    frames = frames.detach().cpu()
    batch_size, seq_len, C, H, W = frames.shape
    batch_size = min(batch_size, max_batch_size)
    seq_len = min(seq_len, max_seq_length)
    frames = frames[:batch_size, :seq_len]

    # denormalize from [-1, 1] to [0, 1]
    frames = (frames + 1) / 2
    frames = torch.clamp(frames, 0, 1)
    fig, axes = plt.subplots(batch_size, seq_len, figsize=(2 * seq_len, 2 * batch_size))

    # single row/column case
    if batch_size == 1:
        axes = axes.reshape(1, -1)
    if seq_len == 1:
        axes = axes.reshape(-1, 1)

    # plot frames
    for i in range(batch_size):
        for j in range(seq_len):
            frame = frames[i, j].permute(1, 2, 0).numpy()  # [H, W, C]

            axes[i, j].imshow(frame)
            axes[i, j].set_title(f'B{i}, T{j}', fontsize=10)
            axes[i, j].axis('off')

    # row and column labels
    for i in range(batch_size):
        axes[i, 0].set_ylabel(f'Batch {i}', fontsize=12, fontweight='bold')
    for j in range(seq_len):
        axes[0, j].set_title(f'Timestep {j}', fontsize=12, fontweight='bold')
    plt.suptitle(title, fontsize=16, fontweight='bold')
    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Visualization saved to: {save_path}")

    plt.show()

def visualize_batch_with_stats(frames, save_path=None, title="Video Sequences Batch with Statistics"):
    # Move to CPU and get dimensions