Skip to content

Commit

Permalink
implement caching support for Boxes3D
Browse files Browse the repository at this point in the history
  • Loading branch information
@teh-cmc
teh-cmc committed Feb 1, 2024
1 parent 13e56bf commit 0aac1e7
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 55 deletions.
152 changes: 108 additions & 44 deletions crates/re_space_view_spatial/src/visualizers/boxes3d.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
use re_entity_db::EntityPath;
use re_query::{ArchetypeView, QueryError};
use re_types::{
archetypes::Boxes3D,
components::{HalfSizes3D, Position3D, Rotation3D},
components::{
ClassId, Color, HalfSizes3D, InstanceKey, KeypointId, Position3D, Radius, Rotation3D, Text,
},
};
use re_viewer_context::{
ApplicableEntities, IdentifiedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection,
Expand All @@ -17,9 +18,9 @@ use crate::{
};

use super::{
entity_iterator::process_archetype_views, filter_visualizable_3d_entities,
picking_id_from_instance_key, process_annotations, process_colors, process_labels,
process_radii, SpatialViewVisualizerData,
filter_visualizable_3d_entities, picking_id_from_instance_key,
process_annotation_and_keypoint_slices, process_color_slice, process_label_slice,
process_radius_slice, SpatialViewVisualizerData,
};

pub struct Boxes3DVisualizer(SpatialViewVisualizerData);
Expand All @@ -33,30 +34,44 @@ impl Default for Boxes3DVisualizer {
}

impl Boxes3DVisualizer {
fn process_arch_view(
fn process_data(
&mut self,
query: &ViewQuery<'_>,
arch_view: &ArchetypeView<Boxes3D>,
data: &Boxes3DComponentData<'_>,
ent_path: &EntityPath,
ent_context: &SpatialSceneEntityContext<'_>,
) -> Result<(), QueryError> {
let annotation_infos = process_annotations::<HalfSizes3D, Boxes3D>(
query,
arch_view,
) {
let (annotation_infos, _) = process_annotation_and_keypoint_slices(
query.latest_at,
data.instance_keys,
data.keypoint_ids,
data.class_ids,
data.half_sizes.iter().map(|_| glam::Vec3::ZERO),
&ent_context.annotations,
)?;

let instance_keys = arch_view.iter_instance_keys();
let half_sizes = arch_view.iter_required_component::<HalfSizes3D>()?;
let positions = arch_view
.iter_optional_component::<Position3D>()?
.map(|position| position.unwrap_or(Position3D::ZERO));
let rotation = arch_view
.iter_optional_component::<Rotation3D>()?
.map(|position| position.unwrap_or(Rotation3D::IDENTITY));
let radii = process_radii(arch_view, ent_path)?;
let colors = process_colors(arch_view, ent_path, &annotation_infos)?;
let labels = process_labels(arch_view, &annotation_infos)?;
);

let centers = || {
data.centers
.as_ref()
.map_or(
itertools::Either::Left(std::iter::repeat(&None).take(data.half_sizes.len())),
|data| itertools::Either::Right(data.iter()),
)
.map(|center| center.unwrap_or(Position3D::ZERO))
};
let rotations = || {
data.rotations
.as_ref()
.map_or(
itertools::Either::Left(std::iter::repeat(&None).take(data.half_sizes.len())),
|data| itertools::Either::Right(data.iter()),
)
.map(|center| center.clone().unwrap_or(Rotation3D::IDENTITY))
};

let radii = process_radius_slice(data.radii, data.half_sizes.len(), ent_path);
let colors = process_color_slice(data.colors, ent_path, &annotation_infos);
let labels = process_label_slice(data.labels, data.half_sizes.len(), &annotation_infos);

let mut line_builder = ent_context.shared_render_builders.lines();
let mut line_batch = line_builder
Expand All @@ -68,31 +83,31 @@ impl Boxes3DVisualizer {

let mut bounding_box = macaw::BoundingBox::nothing();

for (instance_key, half_extent, position, rotation, radius, color, label) in itertools::izip!(
instance_keys,
half_sizes,
positions,
rotation,
for (instance_key, half_size, center, rotation, radius, color, label) in itertools::izip!(
data.instance_keys,
data.half_sizes,
centers(),
rotations(),
radii,
colors,
labels
labels,
) {
let instance_hash = re_entity_db::InstancePathHash::instance(ent_path, instance_key);
let instance_hash = re_entity_db::InstancePathHash::instance(ent_path, *instance_key);

bounding_box.extend(half_extent.box_min(position));
bounding_box.extend(half_extent.box_max(position));
bounding_box.extend(half_size.box_min(center));
bounding_box.extend(half_size.box_max(center));

let position = position.into();
let center = center.into();

let box3d = line_batch
.add_box_outline_from_transform(glam::Affine3A::from_scale_rotation_translation(
glam::Vec3::from(half_extent) * 2.0,
glam::Vec3::from(*half_size) * 2.0,
rotation.into(),
position,
center,
))
.color(color)
.radius(radius)
.picking_instance_id(picking_id_from_instance_key(instance_key));
.picking_instance_id(picking_id_from_instance_key(*instance_key));
if let Some(outline_mask_ids) = ent_context
.highlight
.instances
Expand All @@ -106,7 +121,7 @@ impl Boxes3DVisualizer {
text,
color,
target: UiLabelTarget::Position3D(
ent_context.world_from_entity.transform_point3(position),
ent_context.world_from_entity.transform_point3(center),
),
labeled_instance: instance_hash,
});
Expand All @@ -115,11 +130,23 @@ impl Boxes3DVisualizer {

self.0
.add_bounding_box(ent_path.hash(), bounding_box, ent_context.world_from_entity);

Ok(())
}
}

// ---

struct Boxes3DComponentData<'a> {
pub instance_keys: &'a [InstanceKey],
pub half_sizes: &'a [HalfSizes3D],
pub centers: Option<&'a [Option<Position3D>]>,
pub rotations: Option<&'a [Option<Rotation3D>]>,
pub colors: Option<&'a [Option<Color>]>,
pub radii: Option<&'a [Option<Radius>]>,
pub labels: Option<&'a [Option<Text>]>,
pub keypoint_ids: Option<&'a [Option<KeypointId>]>,
pub class_ids: Option<&'a [Option<ClassId>]>,
}

impl IdentifiedViewSystem for Boxes3DVisualizer {
fn identifier() -> re_viewer_context::ViewSystemIdentifier {
"Boxes3D".into()
Expand All @@ -146,13 +173,50 @@ impl VisualizerSystem for Boxes3DVisualizer {
query: &ViewQuery<'_>,
view_ctx: &ViewContextCollection,
) -> Result<Vec<re_renderer::QueueableDrawData>, SpaceViewSystemExecutionError> {
process_archetype_views::<Boxes3DVisualizer, Boxes3D, { Boxes3D::NUM_COMPONENTS }, _>(
super::entity_iterator::process_archetype_pov1_comp7::<
Boxes3DVisualizer,
Boxes3D,
HalfSizes3D,
Position3D,
Rotation3D,
Color,
Radius,
Text,
re_types::components::KeypointId,
re_types::components::ClassId,
_,
>(
ctx,
query,
view_ctx,
view_ctx.get::<EntityDepthOffsets>()?.box2d,
|_ctx, ent_path, _ent_props, arch_view, ent_context| {
self.process_arch_view(query, &arch_view, ent_path, ent_context)
view_ctx.get::<EntityDepthOffsets>()?.points,
|_ctx,
ent_path,
_ent_props,
ent_context,
(_time, _row_id),
instance_keys,
half_sizes,
centers,
rotations,
colors,
radii,
labels,
keypoint_ids,
class_ids| {
let data = Boxes3DComponentData {
instance_keys,
half_sizes,
centers,
rotations,
colors,
radii,
labels,
keypoint_ids,
class_ids,
};
self.process_data(query, &data, ent_path, ent_context);
Ok(())
},
)?;

Expand Down
25 changes: 14 additions & 11 deletions crates/re_space_view_spatial/src/visualizers/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pub use points3d::{LoadedPoints, Points3DComponentData};
use ahash::HashMap;

use re_entity_db::{EntityPath, InstancePathHash};
use re_types::components::{Color, InstanceKey, Text};
use re_types::components::{Color, InstanceKey};
use re_types::datatypes::{KeypointId, KeypointPair};
use re_types::Archetype;
use re_viewer_context::{
Expand Down Expand Up @@ -173,18 +173,21 @@ pub fn process_color_slice<'a>(
}

/// Process [`Text`] components using annotations.
#[allow(dead_code)]
pub fn process_labels<'a, A: Archetype>(
arch_view: &'a re_query::ArchetypeView<A>,
annotation_infos: &'a ResolvedAnnotationInfos,
) -> Result<impl Iterator<Item = Option<String>> + 'a, re_query::QueryError> {
pub fn process_label_slice(
labels: Option<&[Option<re_types::components::Text>]>,
default_len: usize,
annotation_infos: &ResolvedAnnotationInfos,
) -> Vec<Option<String>> {
re_tracing::profile_function!();

Ok(itertools::izip!(
annotation_infos.iter(),
arch_view.iter_optional_component::<Text>()?,
)
.map(move |(annotation_info, text)| annotation_info.label(text.as_ref().map(|t| t.as_str()))))
match labels {
None => vec![None; default_len],
Some(labels) => itertools::izip!(annotation_infos.iter(), labels)
.map(move |(annotation_info, text)| {
annotation_info.label(text.as_ref().map(|t| t.as_str()))
})
.collect(),
}
}

/// Process [`re_types::components::Radius`] components to [`re_renderer::Size`] using auto size
Expand Down

0 comments on commit 0aac1e7

Please sign in to comment.