Merge pull request #92 from Carifio24/usd-voxel-speed

Improve speed of USD voxel export

glue_ar/common/scatter_usd.py

@@ -14,6 +14,7 @@
 from glue_ar.common.usd_builder import USDBuilder
 from glue_ar.common.shapes import cone_triangles, cone_points, cylinder_points, cylinder_triangles, \
                                   normalize, rectangular_prism_triangulation, sphere_triangles
+from glue_ar.usd_utils import sanitize_path
 from glue_ar.utils import Viewer3DState, export_label_for_layer, iterable_has_nan, hex_to_components, \
                           layer_color, offset_triangles, xyz_for_layer, Bounds, NoneType
 
@@ -94,7 +95,7 @@ def add_scatter_layer_usd(
     color_mode_attr = "color_mode" if vispy_layer_state else "cmap_mode"
     fixed_color = getattr(layer_state, color_mode_attr, "Fixed") == "Fixed"
 
-    identifier = export_label_for_layer(layer_state).replace(" ", "_")
+    identifier = sanitize_path(export_label_for_layer(layer_state))
 
     mask = scatter_layer_mask(viewer_state, layer_state, bounds, clip_to_bounds)
     data = xyz_for_layer(viewer_state, layer_state,

glue_ar/common/volume_export_options.py

@@ -12,4 +12,4 @@ class ARIsosurfaceExportOptions(State):
 
 class ARVoxelExportOptions(State):
     opacity_cutoff = RangedCallbackProperty(default=0.1, min_value=0.01, max_value=1, resolution=0.01)
-    opacity_resolution = RangedCallbackProperty(default=0.02, min_value=0.01, max_value=1, resolution=0.01)
+    opacity_resolution = RangedCallbackProperty(default=0.02, min_value=0.005, max_value=1, resolution=0.005)

glue_ar/common/voxels.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from glue_vispy_viewers.volume.viewer_state import Vispy3DVolumeViewerState
 from numpy import isfinite, argwhere, transpose
 from typing import Iterable, Optional
@@ -9,10 +10,10 @@
 from glue_ar.common.stl_builder import STLBuilder
 from glue_ar.common.usd_builder import USDBuilder
 from glue_ar.common.volume_export_options import ARVoxelExportOptions
-from glue_ar.usd_utils import material_for_color
-from glue_ar.utils import BoundsWithResolution, alpha_composite, binned_opacity, clamp, clamped_opacity, \
+from glue_ar.usd_utils import material_for_color, sanitize_path
+from glue_ar.utils import BoundsWithResolution, alpha_composite, binned_opacity, clamp, \
                           clip_sides, frb_for_layer, hex_to_components, isomin_for_layer, \
-                          isomax_for_layer, layer_color, unique_id, xyz_bounds
+                          isomax_for_layer, layer_color, offset_triangles, unique_id, xyz_bounds
 
 from glue_ar.gltf_utils import add_points_to_bytearray, add_triangles_to_bytearray, \
                                index_mins, index_maxes
@@ -83,7 +84,8 @@ def add_voxel_layers_gltf(builder: GLTFBuilder,
 
         for indices in nonempty_indices:
             value = data[tuple(indices)]
-            adjusted_opacity = clamped_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange)
+            adjusted_opacity = binned_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange,
+                                              opacity_resolution)
             indices_tpl = tuple(indices)
             if indices_tpl in occupied_voxels:
                 current_color = occupied_voxels[indices_tpl]
@@ -126,7 +128,7 @@ def add_voxel_layers_gltf(builder: GLTFBuilder,
             mins=pt_mins,
             maxes=pt_maxes,
         )
-        rgba_tpl = tuple(rgba[:3] + [binned_opacity(rgba[3], opacity_resolution)])
+        rgba_tpl = tuple(rgba)
         if rgba_tpl in materials_map:
             material_index = materials_map[rgba_tpl]
         else:
@@ -162,8 +164,11 @@ def add_voxel_layers_usd(builder: USDBuilder,
 
     triangles = rectangular_prism_triangulation()
 
+    identifier = sanitize_path(f"voxels_{unique_id()}")
+
     opacity_factor = 1
     occupied_voxels = {}
+    colors_map = defaultdict(set)
 
     for layer_state, option in zip(layer_states, options):
         opacity_cutoff = clamp(option.opacity_cutoff, 0, 1)
@@ -185,37 +190,50 @@ def add_voxel_layers_usd(builder: USDBuilder,
 
         for indices in nonempty_indices:
             value = data[tuple(indices)]
-            adjusted_opacity = clamped_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange)
+            adjusted_opacity = binned_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange,
+                                              opacity_resolution)
             indices_tpl = tuple(indices)
             if indices_tpl in occupied_voxels:
                 current_color = occupied_voxels[indices_tpl]
                 adjusted_a_color = color_components[:3] + [adjusted_opacity]
                 new_color = alpha_composite(adjusted_a_color, current_color)
                 occupied_voxels[indices_tpl] = new_color
+                colors_map[current_color].remove(indices_tpl)
+                colors_map[new_color].add(indices_tpl)
             elif adjusted_opacity >= opacity_cutoff:
-                occupied_voxels[indices_tpl] = color_components[:3] + [adjusted_opacity]
+                color = color_components[:3] + [adjusted_opacity]
+                occupied_voxels[indices_tpl] = color
+                colors_map[tuple(color)].add(indices_tpl)
+
     materials_map = {}
-    mesh = None
-    first_point = None
-    for indices, rgba in occupied_voxels.items():
+
+    for rgba, indices_set in colors_map.items():
         if rgba[-1] < opacity_cutoff:
             continue
 
-        center = tuple((index + 0.5) * side for index, side in zip(indices, sides))
-        rgba_tpl = tuple(rgba[:3] + [binned_opacity(rgba[3], opacity_resolution)])
-        if rgba_tpl in materials_map:
-            material = materials_map[rgba_tpl]
+        if rgba in materials_map:
+            material = materials_map[rgba]
         else:
             material = material_for_color(builder.stage, rgba[:3], rgba[3])
-            materials_map[rgba_tpl] = material
-
-        points = rectangular_prism_points(center, sides)
-        if mesh is None:
-            mesh = builder.add_mesh(points, triangles, rgba[:3], rgba[3])
-            first_point = center
-        else:
-            translation = tuple(p - fp for p, fp in zip(center, first_point))
-            builder.add_translated_reference(mesh, translation, material)
+            materials_map[rgba] = material
+        points = []
+        tris = []
+        triangle_offset = 0
+        for indices in indices_set:
+            center = tuple((index + 0.5) * side for index, side in zip(indices, sides))
+            pts = rectangular_prism_points(center, sides)
+            points.append(pts)
+            pt_triangles = offset_triangles(triangles, triangle_offset)
+            triangle_offset += len(pts)
+            tris.append(pt_triangles)
+
+        mesh_points = [pt for pts in points for pt in pts]
+        mesh_triangles = [tri for sphere in tris for tri in sphere]
+        builder.add_mesh(mesh_points,
+                         mesh_triangles,
+                         color=rgba[:3],
+                         opacity=rgba[-1],
+                         identifier=identifier)
 
     return builder
 
@@ -238,6 +256,7 @@ def add_voxel_layers_stl(builder: STLBuilder,
 
     for layer_state, option in zip(layer_states, options):
         opacity_cutoff = clamp(option.opacity_cutoff, 0, 1)
+        opacity_resolution = clamp(option.opacity_resolution, 0, 1)
         data = frb_for_layer(viewer_state, layer_state, bounds)
 
         isomin = isomin_for_layer(viewer_state, layer_state)
@@ -255,7 +274,8 @@ def add_voxel_layers_stl(builder: STLBuilder,
 
         for indices in nonempty_indices:
             value = data[tuple(indices)]
-            adjusted_opacity = clamped_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange)
+            adjusted_opacity = binned_opacity(layer_state.alpha * opacity_factor * (value - isomin) / isorange,
+                                              opacity_resolution)
             indices_tpl = tuple(indices)
             if indices_tpl in occupied_voxels:
                 current_color = occupied_voxels[indices_tpl]

glue_ar/usd_utils.py

@@ -39,4 +39,4 @@ def material_for_mesh(mesh: UsdGeom.Mesh) -> UsdShade.Material:
 
 
 def sanitize_path(path: str) -> str:
-    return sub(r"[-\(\)\[\]]+", "_", path)
+    return sub(r"[-\(\)\[\] ]+", "_", path)