| import torch |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import os |
|
|
| def plot_length_distribution(dataset_name, metadata_path, output_path): |
| print(f"Loading metadata from {metadata_path}...") |
| metadata = torch.load(metadata_path, weights_only=False) |
| |
| lengths = [item['length'] for item in metadata] |
| |
| plt.figure(figsize=(10, 6)) |
| plt.hist(lengths, bins=50, color='skyblue', edgecolor='black', alpha=0.7) |
| |
| mean_len = np.mean(lengths) |
| median_len = np.median(lengths) |
| |
| plt.axvline(mean_len, color='red', linestyle='dashed', linewidth=1, label=f'Mean: {mean_len:.1f}') |
| plt.axvline(median_len, color='green', linestyle='dashed', linewidth=1, label=f'Median: {median_len:.1f}') |
| |
| plt.title(f'Video Length Distribution - {dataset_name}') |
| plt.xlabel('Number of Frames') |
| plt.ylabel('Frequency') |
| plt.grid(axis='y', alpha=0.3) |
| plt.legend() |
| |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| plt.savefig(output_path) |
| print(f"Plot saved to {output_path}") |
| |
| print(f"Stats for {dataset_name}:") |
| print(f" Total Trajectories: {len(lengths)}") |
| print(f" Min Length: {np.min(lengths)}") |
| print(f" Max Length: {np.max(lengths)}") |
| print(f" Mean Length: {mean_len:.1f}") |
| print(f" Median Length: {median_len:.1f}") |
|
|
| if __name__ == "__main__": |
| datasets = { |
| "language_table": "/storage/ice-shared/ae8803che/hxue/data/dataset/language_table/metadata_lite.pt", |
| "rt1": "/storage/ice-shared/ae8803che/hxue/data/dataset/rt1/metadata_lite.pt", |
| "recon": "/storage/ice-shared/ae8803che/hxue/data/dataset/recon_processed/metadata_lite.pt", |
| "dreamer4": "/storage/ice-shared/ae8803che/hxue/data/dataset/dreamer4_processed/metadata_lite.pt" |
| } |
| |
| for name, meta_path in datasets.items(): |
| out_path = f"/storage/ice-shared/ae8803che/hxue/data/world_model/results/stats/{name}_dist.png" |
| plot_length_distribution(name, meta_path, out_path) |
|
|
