[content] generate_arrayrecord_dataset.py (python)
"""
Generate ArrayRecord dataset compatible with train_tokenizer.py

This script creates synthetic video data and saves it in ArrayRecord format
that can be used by the tokenizer training script.
"""

import os
import pickle
import numpy as np
import tensorflow as tf
from array_record.python import array_record_module
import argparse
from typing import Tuple, List
import random


def generate_synthetic_video(
    seq_len: int = 16,
    height: int = 90,
    width: int = 160,
    channels: int = 3
) -> np.ndarray:
    """
    Generate synthetic video data for training.
    
    Args:
        seq_len: Number of frames in the video sequence
        height: Height of each frame
        width: Width of each frame
        channels: Number of color channels
        
    Returns:
        Video array of shape (seq_len, height, width, channels)
    """
    # Generate a simple moving pattern
    video = np.zeros((seq_len, height, width, channels), dtype=np.float32)
    
    # Create a moving circle pattern
    for t in range(seq_len):
        # Circle center moves over time
        center_x = int(width * 0.5 + width * 0.3 * np.sin(t * 0.5))
        center_y = int(height * 0.5 + height * 0.3 * np.cos(t * 0.3))
        radius = 20
        
        # Create frame with moving circle
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                if dist < radius:
                    # Circle color changes over time
                    color = np.array([
                        0.8 + 0.2 * np.sin(t * 0.2),
                        0.6 + 0.4 * np.cos(t * 0.3),
                        0.7 + 0.3 * np.sin(t * 0.4)
                    ], dtype=np.float32)
                    video[t, y, x] = color
                else:
                    # Background with subtle gradient
                    video[t, y, x] = np.array([
                        0.1 + 0.1 * (x / width),
                        0.1 + 0.1 * (y / height),
                        0.1 + 0.1 * (t / seq_len)
                    ], dtype=np.float32)
    
    return video


def create_arrayrecord_dataset(
    output_dir: str,
    num_videos: int = 1000,
    seq_len: int = 16,
    height: int = 90,
    width: int = 160,
    channels: int = 3,
    records_per_file: int = 100,
    seed: int = 42
):
    """
    Create ArrayRecord dataset with synthetic video data.
    
    Args:
        output_dir: Directory to save ArrayRecord files
        num_videos: Total number of videos to generate
        seq_len: Number of frames per video
        height: Frame height
        width: Frame width
        channels: Number of color channels
        records_per_file: Number of records per ArrayRecord file
        seed: Random seed for reproducibility
    """
    # Set random seeds
    random.seed(seed)
    np.random.seed(seed)
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Calculate number of files needed
    num_files = (num_videos + records_per_file - 1) // records_per_file
    
    print(f"Generating {num_videos} videos across {num_files} ArrayRecord files...")
    print(f"Each file will contain up to {records_per_file} videos")
    print(f"Video dimensions: {seq_len} frames × {height}×{width}×{channels}")
    
    total_records = 0
    
    for file_idx in range(num_files):
        # Calculate videos for this file
        start_idx = file_idx * records_per_file
        end_idx = min((file_idx + 1) * records_per_file, num_videos)
        videos_in_file = end_idx - start_idx
        
        # Create ArrayRecord file
        file_path = os.path.join(output_dir, f"videos_{file_idx:04d}.array_record")
        writer = array_record_module.ArrayRecordWriter(file_path, "group_size:1")
        
        print(f"Creating file {file_path} with {videos_in_file} videos...")
        
        for video_idx in range(videos_in_file):
            # Generate synthetic video
            video = generate_synthetic_video(seq_len, height, width, channels)
            
            # Create record in the format expected by train_tokenizer.py
            # The dataloader expects videos with shape (seq_len, height, width, channels)
            record = {
                "videos": video.astype(np.float32)
            }
            
            # Serialize and write record
            writer.write(pickle.dumps(record))
            total_records += 1
            
            if (video_idx + 1) % 100 == 0:
                print(f"  Generated {video_idx + 1}/{videos_in_file} videos")
        
        writer.close()
        print(f"Completed file {file_path}")
    
    print(f"
Dataset generation complete!")
    print(f"Total videos generated: {total_records}")
    print(f"Files created in: {output_dir}")
    print(f"Files: {len([f for f in os.listdir(output_dir) if f.endswith('.array_record')])}")


def verify_dataset(output_dir: str, num_samples: int = 5):
    """
    Verify the generated dataset by loading a few samples.
    
    Args:
        output_dir: Directory containing ArrayRecord files
        num_samples: Number of samples to verify
    """
    print(f"
Verifying dataset in {output_dir}...")
    
    # Find all ArrayRecord files
    array_record_files = [
        os.path.join(output_dir, f) 
        for f in os.listdir(output_dir) 
        if f.endswith('.array_record')
    ]
    
    if not array_record_files:
        print("No ArrayRecord files found!")
        return
    
    print(f"Found {len(array_record_files)} ArrayRecord files")
    
    # Load and verify a few samples
    for i in range(min(num_samples, len(array_record_files))):
        file_path = array_record_files[i]
        print(f"
Verifying file: {os.path.basename(file_path)}")
        
        # Read a few records from the file
        reader = array_record_module.ArrayRecordReader(file_path)
        
        for j, record_bytes in enumerate(reader):
            if j >= 3:  # Only check first 3 records
                break
                
            record = pickle.loads(record_bytes)
            video = record["videos"]
            
            print(f"  Record {j}: video shape = {video.shape}, dtype = {video.dtype}")
            print(f"    Value range: [{video.min():.3f}, {video.max():.3f}]")
            print(f"    Mean: {video.mean():.3f}")
        
        reader.close()


def main():
    parser = argparse.ArgumentParser(description="Generate ArrayRecord dataset for tokenizer training")
    parser.add_argument("--output_dir", type=str, default="data_arrayrecord/coinrun",
                       help="Output directory for ArrayRecord files")
    parser.add_argument("--num_videos", type=int, default=1000,
                       help="Total number of videos to generate")
    parser.add_argument("--seq_len", type=int, default=16,
                       help="Number of frames per video")
    parser.add_argument("--height", type=int, default=90,
                       help="Frame height")
    parser.add_argument("--width", type=int, default=160,
                       help="Frame width")
    parser.add_argument("--channels", type=int, default=3,
                       help="Number of color channels")
    parser.add_argument("--records_per_file", type=int, default=100,
                       help="Number of records per ArrayRecord file")
    parser.add_argument("--seed", type=int, default=42,
                       help="Random seed for reproducibility")
    parser.add_argument("--verify", action="store_true",
                       help="Verify the generated dataset")
    
    args = parser.parse_args()
    
    # Generate the dataset
    create_arrayrecord_dataset(
        output_dir=args.output_dir,
        num_videos=args.num_videos,
        seq_len=args.seq_len,
        height=args.height,
        width=args.width,
        channels=args.channels,
        records_per_file=args.records_per_file,
        seed=args.seed
    )
    
    # Verify if requested
    if args.verify:
        verify_dataset(args.output_dir)


if __name__ == "__main__":
    main()