[content] utils/dataloader.py (python)

def reprocess_arrayrecords_to_chunks(
    input_folder: str,
    output_folder: str,
    chunk_size: int = 160,
    videos_per_file: int = 100,
):
    """
    Reprocesses ArrayRecord files by splitting videos into chunks and creating new files.
    
    This function:
    1. Reads existing ArrayRecord files
    2. Splits each video into chunks of specified size
    3. Creates new ArrayRecord files with specified number of videos per file
    
    Args:
        input_folder: Path to folder containing input ArrayRecord files
        output_folder: Path to output folder for new ArrayRecord files
        chunk_size: Number of frames per video chunk (default: 160)
        videos_per_file: Number of video chunks per output file (default: 100)
    
    Returns:
        List of paths to created ArrayRecord files
    """
    os.makedirs(output_folder, exist_ok=True)
    
    # Find all ArrayRecord files
    arrayrecord_files = list(Path(input_folder).glob("*.array_record"))
    if not arrayrecord_files:
        raise ValueError(f"No ArrayRecord files found in {input_folder}")
    
    print(f"Found {len(arrayrecord_files)} ArrayRecord files")
    
    # Collect all video chunks
    all_chunks = []
    
    for file_path in arrayrecord_files:
        print(f"Processing {file_path.name}...")
        reader = ArrayRecordReader(str(file_path))
        
        for record_bytes in reader:
            # Parse the record
            record = pickle.loads(record_bytes)
            video_data = record['raw_video']
            sequence_length = record['sequence_length']
            
            # Reshape video data
            video_shape = (sequence_length, -1)  # We'll determine actual shape from data
            video_tensor = np.frombuffer(video_data, dtype=np.uint8)
            
            # Try to infer the original shape (assuming square-ish dimensions)
            total_pixels = video_tensor.shape[0] // sequence_length
            # 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
            else:
                # Fallback: assume 3 channels and calculate width
                channels = 3
                side_length = int(np.sqrt(total_pixels // channels))
            
            height = width = side_length
            video_shape = (sequence_length, height, width, channels)
            video_tensor = video_tensor.reshape(video_shape)
            
            # Split video into chunks
            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
            
            # Create overlapping chunks with stride of chunk_size//2 for better coverage
            stride = chunk_size // 2
            for start_idx in range(0, current_episode_len - chunk_size + 1, stride):
                chunk = video_tensor[start_idx:start_idx + chunk_size]
                
                # Create new record for this chunk
                chunk_record = {
                    'raw_video': chunk.tobytes(),
                    'sequence_length': chunk_size,
                    'original_sequence_length': sequence_length,
                    'chunk_start_idx': start_idx,
                }
                
                all_chunks.append(chunk_record)
        
        reader.close()
    
    print(f"Created {len(all_chunks)} video chunks")
    
    # Group chunks into files
    output_files = []
    for i in range(0, len(all_chunks), videos_per_file):
        batch_chunks = all_chunks[i:i + videos_per_file]
        output_filename = f"chunked_videos_{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"Reprocessing complete! Created {len(output_files)} new ArrayRecord files")
    return output_files