p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] data/jasmine_data/ViZDoomPPO/load_model_generate_dataset.py (python)

    # --- Save metadata ---
    metadata = {
        "env": "coinrun",
        "num_actions": get_action_space(),
        "num_episodes_train": args.num_episodes_train,
        "num_episodes_val": args.num_episodes_val,
        "num_episodes_test": args.num_episodes_test,
        "avg_episode_len_train": np.mean(
            [ep["avg_seq_len"] for ep in train_episode_metadata]
        ),
        "avg_episode_len_val": np.mean(
            [ep["avg_seq_len"] for ep in val_episode_metadata]
        ),
        "avg_episode_len_test": np.mean(
            [ep["avg_seq_len"] for ep in test_episode_metadata]
        ),
        "episode_metadata_train": train_episode_metadata,
        "episode_metadata_val": val_episode_metadata,
        "episode_metadata_test": test_episode_metadata,
    }
    with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f)

    print(f"Done generating dataset.")
[tab] slurm/jobs/mihir/horeka/preprocessing/doom_chunked.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=48:00:00
#SBATCH --partition=accelerated
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/doom/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/doom/%x_%j.log
#SBATCH --job-name=preprocess_coinrun_chunked

# export PYTHONUNBUFFERED=1


# source .venv/bin/activate

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/doom_episodes_50m \
    --agent_path /home/hk-project-p0023960/tum_cte0515/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 50_000 \
[tab] dataset_duplicates.ipynb (python)
base = "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/dev"

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.")
[tab] dataset_duplicates.ipynb (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
[tab] dataset_duplicates.ipynb (python)
# --- 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) -> tuple[str, str] | None:
    """
    Reads a single array_record file, extracts the video, hashes it,
    and returns the (hash, video_path) tuple.
    """
    try:
        reader = ArrayRecordReader(video_path)
        record_data = reader.read()
        record_unpickled = pickle.loads(record_data)
        video_bytes = record_unpickled["raw_video"]
        video_hash = hash_byte_data(video_bytes)
        return video_hash, video_path
    except Exception as e:
        print(f"Error processing file {video_path}: {e}")
        return None


def get_frame_level_hashes(video_path: str) -> tuple[list[str], str] | None:
    """
    Reads a single array_record file, extracts the frames, hashes each frame,
    and returns the (list of hashes, video_path) tuple.
    """
    try:
        reader = ArrayRecordReader(video_path)
        record_data = pickle.loads(reader.read())

        # 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) 

        frame_hashes = [hash_numpy_frame(frame) for frame in episode_tensor]

        return 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")
    ]