diff --git a/glue_ar/common/__init__.py b/glue_ar/common/__init__.py index 0ba9dcb..7ead488 100644 --- a/glue_ar/common/__init__.py +++ b/glue_ar/common/__init__.py @@ -5,6 +5,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 diff --git a/glue_ar/common/points_gltf.py b/glue_ar/common/points_gltf.py new file mode 100644 index 0000000..31f834c --- /dev/null +++ b/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: ViewerState3D, + 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) diff --git a/glue_ar/common/points_usd.py b/glue_ar/common/points_usd.py new file mode 100644 index 0000000..8bb8d32 --- /dev/null +++ b/glue_ar/common/points_usd.py @@ -0,0 +1,87 @@ +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.utils import Bounds, NoneType, Viewer3DState, 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) diff --git a/glue_ar/common/scatter_export_options.py b/glue_ar/common/scatter_export_options.py index 681e7e8..33136ff 100644 --- a/glue_ar/common/scatter_export_options.py +++ b/glue_ar/common/scatter_export_options.py @@ -12,3 +12,7 @@ class ARVispyScatterExportOptions(State): class ARIpyvolumeScatterExportOptions(State): pass + + +class ARPointExportOptions(State): + pass diff --git a/glue_ar/common/scatter_usd.py b/glue_ar/common/scatter_usd.py index 98953ee..bee37db 100644 --- a/glue_ar/common/scatter_usd.py +++ b/glue_ar/common/scatter_usd.py @@ -8,8 +8,9 @@ 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, clip_vector_data, radius_for_scatter_layer, \ - scatter_layer_mask, sizes_for_scatter_layer, sphere_points_getter + Scatter3DLayerState, ScatterLayerState3D, box_points_getter, clip_vector_data, \ + 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 from glue_ar.common.shapes import cone_triangles, cone_points, cylinder_points, cylinder_triangles, \ @@ -20,10 +21,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, diff --git a/glue_ar/common/tests/test_base_dialog.py b/glue_ar/common/tests/test_base_dialog.py index 1ed3081..91266db 100644 --- a/glue_ar/common/tests/test_base_dialog.py +++ b/glue_ar/common/tests/test_base_dialog.py @@ -63,7 +63,7 @@ def test_layer_change_state(self): state = self.dialog.state state.layer = "Scatter Data" - assert state.method_helper.choices == ["Scatter"] + assert state.method_helper.choices == ["Scatter", "Points"] assert state.method == "Scatter" state.layer = "Volume Data" @@ -71,7 +71,7 @@ def test_layer_change_state(self): assert state.method in {"Isosurface", "Voxel"} state.layer = "Scatter Data" - assert state.method_helper.choices == ["Scatter"] + assert state.method_helper.choices == ["Scatter", "Points"] assert state.method == "Scatter" def test_method_settings_persistence(self): diff --git a/glue_ar/common/usd_builder.py b/glue_ar/common/usd_builder.py index 28059dd..96ca02d 100644 --- a/glue_ar/common/usd_builder.py +++ b/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.registries import builder from glue_ar.usd_utils import material_for_color, material_for_mesh, sanitize_path @@ -33,7 +34,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) @@ -70,12 +71,13 @@ def add_mesh(self, This breaks the builder pattern but we'll potentially want this reference to it for other meshes that we create. """ + identifier = sanitize_path(identifier or unique_id()) - 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) @@ -95,11 +97,11 @@ def add_translated_reference(self, identifier: Optional[str] = None) -> UsdGeom.Mesh: prim = mesh.GetPrim() identifier = sanitize_path(identifier or unique_id()) - 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() @@ -116,6 +118,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": diff --git a/glue_ar/jupyter/tests/test_dialog.py b/glue_ar/jupyter/tests/test_dialog.py index a778b85..89379ac 100644 --- a/glue_ar/jupyter/tests/test_dialog.py +++ b/glue_ar/jupyter/tests/test_dialog.py @@ -123,7 +123,10 @@ def test_layer_change_ui(self): state.layer = "Scatter Data" assert self.dialog.method_selected == 0 - assert self.dialog.method_items == [{"text": "Scatter", "value": 0}] + assert self.dialog.method_items == [ + {"text": "Scatter", "value": 0}, + {"text": "Points", "value": 1}, + ] assert not self.dialog.has_layer_options state.layer = "Volume Data" @@ -133,7 +136,10 @@ def test_layer_change_ui(self): state.layer = "Scatter Data" assert self.dialog.method_selected == 0 - assert self.dialog.method_items == [{"text": "Scatter", "value": 0}] + assert self.dialog.method_items == [ + {"text": "Scatter", "value": 0}, + {"text": "Points", "value": 1}, + ] assert not self.dialog.has_layer_options def test_on_cancel(self): diff --git a/glue_ar/qt/tests/test_dialog.py b/glue_ar/qt/tests/test_dialog.py index ad7557f..bfaec56 100644 --- a/glue_ar/qt/tests/test_dialog.py +++ b/glue_ar/qt/tests/test_dialog.py @@ -122,9 +122,9 @@ def test_layer_change_ui(self): state.layer = "Scatter Data" assert ui.combosel_method.currentText() == state.method - assert combobox_options(ui.combosel_method) == ["Scatter"] - assert not ui.label_method.isVisible() - assert not ui.combosel_method.isVisible() + assert combobox_options(ui.combosel_method) == ["Scatter", "Points"] + assert ui.label_method.isVisible() + assert ui.combosel_method.isVisible() state.layer = "Volume Data" assert set(combobox_options(ui.combosel_method)) == {"Isosurface", "Voxel"} @@ -134,6 +134,6 @@ def test_layer_change_ui(self): state.layer = "Scatter Data" assert ui.combosel_method.currentText() == state.method - assert combobox_options(ui.combosel_method) == ["Scatter"] - assert not ui.label_method.isVisible() - assert not ui.combosel_method.isVisible() + assert combobox_options(ui.combosel_method) == ["Scatter", "Points"] + assert ui.label_method.isVisible() + assert ui.combosel_method.isVisible() diff --git a/glue_ar/usd_utils.py b/glue_ar/usd_utils.py index c9aa2f6..c3ecd2d 100644 --- a/glue_ar/usd_utils.py +++ b/glue_ar/usd_utils.py @@ -3,9 +3,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, diff --git a/glue_ar/utils.py b/glue_ar/utils.py index 359b97c..5ee5e18 100644 --- a/glue_ar/utils.py +++ b/glue_ar/utils.py @@ -364,5 +364,9 @@ def binned_opacity(raw_opacity: float, resolution: float) -> float: return clamp_with_resolution(raw_opacity, 0, 1, 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]