Skip to content

Commit 59dccf1

Browse files
committed
Add blueprint to Objectron example
- Fixes #3412
1 parent d2ead0c commit 59dccf1

File tree

1 file changed

+10
-60
lines changed

1 file changed

+10
-60
lines changed

examples/python/objectron/main.py

Lines changed: 10 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,16 @@
1717
from typing import Iterable, Iterator
1818

1919
import numpy as np
20-
import numpy.typing as npt
2120
import rerun as rr # pip install rerun-sdk
21+
import rerun.blueprint as rrb
2222
from download_dataset import (
2323
ANNOTATIONS_FILENAME,
2424
AVAILABLE_RECORDINGS,
2525
GEOMETRY_FILENAME,
26-
IMAGE_RESOLUTION,
2726
LOCAL_DATASET_DIR,
2827
ensure_recording_available,
2928
)
30-
from proto.objectron.proto import ARCamera, ARFrame, ARPointCloud, FrameAnnotation, Object, ObjectType, Sequence
29+
from proto.objectron.proto import ARCamera, ARFrame, ARPointCloud, Object, ObjectType, Sequence
3130
from scipy.spatial.transform import Rotation as R
3231

3332

@@ -120,8 +119,6 @@ def log_ar_frames(samples: Iterable[SampleARFrame], seq: Sequence) -> None:
120119
log_camera(sample.frame.camera)
121120
log_point_cloud(sample.frame.raw_feature_points)
122121

123-
log_frame_annotations(frame_times, seq.frame_annotations)
124-
125122

126123
def log_camera(cam: ARCamera) -> None:
127124
"""Logs a camera from an `ARFrame` using the Rerun SDK."""
@@ -186,60 +183,6 @@ def log_annotated_bboxes(bboxes: Iterable[Object]) -> None:
186183
)
187184

188185

189-
def log_frame_annotations(frame_times: list[float], frame_annotations: list[FrameAnnotation]) -> None:
190-
"""Maps annotations to their associated `ARFrame` then logs them using the Rerun SDK."""
191-
192-
for frame_ann in frame_annotations:
193-
frame_idx = frame_ann.frame_id
194-
if frame_idx >= len(frame_times):
195-
continue
196-
197-
time = frame_times[frame_idx]
198-
rr.set_time_sequence("frame", frame_idx)
199-
rr.set_time_seconds("time", time)
200-
201-
for obj_ann in frame_ann.annotations:
202-
keypoint_ids = [kp.id for kp in obj_ann.keypoints]
203-
keypoint_pos2s = np.asarray([[kp.point_2d.x, kp.point_2d.y] for kp in obj_ann.keypoints], dtype=np.float32)
204-
# NOTE: These are normalized points, so we need to bring them back to image space
205-
keypoint_pos2s *= IMAGE_RESOLUTION
206-
207-
if len(keypoint_pos2s) == 9:
208-
log_projected_bbox(f"world/camera/estimates/box-{obj_ann.object_id}", keypoint_pos2s)
209-
else:
210-
for id, pos2 in zip(keypoint_ids, keypoint_pos2s):
211-
rr.log(
212-
f"world/camera/estimates/box-{obj_ann.object_id}/{id}",
213-
rr.Points2D(pos2, colors=[130, 160, 250, 255]),
214-
)
215-
216-
217-
# TODO(#3412): replace once we can auto project 3D bboxes on 2D views (need blueprints)
218-
def log_projected_bbox(path: str, keypoints: npt.NDArray[np.float32]) -> None:
219-
"""
220-
Projects the 3D bounding box to a 2D plane, using line segments.
221-
222-
The 3D bounding box is described by the keypoints of an `ObjectAnnotation`
223-
"""
224-
# fmt: off
225-
segments = np.array([[keypoints[1], keypoints[2]],
226-
[keypoints[1], keypoints[3]],
227-
[keypoints[4], keypoints[2]],
228-
[keypoints[4], keypoints[3]],
229-
230-
[keypoints[5], keypoints[6]],
231-
[keypoints[5], keypoints[7]],
232-
[keypoints[8], keypoints[6]],
233-
[keypoints[8], keypoints[7]],
234-
235-
[keypoints[1], keypoints[5]],
236-
[keypoints[2], keypoints[6]],
237-
[keypoints[3], keypoints[7]],
238-
[keypoints[4], keypoints[8]]], dtype=np.float32)
239-
# fmt: on
240-
rr.log(path, rr.LineStrips2D(segments, colors=[130, 160, 250, 255]))
241-
242-
243186
def main() -> None:
244187
# Ensure the logging in download_dataset.py gets written to stderr:
245188
logging.getLogger().addHandler(logging.StreamHandler())
@@ -272,7 +215,14 @@ def main() -> None:
272215
rr.script_add_args(parser)
273216
args = parser.parse_args()
274217

275-
rr.script_setup(args, "rerun_example_objectron")
218+
rr.script_setup(
219+
args,
220+
"rerun_example_objectron",
221+
blueprint=rrb.Horizontal(
222+
rrb.Spatial3DView(origin="/world", name="World"),
223+
rrb.Spatial2DView(origin="/world/camera", name="Camera", contents=["+ $origin/**", "+ /world/**"]),
224+
),
225+
)
276226

277227
dir = ensure_recording_available(args.recording, args.dataset_dir, args.force_reprocess_video)
278228

0 commit comments

Comments
 (0)