Add point cloud export

glue_ar/common/__init__.py

@@ -2,6 +2,8 @@
 from .scatter_gltf import add_scatter_layer_gltf  # noqa: F401
 from .scatter_stl import add_scatter_layer_stl  # noqa: F401
 from .scatter_usd import add_scatter_layer_usd  # noqa: F401
+from .points_gltf import add_points_layer_gltf  # noqa: F401
+from .points_usd import add_points_layer_usd  # noqa: F401
 from .voxels import add_voxel_layers_gltf, add_voxel_layers_usd  # noqa: F401
 from .scatter_export_options import ARVispyScatterExportOptions  # noqa: F401
 from .volume_export_options import ARIsosurfaceExportOptions, ARVoxelExportOptions  # noqa: F401

glue_ar/common/points_gltf.py

@@ -0,0 +1,154 @@
+from collections import defaultdict
+from gltflib import AccessorType, BufferTarget, ComponentType, PrimitiveMode
+from glue_vispy_viewers.common.viewer_state import Vispy3DViewerState
+from glue_vispy_viewers.scatter.layer_state import ScatterLayerState
+
+from glue_ar.common.export_options import ar_layer_export
+from glue_ar.common.scatter import Scatter3DLayerState, ScatterLayerState3D, scatter_layer_mask
+from glue_ar.gltf_utils import add_points_to_bytearray, index_mins, index_maxes
+from glue_ar.utils import Bounds, NoneType, Viewer3DState, color_identifier, hex_to_components, layer_color, unique_id, xyz_bounds, xyz_for_layer
+from glue_ar.common.gltf_builder import GLTFBuilder
+from glue_ar.common.scatter_export_options import ARPointExportOptions
+
+try:
+    from glue_jupyter.common.state3d import ViewerState3D
+except ImportError:
+    ViewerState3D = NoneType
+
+
+def add_points_layer_gltf(builder: GLTFBuilder,
+                          viewer_state: Viewer3DState,
+                          layer_state: ScatterLayerState3D,
+                          bounds: Bounds,
+                          clip_to_bounds: bool = True):
+
+    if layer_state is None:
+        return
+
+    bounds = xyz_bounds(viewer_state, with_resolution=False)
+
+    vispy_layer_state = isinstance(layer_state, ScatterLayerState)
+    color_mode_attr = "color_mode" if vispy_layer_state else "cmap_mode"
+    fixed_color = getattr(layer_state, color_mode_attr, "Fixed") == "Fixed"
+
+    mask = scatter_layer_mask(viewer_state, layer_state, bounds, clip_to_bounds)
+    data = xyz_for_layer(viewer_state, layer_state,
+                         preserve_aspect=viewer_state.native_aspect,
+                         mask=mask,
+                         scaled=True)
+    data = data[:, [1, 2, 0]]
+
+
+    uri = f"layer_{unique_id()}.bin"
+
+    if fixed_color:
+        color = layer_color(layer_state)
+        color_components = hex_to_components(color)
+        builder.add_material(color=color_components, opacity=layer_state.alpha)
+
+        barr = bytearray()
+        add_points_to_bytearray(barr, data)
+
+        data_mins = index_mins(data)
+        data_maxes = index_maxes(data)
+
+        builder.add_buffer(byte_length=len(barr), uri=uri)
+        builder.add_buffer_view(
+            buffer = builder.buffer_count-1,
+            byte_length=len(barr),
+            byte_offset=0,
+            target=BufferTarget.ARRAY_BUFFER
+        )
+        builder.add_accessor(
+            buffer_view=builder.buffer_view_count-1,
+            component_type=ComponentType.FLOAT,
+            count=len(data),
+            type=AccessorType.VEC3,
+            mins=data_mins,
+            maxes=data_maxes,
+        )
+        builder.add_mesh(
+            position_accessor=builder.accessor_count-1,
+            material=builder.material_count-1,
+            mode=PrimitiveMode.POINTS
+        )
+        builder.add_file_resource(uri, data=barr)
+    else:
+        # If we don't have fixed colors, the idea is to make a different "mesh" for each different color used
+        # So first we need to run through the points and determine which color they have, and group ones with
+        # the same color together
+        points_by_color = defaultdict(list)
+        cmap = layer_state.cmap
+        cmap_attr = "cmap_attribute" if vispy_layer_state else "cmap_att"
+        cmap_att = getattr(layer_state, cmap_attr)
+        cmap_vals = layer_state.layer[cmap_att][mask]
+        crange = layer_state.cmap_vmax - layer_state.cmap_vmin
+        opacity = layer_state.alpha
+
+        for i, point in enumerate(data):
+            cval = cmap_vals[i]
+            normalized = max(min((cval - layer_state.cmap_vmin) / crange, 1), 0)
+            cindex = int(normalized * 255)
+            color = cmap(cindex)
+            points_by_color[color].append(point)
+
+        for color, points in points_by_color.items():
+            builder.add_material(color, opacity)
+            material_index = builder.material_count - 1
+
+            uri = f"layer_{unique_id()}_{color_identifier(color, opacity)}"
+
+            barr = bytearray()
+            add_points_to_bytearray(barr, points)
+            point_mins = index_mins(points)
+            point_maxes = index_maxes(points)
+
+            builder.add_buffer(byte_length=len(barr), uri=uri)
+            builder.add_buffer_view(
+                buffer=builder.buffer_count-1,
+                byte_length=len(barr),
+                byte_offset=0,
+                target=BufferTarget.ARRAY_BUFFER
+            )
+            builder.add_accessor(
+                buffer_view=builder.buffer_view_count-1,
+                component_type=ComponentType.FLOAT,
+                count=len(points),
+                type=AccessorType.VEC3,
+                mins=point_mins,
+                maxes=point_maxes
+            )
+            builder.add_mesh(
+                position_accessor=builder.accessor_count-1,
+                material=material_index,
+                mode=PrimitiveMode.POINTS,
+            )
+            builder.add_file_resource(uri, data=barr)
+
+
+@ar_layer_export(ScatterLayerState, "Points", ARPointExportOptions, ("gltf", "glb"))
+def add_vispy_points_layer_gltf(builder: GLTFBuilder,
+                                viewer_state: Vispy3DViewerState,
+                                layer_state: ScatterLayerState,
+                                options: ARPointExportOptions,
+                                bounds: Bounds,
+                                clip_to_bounds: bool = True):
+    add_points_layer_gltf(builder=builder,
+                          viewer_state=viewer_state,
+                          layer_state=layer_state,
+                          bounds=bounds,
+                          clip_to_bounds=clip_to_bounds)
+
+
+@ar_layer_export(Scatter3DLayerState, "Points", ARPointExportOptions, ("gltf", "glb"))
+def add_ipyvolume_points_layer_gltf(builder: GLTFBuilder,
+                                    viewer_state: Vispy3DViewerState,
+                                    layer_state: ScatterLayerState,
+                                    options: ARPointExportOptions,
+                                    bounds: Bounds,
+                                    clip_to_bounds: bool = True):
+    add_points_layer_gltf(builder=builder,
+                          viewer_state=viewer_state,
+                          layer_state=layer_state,
+                          bounds=bounds,
+                          clip_to_bounds=clip_to_bounds)

glue_ar/common/points_usd.py

@@ -0,0 +1,89 @@
+from collections import defaultdict
+from gltflib import AccessorType, BufferTarget, ComponentType, PrimitiveMode
+from glue_vispy_viewers.common.viewer_state import Vispy3DViewerState
+from glue_vispy_viewers.scatter.layer_state import ScatterLayerState
+
+from glue_ar.common.export_options import ar_layer_export
+from glue_ar.common.scatter import Scatter3DLayerState, ScatterLayerState3D, scatter_layer_mask
+from glue_ar.usd_utils import material_for_color
+from glue_ar.utils import Bounds, NoneType, Viewer3DState, color_identifier, hex_to_components, layer_color, unique_id, xyz_bounds, xyz_for_layer
+from glue_ar.common.usd_builder import USDBuilder
+from glue_ar.common.scatter_export_options import ARPointExportOptions
+
+try:
+    from glue_jupyter.common.state3d import ViewerState3D
+except ImportError:
+    ViewerState3D = NoneType
+
+
+def add_points_layer_usd(builder: USDBuilder,
+                         viewer_state: Viewer3DState,
+                         layer_state: ScatterLayerState3D,
+                         bounds: Bounds,
+                         clip_to_bounds: bool = True):
+
+    if layer_state is None:
+        return
+
+    bounds = xyz_bounds(viewer_state, with_resolution=False)
+
+    vispy_layer_state = isinstance(layer_state, ScatterLayerState)
+    color_mode_attr = "color_mode" if vispy_layer_state else "cmap_mode"
+    fixed_color = getattr(layer_state, color_mode_attr, "Fixed") == "Fixed"
+
+    mask = scatter_layer_mask(viewer_state, layer_state, bounds, clip_to_bounds)
+    data = xyz_for_layer(viewer_state, layer_state,
+                         preserve_aspect=viewer_state.native_aspect,
+                         mask=mask,
+                         scaled=True)
+    data = data[:, [1, 2, 0]]
+
+    identifier = f"layer_{unique_id()}"
+
+    if fixed_color:
+        color = layer_color(layer_state)
+        components = hex_to_components(color)[:3]
+        colors = [components for _ in range(data.shape[0])]
+    else:
+        cmap = layer_state.cmap
+        cmap_attr = "cmap_attribute" if vispy_layer_state else "cmap_att"
+        cmap_att = getattr(layer_state, cmap_attr)
+        cmap_vals = layer_state.layer[cmap_att][mask]
+        crange = layer_state.cmap_vmax - layer_state.cmap_vmin
+
+        def get_color(cval):
+            normalized = max(min((cval - layer_state.cmap_vmin) / crange, 1), 0)
+            cindex = int(normalized * 255)
+            return cmap(cindex)[:3]
+
+        colors = [get_color(cval) for cval in cmap_vals]
+
+    builder.add_points(data, colors, identifier)
+
+
+@ar_layer_export(ScatterLayerState, "Points", ARPointExportOptions, ("usdz", "usdc", "usda"))
+def add_vispy_points_layer_usd(builder: USDBuilder,
+                               viewer_state: Vispy3DViewerState,
+                               layer_state: ScatterLayerState,
+                               options: ARPointExportOptions,
+                               bounds: Bounds,
+                               clip_to_bounds: bool = True):
+    add_points_layer_usd(builder=builder,
+                         viewer_state=viewer_state,
+                         layer_state=layer_state,
+                         bounds=bounds,
+                         clip_to_bounds=clip_to_bounds)
+
+
+@ar_layer_export(Scatter3DLayerState, "Points", ARPointExportOptions, ("usdz", "usdc", "usda"))
+def add_ipyvolume_points_layer_usd(builder: USDBuilder,
+                                   viewer_state: ViewerState3D,
+                                   layer_state: ScatterLayerState,
+                                   options: ARPointExportOptions,
+                                   bounds: Bounds,
+                                   clip_to_bounds: bool = True):
+    add_points_layer_usd(builder=builder,
+                         viewer_state=viewer_state,
+                         layer_state=layer_state,
+                         bounds=bounds,
+                         clip_to_bounds=clip_to_bounds)

glue_ar/common/scatter_export_options.py

@@ -13,3 +13,7 @@ class ARVispyScatterExportOptions(State):
 
 class ARIpyvolumeScatterExportOptions(State):
     pass
+
+
+class ARPointExportOptions(State):
+    pass

glue_ar/common/scatter_usd.py

@@ -8,7 +8,7 @@
 
 from glue_ar.common.export_options import ar_layer_export
 from glue_ar.common.scatter import IPYVOLUME_POINTS_GETTERS, IPYVOLUME_TRIANGLE_GETTERS, VECTOR_OFFSETS, PointsGetter, \
-                                   ScatterLayerState3D, box_points_getter, radius_for_scatter_layer, \
+                                   Scatter3DLayerState, ScatterLayerState3D, box_points_getter, radius_for_scatter_layer, \
                                    scatter_layer_mask, sizes_for_scatter_layer, sphere_points_getter
 from glue_ar.common.scatter_export_options import ARIpyvolumeScatterExportOptions, ARVispyScatterExportOptions
 from glue_ar.common.usd_builder import USDBuilder
@@ -19,10 +19,8 @@
 
 try:
     from glue_jupyter.common.state3d import ViewerState3D
-    from glue_jupyter.ipyvolume.scatter import Scatter3DLayerState
 except ImportError:
     ViewerState3D = NoneType
-    Scatter3DLayerState = NoneType
 
 
 def add_vectors_usd(builder: USDBuilder,

glue_ar/common/usd_builder.py

@@ -2,8 +2,9 @@
 from os import extsep, remove
 from os.path import splitext
 
-from pxr import Usd, UsdGeom, UsdLux, UsdShade, UsdUtils
+from pxr import Gf, Sdf, Usd, UsdGeom, UsdLux, UsdShade, UsdUtils
 from typing import Dict, Iterable, Optional, Tuple
+from glue_ar.common.scatter import Point
 
 from glue_ar.usd_utils import material_for_color, material_for_mesh
 from glue_ar.utils import unique_id
@@ -37,7 +38,7 @@ def _create_stage(self, filepath: str):
         self.default_prim = UsdGeom.Xform.Define(self.stage, self.default_prim_key).GetPrim()
         self.stage.SetDefaultPrim(self.default_prim)
 
-        self._mesh_counts: Dict[str, int] = defaultdict(int)
+        self._geom_counts: Dict[str, int] = defaultdict(int)
 
         light = UsdLux.RectLight.Define(self.stage, "/light")
         light.CreateHeightAttr(-1)
@@ -76,11 +77,11 @@ def add_mesh(self,
         """
         identifier = identifier or unique_id()
         identifier = self._sanitize(identifier)
-        count = self._mesh_counts[identifier]
+        count = self._geom_counts[identifier]
         xform_key = f"{self.default_prim_key}/xform_{identifier}_{count}"
         UsdGeom.Xform.Define(self.stage, xform_key)
         mesh_key = f"{xform_key}/mesh_{identifier}_{count}"
-        self._mesh_counts[identifier] += 1
+        self._geom_counts[identifier] += 1
         mesh = UsdGeom.Mesh.Define(self.stage, mesh_key)
         mesh.CreateSubdivisionSchemeAttr().Set(UsdGeom.Tokens.none)
         mesh.CreatePointsAttr(points)
@@ -101,11 +102,11 @@ def add_translated_reference(self,
         prim = mesh.GetPrim()
         identifier = identifier or unique_id()
         identifier = self._sanitize(identifier)
-        count = self._mesh_counts[identifier]
+        count = self._geom_counts[identifier]
         xform_key = f"{self.default_prim_key}/xform_{identifier}_{count}"
         UsdGeom.Xform.Define(self.stage, xform_key)
         new_mesh_key = f"{xform_key}/mesh_{identifier}_{count}"
-        self._mesh_counts[identifier] += 1
+        self._geom_counts[identifier] += 1
         new_mesh = UsdGeom.Mesh.Define(self.stage, new_mesh_key)
         new_prim = new_mesh.GetPrim()
 
@@ -122,6 +123,30 @@ def add_translated_reference(self,
 
         return mesh
 
+    def add_points(self,
+                   points: Iterable[Point],
+                   colors: Iterable[Tuple[int, int, int]],
+                   identifier: Optional[str] = None) -> UsdGeom.Points:
+
+        identifier = identifier or unique_id()
+        identifier = self._sanitize(identifier)
+        count = self._geom_counts[identifier]
+        xform_key = f"{self.default_prim_key}/xform_{identifier}_{count}"
+        UsdGeom.Xform.Define(self.stage, xform_key)
+        points_key = f"{xform_key}/points_{identifier}_{count}"
+        self._geom_counts[identifier] += 1
+        geom = UsdGeom.Points.Define(self.stage, points_key)
+        point_vecs = [Gf.Vec3f(*point) for point in points]
+        geom.CreatePointsAttr(point_vecs)
+        widths_var = geom.CreateWidthsAttr()
+        widths = [1.0 for _ in points]
+        widths_var.Set(widths)
+        primvars = UsdGeom.PrimvarsAPI(geom.GetPrim())
+        colorvar = primvars.CreatePrimvar("displayColor", Sdf.ValueTypeNames.Color3f)
+        colorvar.Set([Gf.Vec3f(*tuple(c / 255 for c in color)) for color in colors])
+
+        return geom
+
     def export(self, filepath: str):
         base, ext = splitext(filepath)
         if ext == ".usdz":

glue_ar/usd_utils.py

@@ -2,9 +2,7 @@
 
 from pxr import Sdf, Usd, UsdGeom, UsdShade
 
-
-def color_identifier(color: Tuple[int, int, int], opacity: float = 1.0) -> str:
-    return f"{'_'.join(str(c) for c in color)}_{opacity}".replace(".", "_")
+from glue_ar.utils import color_identifier
 
 
 def material_for_color(stage: Usd.Stage,

glue_ar/utils.py

@@ -314,5 +314,9 @@ def binned_opacity(raw_opacity: float, resolution: float) -> float:
     return clamped_opacity(round(raw_opacity / resolution) * resolution)
 
 
+def color_identifier(color: Tuple[int, int, int], opacity: float = 1.0) -> str:
+    return f"{'_'.join(str(c) for c in color)}_{opacity}".replace(".", "_")
+
+
 def offset_triangles(triangle_indices, offset):
     return [tuple(idx + offset for idx in triangle) for triangle in triangle_indices]