edessa/EMG

1#!/usr/bin/env python32# -*- coding: utf-8 -*-3 4"""5Take a .vrs file and create a sibling directory called "frames" that contains the frames.6 7Usage:8  python save_vrs_frames.py /path/to/recording.vrs9 10Behavior:11- Saves frames from a single image stream (default preference: camera-rgb).12- Output directory is: <dir_of_vrs>/frames13- Output images are: frame_000000.jpg, frame_000001.jpg, ...14- Also writes: <dir_of_vrs>/frames/timestamps.csv (frame_idx, timestamp_ns, stream_label)15 16Requires:17  pip install projectaria-tools opencv-python18"""19 20import os21import sys22import csv23import cv224from projectaria_tools.core import data_provider25 26 27PREFERRED_LABELS = [28    "camera-rgb",29    "camera-slam-left",30    "camera-slam-right",31]32 33 34def pick_stream_id(provider):35    """36    Pick the first available image stream from PREFERRED_LABELS.37    Returns (stream_id, label).38    """39    for label in PREFERRED_LABELS:40        sid = provider.get_stream_id_from_label(label)41        if sid is not None:42            # sanity: can we read at least one image?43            try:44                n = provider.get_num_data(sid)45                if n <= 0:46                    continue47                img, rec = provider.get_image_data_by_index(sid, 0)48                if img is not None and rec is not None:49                    return sid, label50            except Exception:51                continue52 53    # If none of the preferred labels exist, try *any* stream in get_all_streams()54    # and see if it yields images.55    try:56        streams = provider.get_all_streams()  # list-like57    except Exception as e:58        raise RuntimeError(59            "Could not find an image stream. Tried camera-rgb and slam cameras, "60            "and provider.get_all_streams() was unavailable."61        ) from e62 63    for s in streams:64        # best-effort to obtain the stream id and label from whatever object type this is65        sid = None66        label = None67        for attr in ("stream_id", "id", "streamId"):68            if hasattr(s, attr):69                sid = getattr(s, attr)70                break71        for attr in ("label", "name"):72            if hasattr(s, attr):73                label = getattr(s, attr)74                break75        if sid is None:76            continue77 78        try:79            n = provider.get_num_data(sid)80            if n <= 0:81                continue82            img, rec = provider.get_image_data_by_index(sid, 0)83            if img is not None and rec is not None:84                return sid, (label or str(sid))85        except Exception:86            continue87 88    raise RuntimeError("No image stream found in this VRS (nothing readable as images).")89 90 91def to_numpy(img_data):92    """93    projectaria_tools sometimes returns an object with .to_numpy_array()94    """95    if hasattr(img_data, "to_numpy_array"):96        return img_data.to_numpy_array()97    return img_data98 99 100def main():101    if len(sys.argv) != 2:102        print("Usage: python save_vrs_frames.py /path/to/recording.vrs")103        sys.exit(2)104 105    vrs_path = sys.argv[1]106    if not (os.path.isfile(vrs_path) and vrs_path.lower().endswith(".vrs")):107        raise FileNotFoundError(f"Not a .vrs file: {vrs_path}")108 109    vrs_dir = os.path.dirname(os.path.abspath(vrs_path))110    out_dir = os.path.join(vrs_dir, "frames")111    os.makedirs(out_dir, exist_ok=True)112 113    provider = data_provider.create_vrs_data_provider(vrs_path)114 115    stream_id, stream_label = pick_stream_id(provider)116    num = provider.get_num_data(stream_id)117 118    ts_csv_path = os.path.join(out_dir, "timestamps.csv")119    with open(ts_csv_path, "w", newline="") as f:120        w = csv.writer(f)121        w.writerow(["frame_idx", "timestamp_ns", "stream_label"])122 123        for i in range(num):124            img_data, rec = provider.get_image_data_by_index(stream_id, i)125            if img_data is None or rec is None:126                continue127 128            img = to_numpy(img_data)129 130            # If RGB -> convert to BGR for OpenCV saving131            if img.ndim == 3 and img.shape[2] == 3:132                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)133 134            out_path = os.path.join(out_dir, f"frame_{i:06d}.jpg")135            ok = cv2.imwrite(out_path, img, [int(cv2.IMWRITE_JPEG_QUALITY), 95])136            if not ok:137                raise RuntimeError(f"Failed to write: {out_path}")138 139            ts_ns = int(getattr(rec, "capture_timestamp_ns", -1))140            w.writerow([i, ts_ns, stream_label])141 142            if i % 500 == 0:143                print(f"  saved {i}/{num} frames...")144 145    print(f"[OK] VRS: {vrs_path}")146    print(f"[OK] Stream: {stream_label} (id={stream_id})")147    print(f"[OK] Frames saved to: {out_dir}")148    print(f"[OK] Timestamps: {ts_csv_path}")149 150 151if __name__ == "__main__":152    main()