wehub/rerun-io--rerun
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:05:14 +08:00
commit 2a547be7fe
7904 changed files with 1000926 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/arkit_scenes/.gitignore
+1
View File
@@ -0,0 +1 @@
dataset/**
examples/python/arkit_scenes/README.md
+132
View File
@@ -0,0 +1,132 @@
<!--[metadata]
title = "ARKit scenes"
tags = ["2D", "3D", "Depth", "Mesh", "Object detection", "Pinhole camera", "Blueprint"]
thumbnail = "https://static.rerun.io/arkit-scenes/6d920eaa42fb86cfd264d47180ecbecbb6dd3e09/480w.png"
thumbnail_dimensions = [480, 480]
channel = "main"
include_in_manifest = true
-->
This example visualizes the [ARKitScenes dataset](https://github.com/apple/ARKitScenes/) using Rerun.
The dataset contains color images, depth images, the reconstructed mesh, and labeled bounding boxes around furniture.
<picture data-inline-viewer="examples/arkit_scenes">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/1200w.png">
<img src="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/full.png" alt="ARKit Scenes screenshot">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image),
[`DepthImage`](https://www.rerun.io/docs/reference/types/archetypes/depth_image), [`Transform3D`](https://www.rerun.io/docs/reference/types/archetypes/transform3d),
[`Pinhole`](https://www.rerun.io/docs/reference/types/archetypes/pinhole), [`Mesh3D`](https://www.rerun.io/docs/reference/types/archetypes/mesh3d),
[`Boxes3D`](https://www.rerun.io/docs/reference/types/archetypes/boxes3d),
[`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
## Background
The ARKitScenes dataset, captured using Apple's ARKit technology, encompasses a diverse array of indoor scenes, offering color and depth images, reconstructed 3D meshes, and labeled bounding boxes around objects like furniture. It's a valuable resource for researchers and developers in computer vision and augmented reality, enabling advancements in object recognition, depth estimation, and spatial understanding.
## Logging and visualizing with Rerun
This visualization through Rerun highlights the dataset's potential in developing immersive AR experiences and enhancing machine learning models for real-world applications while showcasing Reruns visualization capabilities.
### Logging a moving RGB-D camera
To log a moving RGB-D camera, we log four key components: the camera's intrinsics via a pinhole camera model, its pose or extrinsics, along with the color and depth images. The camera intrinsics, which define the camera's lens properties, and the pose, detailing its position and orientation, are logged to create a comprehensive 3D to 2D mapping. Both the RGB and depth images are then logged as child entities, capturing the visual and depth aspects of the scene, respectively. This approach ensures a detailed recording of the camera's viewpoint and the scene it captures, all stored compactly under the same entity path for simplicity.
```python
rr.log("world/camera_lowres", rr.Transform3D(transform=camera_from_world))
rr.log("world/camera_lowres", rr.Pinhole(image_from_camera=intrinsic, resolution=[w, h]))
rr.log(f"{entity_id}/rgb", rr.Image(rgb).compress(jpeg_quality=95))
rr.log(f"{entity_id}/depth", rr.DepthImage(depth, meter=1000))
```
### Ground-truth mesh
The mesh is logged as an [rr.Mesh3D archetype](https://www.rerun.io/docs/reference/types/archetypes/mesh3d).
In this case the mesh is composed of mesh vertices, indices (i.e., which vertices belong to the same face), and vertex
colors.
```python
rr.log(
"world/mesh",
rr.Mesh3D(
vertex_positions=mesh.vertices,
vertex_colors=mesh.visual.vertex_colors,
triangle_indices=mesh.faces,
),
static=True,
)
```
Here, the mesh is logged to the world/mesh entity and is marked as static, since it does not change in the context of this visualization.
### Logging 3D bounding boxes
Here we loop through the data and add bounding boxes to all the items found.
```python
for i, label_info in enumerate(annotation["data"]):
rr.log(
f"world/annotations/box-{uid}-{label}",
rr.Boxes3D(
half_sizes=half_size,
centers=centroid,
labels=label,
colors=colors[i],
),
rr.InstancePoses3D(mat3x3=mat3x3),
static=True,
)
```
### Setting up the default blueprint
This example benefits at lot from having a custom blueprint defined. This happens with the following code:
```python
primary_camera_entity = HIGHRES_ENTITY_PATH if args.include_highres else LOWRES_POSED_ENTITY_PATH
blueprint = rrb.Horizontal(
rrb.Spatial3DView(name="3D"),
rrb.Vertical(
rrb.Tabs(
rrb.Spatial2DView(
name="RGB",
origin=primary_camera_entity,
contents=["$origin/rgb", "/world/annotations/**"],
),
rrb.Spatial2DView(
name="Depth",
origin=primary_camera_entity,
contents=["$origin/depth", "/world/annotations/**"],
),
name="2D",
),
rrb.TextDocumentView(name="Readme"),
),
)
rr.script_setup(args, "rerun_example_arkit_scenes", default_blueprint=blueprint)
```
In particular, we want to reproject 3D annotations onto the 2D camera views. To configure such a view, two things are necessary:
- The view origin must be set to the entity that contains the pinhole transforms. In this example, the entity path is stored in the `primary_camera_entity` variable.
- The view contents must explicitly include the annotations, which are not logged in the subtree defined by the origin. This is done using the `contents` argument, here set to `["$origin/depth", "/world/annotations/**"]`.
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/arkit_scenes
```
To run this example use
```bash
python -m arkit_scenes
```
examples/python/arkit_scenes/arkit_scenes/__init__.py
View File
examples/python/arkit_scenes/arkit_scenes/__main__.py
Executable
+322
View File
@@ -0,0 +1,322 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import os
from typing import TYPE_CHECKING, Any
import cv2
import numpy as np
import trimesh
from scipy.spatial.transform import Rotation as R
from tqdm import tqdm
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
from .download_dataset import AVAILABLE_RECORDINGS, ensure_recording_available
if TYPE_CHECKING:
from pathlib import Path
DESCRIPTION = """
# ARKitScenes
This example visualizes the [ARKitScenes dataset](https://github.com/apple/ARKitScenes/) using Rerun. The dataset
contains color images, depth images, the reconstructed mesh, and labeled bounding boxes around furniture.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/arkit_scenes).
""".strip()
Color = tuple[float, float, float, float]
# hack for now since dataset does not provide orientation information, only known after initial visual inspection
ORIENTATION = {
"48458663": "landscape",
"42444949": "portrait",
"41069046": "portrait",
"41125722": "portrait",
"41125763": "portrait",
"42446167": "portrait",
}
assert len(ORIENTATION) == len(AVAILABLE_RECORDINGS)
assert set(ORIENTATION.keys()) == set(AVAILABLE_RECORDINGS)
LOWRES_POSED_ENTITY_PATH = "world/camera_lowres"
HIGHRES_ENTITY_PATH = "world/camera_highres"
def load_json(js_path: Path) -> dict[str, Any]:
with open(js_path, encoding="utf8") as f:
json_data: dict[str, Any] = json.load(f)
return json_data
def log_annotated_bboxes(annotation: dict[str, Any]) -> None:
"""
Logs annotated oriented bounding boxes to Rerun.
annotation json file
| |-- label: object name of bounding box
| |-- axesLengths[x, y, z]: size of the origin bounding-box before transforming
| |-- centroid[]: the translation matrix (1,3) of bounding-box
| |-- normalizedAxes[]: the rotation matrix (3,3) of bounding-box
"""
for label_info in annotation["data"]:
uid = label_info["uid"]
label = label_info["label"]
half_size = 0.5 * np.array(label_info["segments"]["obbAligned"]["axesLengths"]).reshape(-1, 3)[0]
centroid = np.array(label_info["segments"]["obbAligned"]["centroid"]).reshape(-1, 3)[0]
mat3x3 = np.array(label_info["segments"]["obbAligned"]["normalizedAxes"]).reshape(3, 3).T
rr.log(
f"world/annotations/box-{uid}-{label}",
rr.Boxes3D(
half_sizes=half_size,
labels=label,
),
rr.InstancePoses3D(translations=centroid, mat3x3=mat3x3),
static=True,
)
def log_camera(
intri_path: Path,
frame_id: str,
poses_from_traj: dict[str, rr.Transform3D],
entity_id: str,
) -> None:
"""Logs camera transform and 3D bounding boxes in the image frame."""
w, h, fx, fy, cx, cy = np.loadtxt(intri_path)
intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
camera_from_world = poses_from_traj[frame_id]
# clear previous centroid labels
rr.log(f"{entity_id}/bbox-2D-segments", rr.Clear(recursive=True))
# pathlib makes it easy to get the parent, but log methods requires a string
rr.log(entity_id, camera_from_world)
rr.log(entity_id, rr.Pinhole(image_from_camera=intrinsic, resolution=[w, h]))
def read_camera_from_world(traj_string: str) -> tuple[str, rr.Transform3D]:
"""
Reads out camera_from_world transform from trajectory string.
Args:
----
traj_string:
A space-delimited file where each line represents a camera position at a particular timestamp.
The file has seven columns:
* Column 1: timestamp
* Columns 2-4: rotation (axis-angle representation in radians)
* Columns 5-7: translation (usually in meters)
Returns
-------
timestamp: float
timestamp in seconds
camera_from_world: tuple of two numpy arrays
A tuple containing a translation vector and a quaternion that represent the camera_from_world transform
Raises
------
AssertionError:
If the input string does not contain 7 tokens.
"""
tokens = traj_string.split() # Split the input string into tokens
assert len(tokens) == 7, f"Input string must have 7 tokens, but found {len(tokens)}."
ts: str = tokens[0] # Extract timestamp from the first token
# Extract rotation from the second to fourth tokens
angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
rotation = R.from_rotvec(np.asarray(angle_axis))
# Extract translation from the fifth to seventh tokens
translation = np.asarray([float(tokens[4]), float(tokens[5]), float(tokens[6])])
# Create tuple in format log_transform3d expects
camera_from_world = rr.Transform3D(
translation=translation,
rotation=rr.Quaternion(xyzw=rotation.as_quat()),
relation=rr.TransformRelation.ChildFromParent,
)
return (ts, camera_from_world)
def find_closest_frame_id(target_id: str, frame_ids: dict[str, Any]) -> str:
"""Finds the closest frame id to the target id."""
target_value = float(target_id)
closest_id = min(frame_ids.keys(), key=lambda x: abs(float(x) - target_value))
return closest_id
def log_arkit(recording_path: Path, include_highres: bool) -> None:
"""
Logs ARKit recording data using Rerun.
Args:
----
recording_path (Path):
The path to the ARKit recording.
include_highres (bool):
Whether to include high resolution data.
Returns
-------
None
"""
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
video_id = recording_path.stem
lowres_image_dir = recording_path / "lowres_wide"
image_dir = recording_path / "wide"
lowres_depth_dir = recording_path / "lowres_depth"
depth_dir = recording_path / "highres_depth"
lowres_intrinsics_dir = recording_path / "lowres_wide_intrinsics"
intrinsics_dir = recording_path / "wide_intrinsics"
traj_path = recording_path / "lowres_wide.traj"
# frame_ids are indexed by timestamps, you can see more info here
# https://github.com/apple/ARKitScenes/blob/main/threedod/README.md#data-organization-and-format-of-input-data
depth_filenames = [x.name for x in sorted(lowres_depth_dir.iterdir())]
lowres_frame_ids = [x.split(".png")[0].split("_")[1] for x in depth_filenames]
lowres_frame_ids.sort()
# dict of timestamp to pose which is a tuple of translation and quaternion
camera_from_world_dict = {}
with open(traj_path, encoding="utf-8") as f:
trajectory = f.readlines()
for line in trajectory:
timestamp, camera_from_world = read_camera_from_world(line)
# round timestamp to 3 decimal places as seen in the original repo here
# https://github.com/apple/ARKitScenes/blob/e2e975128a0a9695ea56fa215fe76b4295241538/threedod/benchmark_scripts/utils/tenFpsDataLoader.py#L247
timestamp = f"{round(float(timestamp), 3):.3f}"
camera_from_world_dict[timestamp] = camera_from_world
rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
ply_path = recording_path / f"{recording_path.stem}_3dod_mesh.ply"
print(f"Loading {ply_path}…")
assert os.path.isfile(ply_path), f"Failed to find {ply_path}"
mesh = trimesh.load(str(ply_path))
rr.log(
"world/mesh",
rr.Mesh3D(
vertex_positions=mesh.vertices, # type: ignore[attr-defined]
vertex_colors=mesh.visual.vertex_colors, # type: ignore[attr-defined]
triangle_indices=mesh.faces, # type: ignore[attr-defined]
face_rendering="Front", # We want to hide the front facing faces, but the dataset uses mostly clockwise winding order which is the opposite of what Rerun assumes (CCW).
),
static=True,
)
# load the obb annotations and log them in the world frame
bbox_annotations_path = recording_path / f"{recording_path.stem}_3dod_annotation.json"
annotation = load_json(bbox_annotations_path)
log_annotated_bboxes(annotation)
print("Processing frames…")
for frame_timestamp in tqdm(lowres_frame_ids):
# frame_id is equivalent to timestamp
rr.set_time("time", duration=float(frame_timestamp))
# load the lowres image and depth
bgr = cv2.imread(f"{lowres_image_dir}/{video_id}_{frame_timestamp}.png")
depth = cv2.imread(f"{lowres_depth_dir}/{video_id}_{frame_timestamp}.png", cv2.IMREAD_ANYDEPTH)
high_res_exists: bool = (image_dir / f"{video_id}_{frame_timestamp}.png").exists() and include_highres
# Log the camera transforms:
if frame_timestamp in camera_from_world_dict:
lowres_intri_path = lowres_intrinsics_dir / f"{video_id}_{frame_timestamp}.pincam"
log_camera(
lowres_intri_path,
frame_timestamp,
camera_from_world_dict,
LOWRES_POSED_ENTITY_PATH,
)
rr.log(f"{LOWRES_POSED_ENTITY_PATH}/bgr", rr.Image(bgr, color_model="BGR").compress(jpeg_quality=95))
rr.log(f"{LOWRES_POSED_ENTITY_PATH}/depth", rr.DepthImage(depth, meter=1000))
# log the high res camera
if high_res_exists:
rr.set_time("time high resolution", duration=float(frame_timestamp))
# only low res camera has a trajectory, high res does not so need to find the closest low res frame id
closest_lowres_frame_id = find_closest_frame_id(frame_timestamp, camera_from_world_dict)
highres_intri_path = intrinsics_dir / f"{video_id}_{frame_timestamp}.pincam"
log_camera(
highres_intri_path,
closest_lowres_frame_id,
camera_from_world_dict,
HIGHRES_ENTITY_PATH,
)
# load the highres image and depth if they exist
highres_bgr = cv2.imread(f"{image_dir}/{video_id}_{frame_timestamp}.png")
highres_depth = cv2.imread(f"{depth_dir}/{video_id}_{frame_timestamp}.png", cv2.IMREAD_ANYDEPTH)
rr.log(f"{HIGHRES_ENTITY_PATH}/bgr", rr.Image(highres_bgr, color_model="BGR").compress(jpeg_quality=75))
rr.log(f"{HIGHRES_ENTITY_PATH}/depth", rr.DepthImage(highres_depth, meter=1000))
def main() -> None:
parser = argparse.ArgumentParser(description="Visualizes the ARKitScenes dataset using the Rerun SDK.")
parser.add_argument(
"--video-id",
type=str,
choices=AVAILABLE_RECORDINGS,
default=AVAILABLE_RECORDINGS[0],
help="Video ID of the ARKitScenes Dataset",
)
parser.add_argument(
"--include-highres",
action="store_true",
help="Include the high resolution camera and depth images",
)
rr.script_add_args(parser)
args = parser.parse_args()
primary_camera_entity = HIGHRES_ENTITY_PATH if args.include_highres else LOWRES_POSED_ENTITY_PATH
blueprint = rrb.Horizontal(
rrb.Spatial3DView(name="3D"),
rrb.Vertical(
rrb.Tabs(
# Note that we re-project the annotations into the 2D views:
# For this to work, the origin of the 2D views has to be a pinhole camera,
# this way the viewer knows how to project the 3D annotations into the 2D views.
rrb.Spatial2DView(
name="BGR",
origin=primary_camera_entity,
contents=["$origin/bgr", "/world/annotations/**"],
),
rrb.Spatial2DView(
name="Depth",
origin=primary_camera_entity,
contents=["$origin/depth", "/world/annotations/**"],
),
name="2D",
),
rrb.TextDocumentView(name="Readme"),
row_shares=[2, 1],
),
)
rr.script_setup(args, "rerun_example_arkit_scenes", default_blueprint=blueprint)
recording_path = ensure_recording_available(args.video_id, args.include_highres)
log_arkit(recording_path, args.include_highres)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/arkit_scenes/arkit_scenes/download_dataset.py
+324
View File
@@ -0,0 +1,324 @@
# Copied from https://github.com/apple/ARKitScenes/blob/9ec0b99c3cd55e29fc0724e1229e2e6c2909ab45/download_data.py
# Licensing information: https://github.com/apple/ARKitScenes/blob/9ec0b99c3cd55e29fc0724e1229e2e6c2909ab45/LICENSE
from __future__ import annotations
import math
import os
import subprocess
import zipfile
from pathlib import Path
from typing import Final
import pandas as pd
ARkitscense_url = "https://docs-assets.developer.apple.com/ml-research/datasets/arkitscenes/v1"
TRAINING: Final = "Training"
VALIDATION: Final = "Validation"
HIGRES_DEPTH_ASSET_NAME: Final = "highres_depth"
POINT_CLOUDS_FOLDER: Final = "laser_scanner_point_clouds"
AVAILABLE_RECORDINGS: Final = ["48458663", "42444949", "41069046", "41125722", "41125763", "42446167"]
DATASET_DIR: Final = Path(__file__).parent.parent / "dataset"
default_raw_dataset_assets = [
"mov",
"annotation",
"mesh",
"confidence",
"highres_depth",
"lowres_depth",
"lowres_wide.traj",
"lowres_wide",
"lowres_wide_intrinsics",
"ultrawide",
"ultrawide_intrinsics",
"vga_wide",
"vga_wide_intrinsics",
]
missing_3dod_assets_video_ids = [
"47334522",
"47334523",
"42897421",
"45261582",
"47333152",
"47333155",
"48458535",
"48018733",
"47429677",
"48458541",
"42897848",
"47895482",
"47333960",
"47430089",
"42899148",
"42897612",
"42899153",
"42446164",
"48018149",
"47332198",
"47334515",
"45663223",
"45663226",
"45663227",
]
def raw_files(video_id: str, assets: list[str], metadata: pd.DataFrame) -> list[str]:
file_names = []
for asset in assets:
if HIGRES_DEPTH_ASSET_NAME == asset:
in_upsampling = metadata.loc[metadata["video_id"] == float(video_id), ["is_in_upsampling"]].iat[0, 0]
if not in_upsampling:
print(f"Skipping asset {asset} for video_id {video_id} - Video not in upsampling dataset")
continue # highres_depth asset only available for video ids from upsampling dataset
if asset in [
"confidence",
"highres_depth",
"lowres_depth",
"lowres_wide",
"lowres_wide_intrinsics",
"ultrawide",
"ultrawide_intrinsics",
"wide",
"wide_intrinsics",
"vga_wide",
"vga_wide_intrinsics",
]:
file_names.append(asset + ".zip")
elif asset == "mov":
file_names.append(f"{video_id}.mov")
elif asset == "mesh":
if video_id not in missing_3dod_assets_video_ids:
file_names.append(f"{video_id}_3dod_mesh.ply")
elif asset == "annotation":
if video_id not in missing_3dod_assets_video_ids:
file_names.append(f"{video_id}_3dod_annotation.json")
elif asset == "lowres_wide.traj":
if video_id not in missing_3dod_assets_video_ids:
file_names.append("lowres_wide.traj")
else:
raise Exception(f"No asset = {asset} in raw dataset")
return file_names
def download_file(url: str, file_name: str, dst: Path) -> bool:
os.makedirs(dst, exist_ok=True)
filepath = os.path.join(dst, file_name)
if not os.path.isfile(filepath):
command = f"curl {url} -o {file_name}.tmp --fail"
print(f"Downloading file {filepath}")
try:
subprocess.check_call(command, shell=True, cwd=dst)
except Exception as error:
error_msg = f"Error downloading {url}, error: {error}"
print(error_msg)
raise RuntimeError(error_msg) from error
os.rename(filepath + ".tmp", filepath)
else:
pass # skipping download of existing file
return True
def unzip_file(file_name: str, dst: Path, keep_zip: bool = True) -> bool:
filepath = os.path.join(dst, file_name)
print(f"Unzipping zip file {filepath}")
try:
with zipfile.ZipFile(filepath, "r") as zip_ref:
zip_ref.extractall(dst)
except Exception as error:
error_msg = f"Error unzipping {filepath}, error: {error}"
print(error_msg)
raise RuntimeError(error_msg) from error
if not keep_zip:
os.remove(filepath)
return True
def download_laser_scanner_point_clouds_for_video(video_id: str, metadata: pd.DataFrame, download_dir: Path) -> None:
video_metadata = metadata.loc[metadata["video_id"] == float(video_id)]
visit_id: float = video_metadata["visit_id"].iat[0] # type: ignore[assignment]
has_laser_scanner_point_clouds = video_metadata["has_laser_scanner_point_clouds"].iat[0]
if not has_laser_scanner_point_clouds:
print(f"Warning: Laser scanner point clouds for video {video_id} are not available")
return
if math.isnan(visit_id) or not visit_id.is_integer():
print(f"Warning: Downloading laser scanner point clouds for video {video_id} failed - Bad visit id {visit_id}")
return
visit_id = int(visit_id)
laser_scanner_point_clouds_ids = laser_scanner_point_clouds_for_visit_id(visit_id, download_dir)
for point_cloud_id in laser_scanner_point_clouds_ids:
download_laser_scanner_point_clouds(point_cloud_id, visit_id, download_dir)
def laser_scanner_point_clouds_for_visit_id(visit_id: int, download_dir: Path) -> list[str]:
point_cloud_to_visit_id_mapping_filename = "laser_scanner_point_clouds_mapping.csv"
if not os.path.exists(point_cloud_to_visit_id_mapping_filename):
point_cloud_to_visit_id_mapping_url = (
f"{ARkitscense_url}/raw/laser_scanner_point_clouds/{point_cloud_to_visit_id_mapping_filename}"
)
download_file(
point_cloud_to_visit_id_mapping_url,
point_cloud_to_visit_id_mapping_filename,
download_dir,
)
point_cloud_to_visit_id_mapping_filepath = os.path.join(download_dir, point_cloud_to_visit_id_mapping_filename)
point_cloud_to_visit_id_mapping = pd.read_csv(point_cloud_to_visit_id_mapping_filepath)
point_cloud_ids = point_cloud_to_visit_id_mapping.loc[
point_cloud_to_visit_id_mapping["visit_id"] == visit_id,
["laser_scanner_point_clouds_id"],
]
point_cloud_ids_list = [scan_id[0] for scan_id in point_cloud_ids.values]
return point_cloud_ids_list
def download_laser_scanner_point_clouds(laser_scanner_point_cloud_id: str, visit_id: int, download_dir: Path) -> None:
laser_scanner_point_clouds_folder_path = download_dir / POINT_CLOUDS_FOLDER / str(visit_id)
os.makedirs(laser_scanner_point_clouds_folder_path, exist_ok=True)
for extension in [".ply", "_pose.txt"]:
filename = f"{laser_scanner_point_cloud_id}{extension}"
filepath = os.path.join(laser_scanner_point_clouds_folder_path, filename)
if os.path.exists(filepath):
return
file_url = f"{ARkitscense_url}/raw/laser_scanner_point_clouds/{visit_id}/{filename}"
download_file(file_url, filename, laser_scanner_point_clouds_folder_path)
def get_metadata(dataset: str, download_dir: Path) -> pd.DataFrame:
filename = "metadata.csv"
url = f"{ARkitscense_url}/threedod/{filename}" if "3dod" == dataset else f"{ARkitscense_url}/{dataset}/{filename}"
dst_folder = download_dir / dataset
dst_file = dst_folder / filename
download_file(url, filename, dst_folder)
metadata = pd.read_csv(dst_file)
return metadata
def download_data(
dataset: str,
video_ids: list[str],
dataset_splits: list[str],
download_dir: Path,
keep_zip: bool,
raw_dataset_assets: list[str] | None = None,
should_download_laser_scanner_point_cloud: bool = False,
) -> None:
"""
Downloads data from the specified dataset and video IDs to the given download directory.
Args:
----
dataset: the name of the dataset to download from (raw, 3dod, or upsampling)
video_ids: the list of video IDs to download data for
dataset_splits: the list of splits for each video ID (train, validation, or test)
download_dir: the directory to download data to
keep_zip: whether to keep the downloaded zip files after extracting them
raw_dataset_assets: a list of asset types to download from the raw dataset, if dataset is "raw"
should_download_laser_scanner_point_cloud: whether to download the laser scanner point cloud data, if available
Returns: None
"""
metadata = get_metadata(dataset, download_dir)
for video_id in sorted(set(video_ids)):
split = dataset_splits[video_ids.index(video_id)]
dst_dir = download_dir / dataset / split
if dataset == "raw":
url_prefix = ""
file_names = []
if not raw_dataset_assets:
print(f"Warning: No raw assets given for video id {video_id}")
else:
dst_dir = dst_dir / str(video_id)
url_prefix = f"{ARkitscense_url}/raw/{split}/{video_id}" + "/{}"
file_names = raw_files(video_id, raw_dataset_assets, metadata)
elif dataset == "3dod":
url_prefix = f"{ARkitscense_url}/threedod/{split}" + "/{}"
file_names = [
f"{video_id}.zip",
]
elif dataset == "upsampling":
url_prefix = f"{ARkitscense_url}/upsampling/{split}" + "/{}"
file_names = [
f"{video_id}.zip",
]
else:
raise Exception(f"No such dataset = {dataset}")
if should_download_laser_scanner_point_cloud and dataset == "raw":
# Point clouds only available for the raw dataset
download_laser_scanner_point_clouds_for_video(video_id, metadata, download_dir)
for file_name in file_names:
dst_path = os.path.join(dst_dir, file_name)
url = url_prefix.format(file_name)
if not file_name.endswith(".zip") or not os.path.isdir(dst_path[: -len(".zip")]):
download_file(url, dst_path, dst_dir)
else:
pass # skipping download of existing zip file
if file_name.endswith(".zip") and os.path.isfile(dst_path):
unzip_file(file_name, dst_dir, keep_zip)
def ensure_recording_downloaded(video_id: str, include_highres: bool) -> Path:
"""Only downloads from validation set."""
data_path = DATASET_DIR / "raw" / "Validation" / video_id
assets_to_download = [
"lowres_wide",
"lowres_depth",
"lowres_wide_intrinsics",
"lowres_wide.traj",
"annotation",
"mesh",
]
if include_highres:
assets_to_download.extend(["highres_depth", "wide", "wide_intrinsics"])
download_data(
dataset="raw",
video_ids=[video_id],
dataset_splits=[VALIDATION],
download_dir=DATASET_DIR,
keep_zip=False,
raw_dataset_assets=assets_to_download,
should_download_laser_scanner_point_cloud=False,
)
return data_path
def ensure_recording_available(video_id: str, include_highres: bool) -> Path:
"""
Returns the path to the recording for a given video_id.
Args:
----
video_id (str):
Identifier for the recording.
include_highres (bool):
Whether to include the high resolution recording.
Returns
-------
Path: Path object representing the path to the recording.
Raises
------
AssertionError:
If the recording path does not exist.
"""
recording_path = ensure_recording_downloaded(video_id, include_highres)
assert recording_path.exists(), f"Recording path {recording_path} does not exist."
return recording_path # Return the path to the recording
examples/python/arkit_scenes/pyproject.toml
+22
View File
@@ -0,0 +1,22 @@
[project]
name = "arkit_scenes"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = [
"matplotlib",
"numpy",
"opencv-python",
"pandas",
"rerun-sdk",
"scipy",
"tqdm",
"trimesh",
]
[project.scripts]
arkit_scenes = "arkit_scenes.__main__:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"