wehub-resource-sync
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import rerun as rr
|
|
import rerun.blueprint as rrb
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Simple benchmark for many transforms over time & space.")
|
|
rr.script_add_args(parser)
|
|
|
|
parser.add_argument("--branching-factor", type=int, default=2, help="How many children each node has")
|
|
parser.add_argument("--hierarchy-depth", type=int, default=10, help="How many levels of hierarchy we want")
|
|
parser.add_argument(
|
|
"--transforms-every-n-levels", type=int, default=2, help="At which level in the hierarchies we add transforms"
|
|
)
|
|
parser.add_argument(
|
|
"--num-timestamps",
|
|
type=int,
|
|
default=100,
|
|
help="Number of timestamps to log. Stamps shift for each entity a bit.",
|
|
)
|
|
parser.add_argument("--transforms-only", action="store_true", help="If set, don't log a point at each leaf")
|
|
parser.add_argument(
|
|
"--num-views", type=int, default=6, help="Number of 3D views to create (each will use a different origin)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
rr.script_setup(args, "rerun_example_benchmark_many_transforms")
|
|
rr.set_time("sim_time", duration=0)
|
|
|
|
entity_paths = []
|
|
call_id = 0
|
|
|
|
def log_hierarchy(entity_path: str, level: int) -> None:
|
|
nonlocal call_id
|
|
call_id += 1
|
|
entity_paths.append(entity_path)
|
|
|
|
# Add a transform at every 'transforms_every_n_levels' level except root
|
|
if level > 0 and level % args.transforms_every_n_levels == 0:
|
|
# Add a static transform that has to be combined in to stress the per-timestamp transform resolve.
|
|
rr.log(
|
|
entity_path,
|
|
# Have to be careful to not override all other transforms, therefore, use `from_fields`.
|
|
rr.Transform3D.from_fields(
|
|
mat3x3=[
|
|
[1.0 + level * 0.1, 0.0, 0.0],
|
|
[0.0, 1.0 + level * 0.1, 0.0],
|
|
[0.0, 0.0, 1.0 + level * 0.1],
|
|
]
|
|
),
|
|
static=True,
|
|
)
|
|
|
|
# Add a transform that changes for each timestamp.
|
|
for i in range(args.num_timestamps):
|
|
call_id_factor = call_id * 0.02
|
|
rr.set_time("sim_time", duration=i + call_id_factor)
|
|
rr.log(
|
|
entity_path,
|
|
rr.Transform3D(
|
|
translation=[i * 0.1 * level + call_id_factor, call_id_factor * level, 0.0],
|
|
rotation_axis_angle=rr.RotationAxisAngle(axis=(0.0, 1.0, 0.0), degrees=i * 0.1),
|
|
),
|
|
)
|
|
|
|
if level == args.hierarchy_depth:
|
|
if not args.transforms_only:
|
|
# Log a single point at the leaf
|
|
rr.set_time("sim_time", duration=0)
|
|
rr.log(entity_path, rr.Points3D([[0.0, 0.0, 0.0]]))
|
|
return
|
|
|
|
for i in range(args.branching_factor):
|
|
child_path = f"{entity_path}/{i}_at_{level}"
|
|
log_hierarchy(child_path, level + 1)
|
|
|
|
log_hierarchy("root", 0)
|
|
|
|
# All views display all entities.
|
|
rr.send_blueprint(
|
|
rrb.Blueprint(
|
|
rrb.Grid(
|
|
contents=[rrb.Spatial3DView(origin=path, contents="/**") for path in entity_paths[: args.num_views]]
|
|
),
|
|
collapse_panels=True, # Collapse panels, so perf is mostly about the data & the views.
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|