[tab] data/jasmine_data/atari/visualize_array_record.py (python)
import os
import math
import argparse
import pickle
import grain
from typing import Optional, List

import numpy as np

from PIL import Image, ImageDraw


def infer_hw_from_bytes(total_elements: int, seq_len: int, channels: int) -> tuple[int, int]:
    assert seq_len > 0 and channels > 0, "sequence_length and channels must be positive"
    base = total_elements // (seq_len * channels)
    side = int(math.isqrt(base))
    if side * side != base:
        raise ValueError(
            f"Could not infer square HxW from buffer. elements={total_elements}, seq_len={seq_len}, channels={channels}"
        )
    return side, side


def load_one_sequence_via_grain(
    data_dir: str,
    seq_len: int,
    image_h: int,
    image_w: int,
    image_c: int,
    seed: Optional[int] = None,
) -> dict:
    array_record_files = [
        os.path.join(data_dir, x)
        for x in os.listdir(data_dir)
        if x.endswith(".array_record")
    ]
    if not array_record_files:
        raise RuntimeError(f"No .array_record files found in {data_dir}")

    dataloader = get_dataloader(
        array_record_files,
        seq_len,
        1,  # global batch size 1 for visualization
        image_h,
        image_w,
        image_c,
        num_workers=1,
        prefetch_buffer_size=1,
        seed=0 if seed is None else int(seed),
    )

    state = dataloader._create_initial_state()
    it = grain.DataLoaderIterator(dataloader, state)
    batch = next(it)
    videos = np.asarray(batch["videos"][0])  # (T,H,W,C) uint8
    actions = None
    if "actions" in batch:
        actions = np.asarray(batch["actions"][0]).reshape(-1)

    return {"videos": videos, "actions": actions, "sequence_length": int(videos.shape[0])}


def get_action_meanings(env_id: Optional[str]) -> Optional[List[str]]:
    if env_id is None:
        return None
    try:
        import gymnasium as gym
        import ale_py  # noqa: F401

        gym.register_envs(ale_py)  # ensure Atari envs are registered
        env = gym.make(env_id)
        try:
            meanings = list(env.unwrapped.get_action_meanings())  # type: ignore[attr-defined]
        finally:
            env.close()
        return meanings
    except Exception:
        return None


def save_frames_with_actions(
    record: dict,
    output_dir: str,
    channels: int,
    height: Optional[int],
    width: Optional[int],
    action_meanings: Optional[List[str]],
    fps: int,
) -> None:
    os.makedirs(output_dir, exist_ok=True)

    actions: Optional[np.ndarray] = record.get("actions")
    if actions is not None:
        actions = np.asarray(actions).reshape(-1)

    seq_len = int(record["sequence_length"])
    raw_video: bytes = record["raw_video"]
    arr = np.frombuffer(raw_video, dtype=np.uint8)
    total_elements = arr.size
    if height is not None and width is not None:
        h, w = int(height), int(width)
        expected = seq_len * h * w * channels
        assert expected == total_elements, f"Expected {expected} elements, got {total_elements}"
    else:
        h, w = infer_hw_from_bytes(total_elements, seq_len, channels)
    frames = arr.reshape(seq_len, h, w, channels)

    if actions is not None:
        assert actions.shape[0] == seq_len, f"Expected {seq_len} actions, got {actions.shape[0]}"

    # Save an actions.txt index for quick inspection
    if actions is not None:
        with open(os.path.join(output_dir, "actions.txt"), "w") as f:
            for t in range(seq_len):
                a = int(actions[t])
                name = (
                    action_meanings[a]
                    if action_meanings is not None and 0 <= a < len(action_meanings)
                    else None
                )
                f.write(f"{t}	{a}" + (f"	{name}" if name is not None else "") + "
")

    # Build frames with overlays and save as GIF
    duration_ms = max(1, int(1000 / max(1, fps)))
    imgs: List[Image.Image] = []
    for t in range(seq_len):
        # frames is (T, H, W, C) so we can use it directly for RGB
        img_np_rgb = frames[t]
        img = Image.fromarray(img_np_rgb, mode="RGB")
        draw = ImageDraw.Draw(img)
        if actions is not None:
            a = int(actions[t])
            name = (
                action_meanings[a]
                if action_meanings is not None and 0 <= a < len(action_meanings)
                else None
            )
            text = f"{a}" if name is None else f"{a} {name}"
        else:
            text = "?"
        draw.text((2, 2), text, fill=(255, 255, 255))
        imgs.append(img)

    if not imgs:
        raise RuntimeError("No frames to render.")

    os.makedirs(output_dir, exist_ok=True)
    gif_path = os.path.join(output_dir, "sequence.gif")

    imgs[0].save(
        gif_path,
        save_all=True,
        append_images=imgs[1:],
        duration=duration_ms,
        loop=0,
    )

    print(f"Saved GIF to {gif_path} with {seq_len} frames (H={h}, W={w}, C={channels}, fps={fps}).")


class EpisodeLengthFilter(grain.transforms.Filter):
    """
    A Grain Filter that keeps only episodes with sufficient length.
    """

    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with sequence length requirements."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c

    def filter(self, element: Any) -> bool:
        """
        Filters episodes based on length.

        Args:
            element: A dictionary representing one record from the DataSource.
                     Expected to contain 'raw_video' (bytes) and 'sequence_length' (int)

        Returns:
            True if the episode has sufficient length, False otherwise.
        """
        assert isinstance(element, bytes)
        element = pickle.loads(element)

        current_episode_len = element["sequence_length"]
        if current_episode_len < self.seq_len:
            print(
                f"Filtering out episode with length {current_episode_len}, which is "
                f"shorter than the requested sequence length {self.seq_len}."
            )
            return False

        return True


class ProcessEpisodeAndSlice(grain.transforms.RandomMap):
    """
    A Grain Transformation that combines parsing, slicing, and normalizing.
    """

    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the transformation with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c

    def random_map(self, element: dict, rng: np.random.Generator) -> Any:
        """
        Processes a single raw episode from the data source.

        Args:
            element: A dictionary representing one record from the DataSource.
                     Expected to contain 'raw_video' (bytes) and 'sequence_length' (int)
            rng: A per-record random number generator provided by the Grain sampler.

        Returns:
            A processed video sequence as a NumPy array with shape
            (seq_len, height, width, channels) and dtype float32.
        """
        assert isinstance(element, bytes)
        element = pickle.loads(element)

        video_shape = (
            element["sequence_length"],
            self.image_h,
            self.image_w,
            self.image_c,
        )
        episode_tensor = np.frombuffer(element["raw_video"], dtype=np.uint8)
        episode_tensor = episode_tensor.reshape(video_shape)

        current_episode_len = episode_tensor.shape[0]
        if current_episode_len < self.seq_len:
            raise ValueError(
                f"Episode length {current_episode_len} is shorter than "
                f"requested sequence length {self.seq_len}. This should "
                f"have been filtered out."
            )

        max_start_idx = current_episode_len - self.seq_len

        start_idx = rng.integers(0, max_start_idx + 1)

        seq = episode_tensor[start_idx : start_idx + self.seq_len]

        data_dict = {"videos": seq}
        if "actions" in element.keys():
            actions_tensor = np.array(element["actions"])
            actions = actions_tensor[start_idx : start_idx + self.seq_len]
            data_dict["actions"] = actions

        return data_dict


def get_dataloader(
    array_record_paths: list[str],
    seq_len: int,
    global_batch_size: int,
    image_h: int,
    image_w: int,
    image_c: int,
    num_workers: int = 1,
    prefetch_buffer_size: int = 1,
    seed: int = 42,
):
    """
    Creates a data loading pipeline using Grain.
    """
    if not array_record_paths:
        raise ValueError("array_record_paths list cannot be empty.")

    per_process_batch_size = global_batch_size

    source = grain.sources.ArrayRecordDataSource(array_record_paths)

    sampler = grain.samplers.IndexSampler(
        num_records=len(source),
        shard_options=grain.sharding.ShardByJaxProcess(drop_remainder=True),
        shuffle=True,
        num_epochs=None,
        seed=seed,
    )

    operations = [
        EpisodeLengthFilter(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        ProcessEpisodeAndSlice(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        grain.transforms.Batch(batch_size=per_process_batch_size, drop_remainder=True),
    ]

    read_options = grain.ReadOptions(
        prefetch_buffer_size=prefetch_buffer_size,
        num_threads=1,
    )
    dataloader = grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=num_workers,
        worker_buffer_size=1,
        read_options=read_options,
    )

    return dataloader


def main():
    parser = argparse.ArgumentParser(
        description="Visualize a random sequence from an ArrayRecord by saving frames with actions."
    )
    parser.add_argument("--data_dir", required=True, help="Directory containing .array_record files")
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Directory to save output images and actions.txt",
    )
    parser.add_argument("--seq_len", type=int, default=16, help="Sequence length to sample with Grain")
    parser.add_argument("--seed", type=int, default=0, help="Random seed for sampling")
    parser.add_argument("--channels", type=int, default=3, help="Number of channels in frames (default 3 for RGB)")
    parser.add_argument("--height", type=int, default=84, help="Frame height (if mismatch, will be inferred)")
    parser.add_argument("--width", type=int, default=84, help="Frame width (if mismatch, will be inferred)")
    parser.add_argument("--env_id", type=str, default=None, help="Gymnasium env id to map action indices to names")
    parser.add_argument("--fps", type=int, default=10, help="GIF frames per second")

    args = parser.parse_args()

    assert args.channels == 3, "Only 3 channels are currently supported"

    record = load_one_sequence_via_grain(
        data_dir=args.data_dir,
        seq_len=int(args.seq_len),
        image_h=int(args.height),
        image_w=int(args.width),
        image_c=int(args.channels),
        seed=int(args.seed),
    )

    # Print quick summary
    seq_len = int(record.get("sequence_length", -1))
    has_actions = "actions" in record and record["actions"] is not None
    print(
        f"Loaded record with sequence_length={seq_len}, actions_present={has_actions}"
    )

    action_meanings = get_action_meanings(args.env_id)

    save_frames_with_actions(
        record=record,
        output_dir=args.output_dir,
        channels=int(args.channels),
        height=int(args.height) if args.height else None,
        width=int(args.width) if args.width else None,
        action_meanings=action_meanings,
        fps=int(args.fps),
    )


if __name__ == "__main__":
    main()