p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] reprocess_knoms_to_optimized_arrayrecord_layout.py (python)
from utils.dataloader import convert_tfrecords_to_arrayrecords

convert_tfrecords_to_arrayrecords(
    tfrecord_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards",
    output_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_group_size_1",
)
[content] utils/dataloader.py (python)
                pixels_per_channel = total_pixels // channels
                
                # Try common aspect ratios for 160x90
                if pixels_per_channel == 160 * 90:
                    height, width = 90, 160
                else:
                    # Generic fallback - assume square-ish
                    side_length = int(np.sqrt(pixels_per_channel))
                    height = width = side_length
                    print(f"Warning: Using fallback dimensions {height}x{width} for {total_pixels} pixels")
[content] utils/dataloader.py (python)
            
            # For 160x90 dimensions, we know the exact shape
            # Check if it matches our expected dimensions
            expected_pixels_160x90_rgb = 160 * 90 * 3  # RGB
            expected_pixels_160x90_gray = 160 * 90 * 1  # Grayscale
            
            if total_pixels == expected_pixels_160x90_rgb:
                height, width, channels = 90, 160, 3
            elif total_pixels == expected_pixels_160x90_gray:
                height, width, channels = 90, 160, 1
[content] utils/dataloader.py (python)
            # Assume square-ish dimensions, find closest square root
            side_length = int(np.sqrt(total_pixels))
            if side_length * side_length * 3 == total_pixels:  # RGB
                channels = 3
            elif side_length * side_length == total_pixels:  # Grayscale
                channels = 1
[content] utils/dataloader.py (python)
            
            # Use the provided dimensions
            expected_pixels = image_h * image_w * image_c
            
            if total_pixels != expected_pixels:
                print(f"Warning: Expected {expected_pixels} pixels per frame, got {total_pixels}")
                print(f"Expected shape: ({sequence_length}, {image_h}, {image_w}, {image_c})")
                print(f"Actual total pixels: {total_pixels}")
                # Try to infer dimensions as fallback
                if total_pixels % image_c == 0:
                    pixels_per_channel = total_pixels // image_c
                    if pixels_per_channel == image_h * image_w:
                        # Dimensions match, use provided values
                        height, width, channels = image_h, image_w, image_c
                    else:
                        # Calculate closest dimensions
                        side_length = int(np.sqrt(pixels_per_channel))
                        height = width = side_length
                        channels = image_c
                        print(f"Using fallback dimensions {height}x{width}x{channels}")
                else:
                    # Assume grayscale if channel count doesn't divide evenly
                    height, width, channels = image_h, image_w, 1
                    print(f"Using grayscale fallback: {height}x{width}x{channels}")
[content] utils/dataloader.py (python)
    if num_workers is None:
        num_workers = min(mp.cpu_count(), len(arrayrecord_files))
    
    print(f"Using {num_workers} worker processes for reprocessing")
    
    # Create a list of arguments for each file processing task
    process_args = []
    for i, file_path in enumerate(arrayrecord_files):
        process_args.append((
            file_path,
            output_folder,
            chunk_size,
            videos_per_file,
            image_h,
            image_w,
            image_c,
            i,
        ))
    
    # Process files in parallel
    with mp.Pool(processes=num_workers) as pool:
        results = pool.starmap(_reprocess_single_arrayrecord, process_args)
    
    # Flatten the results
    all_output_files = []
    for result in results:
        all_output_files.extend(result)
[content] utils/dataloader.py (python)
def _reprocess_single_arrayrecord(
    arrayrecord_file: Path,
    output_folder: str,
    chunk_size: int,
    videos_per_file: int,
    image_h: int,
    image_w: int,
    image_c: int,
    file_index: int,
) -> list[str]:
    """
    Reprocess a single ArrayRecord file by splitting videos into chunks.
    
    Args:
        arrayrecord_file: Path to the ArrayRecord file
        output_folder: Output folder for the chunked files
        chunk_size: Number of frames per video chunk
        videos_per_file: Number of video chunks per output file
        image_h: Image height in pixels
        image_w: Image width in pixels
        image_c: Number of image channels
        file_index: Index for naming output files
    
    Returns:
        List of paths to created ArrayRecord files
    """
    print(f"Processing {arrayrecord_file.name}...")
    reader = ArrayRecordReader(str(arrayrecord_file))
    
    all_record_bytes = reader.read_all()
    file_chunks = []
    
    for record_bytes in all_record_bytes:
        record = pickle.loads(record_bytes)
        video_data = record['raw_video']
        sequence_length = record['sequence_length']
        
        video_tensor = np.frombuffer(video_data, dtype=np.uint8)
        
        assert video_tensor.shape[0] == sequence_length * image_h * image_w * image_c, \
            f"Video tensor shape {video_tensor.shape} does not match expected shape {sequence_length * image_h * image_w * image_c}"
        video_tensor = video_tensor.reshape(sequence_length, image_h, image_w, image_c)
        
        current_episode_len = video_tensor.shape[0]
        if current_episode_len < chunk_size:
            print(f"Warning: Video has {current_episode_len} frames, skipping (need {chunk_size})")
            continue
        
        for start_idx in range(0, current_episode_len - chunk_size + 1, chunk_size):
            chunk = video_tensor[start_idx:start_idx + chunk_size]
            
            # FIXME: currently no way of correlating the chunk with the original video
            chunk_record = {
                'raw_video': chunk.tobytes(),
                'sequence_length': chunk_size,
            }
            
            file_chunks.append(chunk_record)
    
    reader.close()
    
    # Write chunks to output files
    output_files = []
    for i in range(0, len(file_chunks), videos_per_file):
        batch_chunks = file_chunks[i:i + videos_per_file]
        output_filename = f"chunked_videos_{file_index:04d}_{i//videos_per_file:04d}.array_record"
        output_file = os.path.join(output_folder, output_filename)
        
        writer = ArrayRecordWriter(output_file, "group_size:1")
        for chunk in batch_chunks:
            writer.write(pickle.dumps(chunk))
        writer.close()
        
        output_files.append(output_file)
        print(f"Created {output_filename} with {len(batch_chunks)} video chunks")
    
    print(f"Processed {arrayrecord_file.name}: {len(file_chunks)} chunks -> {len(output_files)} files")
    return output_files
[tab] experiments/tokenizer_dummy_time_measurement_grain.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy_arrayrecords_600_shards'

# FIXME: the `time_measurement_interval` line is only for debugging and should be removed
srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] delete_corrupted_files.py (python)
#!/usr/bin/env python3

import os
from pathlib import Path

def delete_corrupted_files():
    """Delete the corrupted array record files."""
    corrupted_files = [
        "chunked_videos_0198_0001.array_record",
        "chunked_videos_0261_0001.array_record", 
        "chunked_videos_0165_0001.array_record",
        "chunked_videos_0063_0001.array_record",
        "chunked_videos_0258_0000.array_record",
        "chunked_videos_0358_0000.array_record",
        "chunked_videos_0271_0000.array_record",
        "chunked_videos_0399_0000.array_record",
        "chunked_videos_0373_0000.array_record"
    ]
    
    output_folder = "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_optimized_layout"
    
    print("=== Deleting Corrupted Files ===")
    deleted_count = 0
    total_size_mb = 0
    
    for filename in corrupted_files:
        file_path = os.path.join(output_folder, filename)
        if os.path.exists(file_path):
            size_mb = os.path.getsize(file_path) / (1024 * 1024)
            print(f"Deleting {filename} ({size_mb:.1f} MB)...")
            os.remove(file_path)
            deleted_count += 1
            total_size_mb += size_mb
        else:
            print(f"Warning: {filename} not found, skipping...")
    
    print(f"
=== Deletion Summary ===")
    print(f"Files deleted: {deleted_count}")
    print(f"Total space freed: {total_size_mb:.1f} MB")
    print(f"Remaining files: 856 (out of original 865)")
    
    # Verify deletion
    print(f"
=== Verification ===")
    remaining_corrupted = 0
    for filename in corrupted_files:
        file_path = os.path.join(output_folder, filename)
        if os.path.exists(file_path):
            print(f"Warning: {filename} still exists!")
            remaining_corrupted += 1
    
    if remaining_corrupted == 0:
        print("✓ All corrupted files successfully deleted!")
    else:
        print(f"✗ {remaining_corrupted} corrupted files still exist")

if __name__ == "__main__":
    delete_corrupted_files()
[tab] calculate_avg_frames.py (python)
#!/usr/bin/env python3
"""
Script to calculate the average number of frames per video in an ArrayRecord file.
"""

import grain
import pickle
import sys

def calculate_avg_frames(array_record_path):
    """
    Calculate the average number of frames per video in an ArrayRecord file.
    
    Args:
        array_record_path: Path to the ArrayRecord file
        
    Returns:
        Average number of frames per video
    """
    total_frames = 0
    num_videos = 0
    
    print(f"Reading ArrayRecord file: {array_record_path}")
    
    # Use Grain's ArrayRecordDataSource to read the file
    data_source = grain.sources.ArrayRecordDataSource([array_record_path])
    num_records = len(data_source)
    
    print(f"Found {num_records} records in the file")
    
    for i in range(num_records):
        record_bytes = data_source[i]
        data = pickle.loads(record_bytes)
        total_frames += data["sequence_length"]
        num_videos += 1
        
        # Print progress every 1000 videos
        if num_videos % 1000 == 0:
            print(f"Processed {num_videos}/{num_records} videos...")
    
    if num_videos > 0:
        avg_frames = total_frames / num_videos
        print(f"
Results:")
        print(f"Total videos: {num_videos}")
        print(f"Total frames: {total_frames}")
        print(f"Average frames per video: {avg_frames:.2f}")
        return avg_frames
    else:
        print("No videos found in the ArrayRecord file.")
        return 0

if __name__ == "__main__":
    # Default path to the ArrayRecord file
    array_record_path = "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy_arrayrecords_600_shards/videos_0010.array_record"
    
    # Allow command line argument to override the path
    if len(sys.argv) > 1:
        array_record_path = sys.argv[1]
    
    try:
        calculate_avg_frames(array_record_path)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)