select_examples.py

#!/usr/bin/env python3
"""Select 5 random diverse test samples and copy videos to examples folder."""
import pandas as pd
import numpy as np
from pathlib import Path
import shutil

# Set seed for reproducibility
np.random.seed(42)

# Read test answers
df = pd.read_csv('videos/labels/test-answers.csv', sep=';', header=None, names=['video_id', 'label'])

# Get unique labels
unique_labels = df['label'].unique()
print(f"Total unique labels: {len(unique_labels)}")

# Select 5 different labels randomly
selected_labels = np.random.choice(unique_labels, size=min(5, len(unique_labels)), replace=False)

# For each selected label, pick one sample
selected_samples = []
for label in selected_labels:
    samples_with_label = df[df['label'] == label]
    selected_sample = samples_with_label.sample(n=1, random_state=42).iloc[0]
    selected_samples.append(selected_sample)

# Create examples directory
examples_dir = Path('notebooks/examples')
examples_dir.mkdir(parents=True, exist_ok=True)

# Copy videos and save selection info
selection_info = []
for i, sample in enumerate(selected_samples):
    video_id = sample['video_id']
    label = sample['label']
    
    # Find the video file
    video_path = Path(f'videos/20bn-something-something-v2/{video_id}.webm')
    
    if video_path.exists():
        # Copy to examples
        dest_path = examples_dir / f'{video_id}.webm'
        shutil.copy2(video_path, dest_path)
        print(f"Copied {video_id}.webm - Label: {label}")
        selection_info.append({'video_id': str(video_id), 'label': label})
    else:
        print(f"Warning: Video {video_id}.webm not found")

# Save selection info
import json
with open(examples_dir / 'selected_samples.json', 'w') as f:
    json.dump(selection_info, f, indent=2)

print(f"\nSelected {len(selection_info)} samples")