[tab] verify_optimized_arrayrecords.py (python)
#!/usr/bin/env python3

import os
import pickle
import numpy as np
from pathlib import Path
from array_record.python.array_record_module import ArrayRecordReader
import random

def verify_arrayrecord_file(file_path: str, expected_chunk_size: int = 160, 
                           expected_videos_per_file: int = 100,
                           image_h: int = 90, image_w: int = 160, image_c: int = 3):
    """
    Verify a single array record file for correctness.
    
    Args:
        file_path: Path to the array record file
        expected_chunk_size: Expected number of frames per video chunk
        expected_videos_per_file: Expected number of videos per file
        image_h, image_w, image_c: Image dimensions
    
    Returns:
        dict with verification results
    """
    print(f"Verifying {os.path.basename(file_path)}...")
    
    reader = ArrayRecordReader(file_path)
    all_record_bytes = reader.read_all()
    reader.close()
    
    num_records = len(all_record_bytes)
    print(f"  - Found {num_records} records")
    
    # Check if number of records is reasonable (should be <= expected_videos_per_file)
    if num_records > expected_videos_per_file:
        print(f"  - WARNING: More records than expected ({num_records} > {expected_videos_per_file})")
    
    # Sample a few records to verify structure
    sample_indices = random.sample(range(min(num_records, 5)), min(num_records, 5))
    
    verification_results = {
        'file_path': file_path,
        'num_records': num_records,
        'sample_checks': [],
        'overall_valid': True
    }
    
    for idx in sample_indices:
        try:
            record_bytes = all_record_bytes[idx]
            record = pickle.loads(record_bytes)
            
            # Check required keys
            if 'raw_video' not in record or 'sequence_length' not in record:
                print(f"  - ERROR: Record {idx} missing required keys")
                verification_results['overall_valid'] = False
                continue
            
            # Check sequence length
            sequence_length = record['sequence_length']
            if sequence_length != expected_chunk_size:
                print(f"  - ERROR: Record {idx} has wrong sequence length: {sequence_length} != {expected_chunk_size}")
                verification_results['overall_valid'] = False
                continue
            
            # Check video data
            video_data = record['raw_video']
            if not isinstance(video_data, bytes):
                print(f"  - ERROR: Record {idx} video data is not bytes")
                verification_results['overall_valid'] = False
                continue
            
            # Check video tensor shape
            video_tensor = np.frombuffer(video_data, dtype=np.uint8)
            expected_shape = expected_chunk_size * image_h * image_w * image_c
            if video_tensor.shape[0] != expected_shape:
                print(f"  - ERROR: Record {idx} video tensor shape mismatch: {video_tensor.shape[0]} != {expected_shape}")
                verification_results['overall_valid'] = False
                continue
            
            # Check value range (should be 0-255 for uint8)
            if video_tensor.min() < 0 or video_tensor.max() > 255:
                print(f"  - ERROR: Record {idx} video values out of range: [{video_tensor.min()}, {video_tensor.max()}]")
                verification_results['overall_valid'] = False
                continue
            
            # Reshape and verify dimensions
            reshaped_video = video_tensor.reshape(expected_chunk_size, image_h, image_w, image_c)
            if reshaped_video.shape != (expected_chunk_size, image_h, image_w, image_c):
                print(f"  - ERROR: Record {idx} reshaped video has wrong shape: {reshaped_video.shape}")
                verification_results['overall_valid'] = False
                continue
            
            verification_results['sample_checks'].append({
                'record_idx': idx,
                'sequence_length': sequence_length,
                'video_shape': reshaped_video.shape,
                'value_range': (video_tensor.min(), video_tensor.max()),
                'valid': True
            })
            
            print(f"  - Record {idx}: ✓ {sequence_length} frames, shape {reshaped_video.shape}, range [{video_tensor.min()}, {video_tensor.max()}]")
            
        except Exception as e:
            print(f"  - ERROR: Failed to process record {idx}: {e}")
            verification_results['overall_valid'] = False
            verification_results['sample_checks'].append({
                'record_idx': idx,
                'error': str(e),
                'valid': False
            })
    
    return verification_results

def main():
    """Main verification function."""
    output_folder = "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_optimized_layout"
    
    # Expected parameters from the reprocessing function
    expected_chunk_size = 160
    expected_videos_per_file = 100
    image_h, image_w, image_c = 90, 160, 3
    
    print("=== ArrayRecord Verification Report ===")
    print(f"Output folder: {output_folder}")
    print(f"Expected chunk size: {expected_chunk_size}")
    print(f"Expected videos per file: {expected_videos_per_file}")
    print(f"Image dimensions: {image_h}x{image_w}x{image_c}")
    print()
    
    # Get all array record files
    arrayrecord_files = list(Path(output_folder).glob("*.array_record"))
    if not arrayrecord_files:
        print("ERROR: No array record files found!")
        return
    
    print(f"Found {len(arrayrecord_files)} array record files")
    print()
    
    # Verify a sample of files (first 5 and last 5)
    sample_files = arrayrecord_files[:5] + arrayrecord_files[-5:]
    
    all_results = []
    for file_path in sample_files:
        result = verify_arrayrecord_file(
            str(file_path), 
            expected_chunk_size, 
            expected_videos_per_file,
            image_h, image_w, image_c
        )
        all_results.append(result)
        print()
    
    # Summary
    print("=== Summary ===")
    valid_files = sum(1 for r in all_results if r['overall_valid'])
    total_files_checked = len(all_results)
    
    print(f"Files checked: {total_files_checked}")
    print(f"Valid files: {valid_files}")
    print(f"Invalid files: {total_files_checked - valid_files}")
    
    if valid_files == total_files_checked:
        print("✓ All checked files are valid!")
    else:
        print("✗ Some files have issues!")
    
    # Check file size distribution
    print("
=== File Size Analysis ===")
    file_sizes = []
    for file_path in arrayrecord_files:
        size_mb = os.path.getsize(file_path) / (1024 * 1024)
        file_sizes.append(size_mb)
    
    print(f"Total files: {len(file_sizes)}")
    print(f"Average file size: {np.mean(file_sizes):.1f} MB")
    print(f"Min file size: {np.min(file_sizes):.1f} MB")
    print(f"Max file size: {np.max(file_sizes):.1f} MB")
    print(f"Total data size: {np.sum(file_sizes):.1f} MB ({np.sum(file_sizes)/1024:.1f} GB)")

if __name__ == "__main__":
    main()