[tab] dataset_duplicates.py (python)
import os
import pickle
import hashlib
import numpy as np
from collections import defaultdict
from array_record.python.array_record_module import ArrayRecordReader
import multiprocessing
from tqdm import tqdm # Using tqdm for a nice progress bar

# --- Helper Functions ---

def hash_byte_data(bytes_data: bytes) -> str:
    """Calculates the SHA256 hash of a byte string."""
    return hashlib.sha256(bytes_data).hexdigest()


def hash_numpy_frame(frame: np.ndarray) -> str:
    """Calculates the SHA256 hash of a numpy array."""
    return hashlib.sha256(np.ascontiguousarray(frame)).hexdigest()


def get_episode_level_hash(video_path: str) -> list[tuple[str, str]] | None:
    """
    Reads all records (chunks) from a single array_record file, concatenates the episode
    video bytes across chunks, hashes the full episode, and returns a single
    (hash, video_path) entry in a list.
    """
    try:
        reader = ArrayRecordReader(video_path)
        concatenated_bytes_parts: list[bytes] = []
        for _ in range(reader.num_records()):
            record_data = reader.read()
            if not record_data:
                break
            record_unpickled = pickle.loads(record_data)
            concatenated_bytes_parts.append(record_unpickled["raw_video"])
        if not concatenated_bytes_parts:
            return []
        full_video_bytes = b"".join(concatenated_bytes_parts)
        video_hash = hash_byte_data(full_video_bytes)
        return [(video_hash, video_path)]
    except Exception as e:
        return None


def get_frame_level_hashes(video_path: str) -> list[tuple[list[str], str]] | None:
    """
    Reads all records (chunks) from a single array_record file, decodes frames per chunk,
    concatenates frames to form the full episode, hashes each frame, and returns a single
    (list of frame hashes, video_path) entry in a list.
    """
    try:
        reader = ArrayRecordReader(video_path)
        all_frame_hashes: list[str] = []
        for _ in range(reader.num_records()):
            raw = reader.read()
            if not raw:
                break
            record_data = pickle.loads(raw)

            # video shape (seq_len, 64, 64, 3)
            video_shape = (record_data["sequence_length"], 64, 64, 3)
            episode_tensor = np.frombuffer(record_data["raw_video"], dtype=np.uint8)
            episode_tensor = episode_tensor.reshape(video_shape)

            for frame in episode_tensor:
                all_frame_hashes.append(hash_numpy_frame(frame))

        if not all_frame_hashes:
            return []
        return [(all_frame_hashes, video_path)]
    except Exception as e:
        print(f"Error processing file {video_path}: {e}")
        return None

def get_array_record_files(dir):
    return [
        os.path.join(dir, x)
        for x in os.listdir(dir)
        if x.endswith(".array_record")
    ]


def get_split_from_path(path: str) -> str:
    """
    Returns the dataset split inferred from the file path (e.g., 'train', 'val', or 'test').
    Assumes files are organized as .../array_record/<split>/<file>.array_record
    """
    return os.path.basename(os.path.dirname(path))

base = "/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/"

train_dir = os.path.join(base, "train")
test_dir = os.path.join(base, "test")
val_dir = os.path.join(base, "val")

train_array_record_files = get_array_record_files(train_dir)
test_array_record_files = get_array_record_files(test_dir)
val_array_record_files = get_array_record_files(val_dir)

array_record_files = train_array_record_files + test_array_record_files + val_array_record_files
print(f"Found {len(array_record_files)} files to process.")

num_processes = multiprocessing.cpu_count()
print(f"Using {num_processes} worker processes.")

# Find episode level duplicates
duplicate_episode = defaultdict(list)

# The 'with' statement ensures the pool is properly closed
with multiprocessing.Pool(processes=num_processes) as pool:
    # Use pool.imap_unordered for efficiency. It processes items as they are submitted
    # and returns results as they complete, which is perfect for progress bars.
    # We wrap the result iterator with tqdm to show progress.
    results = pool.imap_unordered(get_episode_level_hash, array_record_files)

    print("Processing files and calculating hashes...")
    for result in tqdm(results, total=len(array_record_files)):
        if result:  # Ensure the worker didn't return None due to an error
            for video_hash, video_identifier in result:
                duplicate_episode[video_hash].append(video_identifier)

print("
Aggregation complete. Finding duplicates...")
duplicates = {h: paths for h, paths in duplicate_episode.items() if len(paths) > 1}


print(f"
Found {len(duplicates)} sets of duplicate videos.")
if duplicates:
    # Print the first 5 duplicate sets as an example
    for i, (h, paths) in enumerate(duplicates.items()):
        print(f"  - Hash: {h[:10]}... ({len(paths)} files)")
        for path in paths:
            print(f"    - {os.path.basename(path)}")


# Frame level duplicates
# This dictionary will map a frame hash to a list of (video_path, frame_index)
frame_dup_dict = defaultdict(list)

with multiprocessing.Pool(processes=num_processes) as pool:
    results = pool.imap_unordered(get_frame_level_hashes, array_record_files)

    print("Processing files and calculating frame hashes...")
    for result in tqdm(results, total=len(array_record_files)):
        if result:
            for frame_hashes, video_identifier in result:
                for frame_idx, frame_hash in enumerate(frame_hashes):
                    frame_dup_dict[frame_hash].append((video_identifier, frame_idx))

print("
Aggregation complete. Finding duplicate frames...")
duplicate_frames = {hash: location for hash, location in frame_dup_dict.items() if len(location) > 1}

total_frames = sum(len(locations) for locations in frame_dup_dict.values())
print(f"Total frames: {total_frames}")

num_duplicate_frames = len(duplicate_frames.keys())
print(f"Number of duplicate frames: {num_duplicate_frames}")

percentage = num_duplicate_frames / total_frames
print(f"Percentage of duplicate frames: {percentage:.2%}")


# --- Train leakage checks ---
# Episode-level leakage: any episode hash that appears in both train and val
leaking_episode_hashes = {}
for h, paths in duplicate_episode.items():
    splits = {get_split_from_path(p) for p in paths}
    if "train" in splits and "val" in splits:
        leaking_episode_hashes[h] = paths

print("
Train→Val leakage at episode level:")
if leaking_episode_hashes:
    print(f"  - {len(leaking_episode_hashes)} episode hashes appear in both train and val.")
    example_count = 5
    for i, (h, paths) in enumerate(list(leaking_episode_hashes.items())[:example_count]):
        print(f"    * Hash: {h[:10]}... ({len(paths)} files)")
        for path in paths:
            print(f"      - [{get_split_from_path(path)}] {os.path.basename(path)}")
else:
    print("  - None detected.")

# Frame-level leakage: any frame hash that appears in both train and val
leaking_frame_hashes = {}
for fh, locations in duplicate_frames.items():
    splits = {get_split_from_path(p) for (p, _) in locations}
    if "train" in splits and "val" in splits:
        leaking_frame_hashes[fh] = locations

num_leaking_frame_hashes = len(leaking_frame_hashes)
num_leaking_val_frames = sum(1 for locations in leaking_frame_hashes.values() for (p, _) in locations if get_split_from_path(p) == "val")
num_leaking_train_frames = sum(1 for locations in leaking_frame_hashes.values() for (p, _) in locations if get_split_from_path(p) == "train")

print("
Train→Val leakage at frame level:")
if num_leaking_frame_hashes > 0:
    print(f"  - {num_leaking_frame_hashes} frame hashes appear in both train and val.")
    print(f"  - Val frames involved: {num_leaking_val_frames}")
    print(f"  - Train frames involved: {num_leaking_train_frames}")
    example_count = 5
    for i, (fh, locations) in enumerate(list(leaking_frame_hashes.items())[:example_count]):
        print(f"    * Frame hash: {fh[:10]}... ({len(locations)} occurrences)")
        for (p, idx) in locations[:10]:
            print(f"      - [{get_split_from_path(p)}] {os.path.basename(p)} frame {idx}")
else:
    print("  - None detected.")