diff --git a/glue_ar/common/__init__.py b/glue_ar/common/__init__.py index 1df2b48..0279bf7 100644 --- a/glue_ar/common/__init__.py +++ b/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 diff --git a/glue_ar/common/points_gltf.py b/glue_ar/common/points_gltf.py new file mode 100644 index 0000000..995ac17 --- /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: 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) diff --git a/glue_ar/common/points_usd.py b/glue_ar/common/points_usd.py new file mode 100644 index 0000000..971adee --- /dev/null +++ b/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) diff --git a/glue_ar/common/scatter_export_options.py b/glue_ar/common/scatter_export_options.py index 4cc1db9..34a8694 100644 --- a/glue_ar/common/scatter_export_options.py +++ b/glue_ar/common/scatter_export_options.py @@ -13,3 +13,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 ec6a6be..23854b4 100644 --- a/glue_ar/common/scatter_usd.py +++ b/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, diff --git a/glue_ar/common/usd_builder.py b/glue_ar/common/usd_builder.py index 106e2a2..7cbc5b0 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.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": diff --git a/glue_ar/usd_utils.py b/glue_ar/usd_utils.py index 9d57b53..29238fc 100644 --- a/glue_ar/usd_utils.py +++ b/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, diff --git a/glue_ar/utils.py b/glue_ar/utils.py index 332f21e..9bfa991 100644 --- a/glue_ar/utils.py +++ b/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]