chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
Executable
+221
@@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example using an example depth dataset from NYU.
|
||||
|
||||
https://cs.nyu.edu/~fergus/datasets/nyu_depth_v2.html
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import zipfile
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Final
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
import rerun as rr # pip install rerun-sdk
|
||||
import rerun.blueprint as rrb
|
||||
|
||||
DESCRIPTION = """
|
||||
# RGBD
|
||||
Visualizes an example recording from [the NYUD dataset](https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html) with RGB and Depth channels.
|
||||
|
||||
The full source code for this example is available [on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/rgbd).
|
||||
"""
|
||||
|
||||
DEPTH_IMAGE_SCALING: Final = 1e4
|
||||
DATASET_DIR: Final = Path(os.path.dirname(__file__)) / "dataset"
|
||||
DATASET_URL_BASE: Final = "https://static.rerun.io/rgbd_dataset"
|
||||
DATASET_URL_BASE_ALTERNATE: Final = "https://cs.nyu.edu/~fergus/datasets/nyu_depth_v2.html"
|
||||
AVAILABLE_RECORDINGS: Final = ["cafe", "basements", "studies", "office_kitchens", "playroooms"]
|
||||
|
||||
|
||||
def parse_timestamp(filename: str) -> datetime:
|
||||
"""Parse the timestamp portion of the filename."""
|
||||
file_name_parts = filename.split("-")
|
||||
time = file_name_parts[len(file_name_parts) - 2]
|
||||
return datetime.fromtimestamp(float(time))
|
||||
|
||||
|
||||
def read_image_bgr(buf: bytes) -> npt.NDArray[np.uint8]:
|
||||
"""Decode an image provided in `buf`, and interpret it as RGB data."""
|
||||
np_buf: npt.NDArray[np.uint8] = np.ndarray(shape=(1, len(buf)), dtype=np.uint8, buffer=buf)
|
||||
img_bgr: npt.NDArray[Any] = cv2.imdecode(np_buf, cv2.IMREAD_COLOR)
|
||||
return img_bgr
|
||||
|
||||
|
||||
def read_depth_image(buf: bytes) -> npt.NDArray[Any]:
|
||||
"""Decode an image provided in `buf`."""
|
||||
np_buf: npt.NDArray[np.uint8] = np.ndarray(shape=(1, len(buf)), dtype=np.uint8, buffer=buf)
|
||||
img: npt.NDArray[Any] = cv2.imdecode(np_buf, cv2.IMREAD_UNCHANGED)
|
||||
return img
|
||||
|
||||
|
||||
def log_nyud_data(recording_path: Path, subset_idx: int, frames: int) -> None:
|
||||
rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)
|
||||
|
||||
with zipfile.ZipFile(recording_path, "r") as archive:
|
||||
archive_dirs = [f.filename for f in archive.filelist if f.is_dir()]
|
||||
|
||||
print(f"Using recording subset {subset_idx} ([0 - {len(archive_dirs) - 1}] available).")
|
||||
|
||||
dir_to_log = archive_dirs[subset_idx]
|
||||
subset = [
|
||||
f
|
||||
for f in archive.filelist
|
||||
if f.filename.startswith(dir_to_log) and (f.filename.endswith(".ppm") or f.filename.endswith(".pgm"))
|
||||
]
|
||||
files_with_timestamps = [(parse_timestamp(f.filename), f) for f in subset]
|
||||
files_with_timestamps.sort(key=lambda t: t[0])
|
||||
|
||||
if len(files_with_timestamps) > frames:
|
||||
files_with_timestamps = files_with_timestamps[:frames]
|
||||
|
||||
for time, f in files_with_timestamps:
|
||||
rr.set_time("time", timestamp=time.timestamp())
|
||||
|
||||
if f.filename.endswith(".ppm"):
|
||||
buf = archive.read(f)
|
||||
img_bgr = read_image_bgr(buf)
|
||||
rr.log("world/camera/image/rgb", rr.Image(img_bgr, color_model="BGR").compress(jpeg_quality=95))
|
||||
|
||||
elif f.filename.endswith(".pgm"):
|
||||
buf = archive.read(f)
|
||||
img_depth = read_depth_image(buf)
|
||||
|
||||
# Log the camera transforms:
|
||||
rr.log(
|
||||
"world/camera/image",
|
||||
rr.Pinhole(
|
||||
resolution=[img_depth.shape[1], img_depth.shape[0]],
|
||||
focal_length=0.7 * img_depth.shape[1],
|
||||
# Intentionally off-center to demonstrate that we support it
|
||||
principal_point=[0.45 * img_depth.shape[1], 0.55 * img_depth.shape[0]],
|
||||
),
|
||||
)
|
||||
|
||||
# Log the depth image to the cameras image-space:
|
||||
rr.log("world/camera/image/depth", rr.DepthImage(img_depth, meter=DEPTH_IMAGE_SCALING))
|
||||
|
||||
|
||||
def ensure_recording_downloaded(name: str) -> Path:
|
||||
recording_filename = f"{name}.zip"
|
||||
recording_path = DATASET_DIR / recording_filename
|
||||
if recording_path.exists():
|
||||
return recording_path
|
||||
|
||||
url = f"{DATASET_URL_BASE}/{recording_filename}"
|
||||
alternate_url = f"{DATASET_URL_BASE_ALTERNATE}/{recording_filename}"
|
||||
|
||||
os.makedirs(DATASET_DIR, exist_ok=True)
|
||||
try:
|
||||
try:
|
||||
print(f"downloading {url} to {recording_path}")
|
||||
download_progress(url, recording_path)
|
||||
except ValueError:
|
||||
print(f"Failed to download from {url}, trying backup URL {alternate_url} instead")
|
||||
download_progress(alternate_url, recording_path)
|
||||
except BaseException:
|
||||
recording_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
return recording_path
|
||||
|
||||
|
||||
def download_progress(url: str, dst: Path) -> None:
|
||||
"""
|
||||
Download file with tqdm progress bar.
|
||||
|
||||
From: <https://gist.github.com/yanqd0/c13ed29e29432e3cf3e7c38467f42f51>
|
||||
"""
|
||||
resp = requests.get(url, stream=True)
|
||||
if resp.status_code != 200:
|
||||
raise ValueError(f"Failed to download file (status code: {resp.status_code})")
|
||||
total = int(resp.headers.get("content-length", 0))
|
||||
chunk_size = 1024 * 1024
|
||||
# Can also replace 'file' with a io.BytesIO object
|
||||
with (
|
||||
open(dst, "wb") as file,
|
||||
tqdm(
|
||||
desc=dst.name,
|
||||
total=total,
|
||||
unit="iB",
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
) as bar,
|
||||
):
|
||||
for data in resp.iter_content(chunk_size=chunk_size):
|
||||
size = file.write(data)
|
||||
bar.update(size)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Example using an example depth dataset from NYU.")
|
||||
parser.add_argument(
|
||||
"--recording",
|
||||
type=str,
|
||||
choices=AVAILABLE_RECORDINGS,
|
||||
default=AVAILABLE_RECORDINGS[0],
|
||||
help="Name of the NYU Depth Dataset V2 recording",
|
||||
)
|
||||
parser.add_argument("--subset-idx", type=int, default=0, help="The index of the subset of the recording to use.")
|
||||
parser.add_argument(
|
||||
"--frames",
|
||||
type=int,
|
||||
default=sys.maxsize,
|
||||
help="If specified, limits the number of frames logged",
|
||||
)
|
||||
rr.script_add_args(parser)
|
||||
args = parser.parse_args()
|
||||
|
||||
rr.script_setup(
|
||||
args,
|
||||
"rerun_example_rgbd",
|
||||
default_blueprint=rrb.Horizontal(
|
||||
rrb.Vertical(
|
||||
rrb.Spatial3DView(name="3D", origin="world"),
|
||||
rrb.TextDocumentView(name="Description", origin="/description"),
|
||||
row_shares=[7, 3],
|
||||
),
|
||||
rrb.Vertical(
|
||||
# Put the origin for both 2D spaces where the pinhole is logged. Doing so allows them to understand how they're connected to the 3D space.
|
||||
# This enables interactions like clicking on a point in the 3D space to show the corresponding point in the 2D spaces and vice versa.
|
||||
rrb.Spatial2DView(
|
||||
name="RGB & Depth",
|
||||
origin="world/camera/image",
|
||||
overrides={"world/camera/image/rgb": rr.Image.from_fields(opacity=0.5)},
|
||||
),
|
||||
rrb.Tabs(
|
||||
rrb.Spatial2DView(name="RGB", origin="world/camera/image", contents="world/camera/image/rgb"),
|
||||
rrb.Spatial2DView(name="Depth", origin="world/camera/image", contents="world/camera/image/depth"),
|
||||
),
|
||||
name="2D",
|
||||
row_shares=[3, 3, 2],
|
||||
),
|
||||
column_shares=[2, 1],
|
||||
),
|
||||
)
|
||||
|
||||
recording_path = ensure_recording_downloaded(args.recording)
|
||||
|
||||
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
|
||||
|
||||
log_nyud_data(
|
||||
recording_path=recording_path,
|
||||
subset_idx=args.subset_idx,
|
||||
frames=args.frames,
|
||||
)
|
||||
|
||||
rr.script_teardown(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user