diff --git a/glue_ar/common/export.py b/glue_ar/common/export.py index 675220e..1d8e377 100644 --- a/glue_ar/common/export.py +++ b/glue_ar/common/export.py @@ -28,7 +28,8 @@ "gltf": GLTFBuilder, "glb": GLTFBuilder, "usda": USDBuilder, - "usdc": USDBuilder + "usdc": USDBuilder, + "usdz": USDBuilder, } diff --git a/glue_ar/common/export_state.py b/glue_ar/common/export_state.py index e400db6..0aa865a 100644 --- a/glue_ar/common/export_state.py +++ b/glue_ar/common/export_state.py @@ -23,7 +23,7 @@ def __init__(self, layers: Iterable[LayerState]): super(ARExportDialogState, self).__init__() self.filetype_helper = ComboHelper(self, 'filetype') - self.filetype_helper.choices = ['glB', 'glTF', 'USDC', 'USDA'] + self.filetype_helper.choices = ['glB', 'glTF', 'USDZ', 'USDC', 'USDA'] self.compression_helper = ComboHelper(self, 'compression') self.compression_helper.choices = ['None', 'Draco', 'Meshoptimizer'] diff --git a/glue_ar/common/marching_cubes.py b/glue_ar/common/marching_cubes.py index 54f1eac..9cbca0f 100644 --- a/glue_ar/common/marching_cubes.py +++ b/glue_ar/common/marching_cubes.py @@ -114,7 +114,7 @@ def add_isosurface_layer_gltf(builder: GLTFBuilder, builder.add_file_resource(level_bin, data=barr) -@ar_layer_export(VolumeLayerState, "Isosurface", ARIsosurfaceExportOptions, ("usdc", "usda")) +@ar_layer_export(VolumeLayerState, "Isosurface", ARIsosurfaceExportOptions, ("usdz", "usdc", "usda")) def add_isosurface_layer_usd( builder: USDBuilder, viewer_state: Vispy3DVolumeViewerState, @@ -171,6 +171,6 @@ def add_isosurface_layer_usd( ar_layer_export.add(IPVVolumeLayerState, "Isosurface", ARIsosurfaceExportOptions, ("gltf", "glb"), False, add_isosurface_layer_gltf) ar_layer_export.add(IPVVolumeLayerState, "Isosurface", ARIsosurfaceExportOptions, - ("usda", "usdc"), False, add_isosurface_layer_usd) + ("usda", "usdc", "usdz"), False, add_isosurface_layer_usd) except ImportError: pass diff --git a/glue_ar/common/scatter_usd.py b/glue_ar/common/scatter_usd.py index 307488f..717f729 100644 --- a/glue_ar/common/scatter_usd.py +++ b/glue_ar/common/scatter_usd.py @@ -188,7 +188,7 @@ def add_scatter_layer_usd( ) -@ar_layer_export(ScatterLayerState, "Scatter", ARVispyScatterExportOptions, ("usdc", "usda")) +@ar_layer_export(ScatterLayerState, "Scatter", ARVispyScatterExportOptions, ("usdz", "usdc", "usda")) def add_vispy_scatter_layer_usd(builder: USDBuilder, viewer_state: Vispy3DViewerState, layer_state: ScatterLayerState, @@ -211,7 +211,7 @@ def add_vispy_scatter_layer_usd(builder: USDBuilder, clip_to_bounds=clip_to_bounds) -@ar_layer_export(Scatter3DLayerState, "Scatter", ARIpyvolumeScatterExportOptions, ("usdc", "usda")) +@ar_layer_export(Scatter3DLayerState, "Scatter", ARIpyvolumeScatterExportOptions, ("usdz", "usdc", "usda")) def add_ipyvolume_scatter_layer_usd(builder: USDBuilder, viewer_state: ViewerState3D, layer_state: Scatter3DLayerState, diff --git a/glue_ar/common/tests/test_base_dialog.py b/glue_ar/common/tests/test_base_dialog.py index 898c060..25c5996 100644 --- a/glue_ar/common/tests/test_base_dialog.py +++ b/glue_ar/common/tests/test_base_dialog.py @@ -49,7 +49,7 @@ def test_default_state(self): assert state.layer == "Volume Data" assert state.method in {"Isosurface", "Voxel"} - assert state.filetype_helper.choices == ['glB', 'glTF', 'USDC', 'USDA'] + assert state.filetype_helper.choices == ['glB', 'glTF', 'USDZ', 'USDC', 'USDA'] assert state.compression_helper.choices == ['None', 'Draco', 'Meshoptimizer'] assert state.layer_helper.choices == ["Volume Data", "Scatter Data"] assert set(state.method_helper.choices) == {"Isosurface", "Voxel"} diff --git a/glue_ar/common/usd_builder.py b/glue_ar/common/usd_builder.py index 1acbb40..43556d6 100644 --- a/glue_ar/common/usd_builder.py +++ b/glue_ar/common/usd_builder.py @@ -1,8 +1,11 @@ from collections import defaultdict -from pxr import Usd, UsdGeom, UsdLux, UsdShade +from os import extsep, remove +from os.path import splitext + +from pxr import Usd, UsdGeom, UsdLux, UsdShade, UsdUtils from typing import Dict, Iterable, Optional, Tuple -from glue_ar.usd_utils import material_for_color, material_for_mesh +from glue_ar.usd_utils import material_for_color, material_for_mesh from glue_ar.utils import unique_id @@ -12,6 +15,9 @@ class USDBuilder: def __init__(self, filepath: str): + base, ext = splitext(filepath) + if ext == ".usdz": + filepath = f"{base}{extsep}usdc" self._create_stage(filepath) self._material_map: Dict[MaterialInfo, UsdShade.Shader] = {} @@ -113,8 +119,15 @@ def add_translated_reference(self, return mesh - def export(self, filepath): - self.stage.GetRootLayer().Export(filepath) - - def build_and_export(self, filepath): + def export(self, filepath: str): + base, ext = splitext(filepath) + if ext == ".usdz": + usdc_path = f"{base}{extsep}usdc" + self.stage.GetRootLayer().Export(usdc_path) + UsdUtils.CreateNewUsdzPackage(usdc_path, filepath) + remove(usdc_path) + else: + self.stage.GetRootLayer().Export(filepath) + + def build_and_export(self, filepath: str): self.export(filepath) diff --git a/glue_ar/common/voxels.py b/glue_ar/common/voxels.py index c5f37e7..061fa4b 100644 --- a/glue_ar/common/voxels.py +++ b/glue_ar/common/voxels.py @@ -162,7 +162,7 @@ def add_voxel_layers_gltf(builder: GLTFBuilder, builder.add_file_resource(triangles_bin, data=triangles_barr) -@ar_layer_export(VolumeLayerState, "Voxel", ARVoxelExportOptions, ("usda", "usdc"), multiple=True) +@ar_layer_export(VolumeLayerState, "Voxel", ARVoxelExportOptions, ("usda", "usdc", "usdz"), multiple=True) def add_voxel_layers_usd(builder: USDBuilder, viewer_state: Vispy3DVolumeViewerState, layer_states: Iterable[VolumeLayerState], @@ -252,6 +252,6 @@ def add_voxel_layers_usd(builder: USDBuilder, ar_layer_export.add(IPVVolumeLayerState, "Voxel", ARVoxelExportOptions, ("gltf", "glb"), True, add_voxel_layers_gltf) ar_layer_export.add(IPVVolumeLayerState, "Voxel", ARVoxelExportOptions, - ("usda", "usdc"), True, add_voxel_layers_usd) + ("usda", "usdc", "usdz"), True, add_voxel_layers_usd) except ImportError: pass diff --git a/glue_ar/jupyter/tests/test_dialog.py b/glue_ar/jupyter/tests/test_dialog.py index c7041df..d9dc962 100644 --- a/glue_ar/jupyter/tests/test_dialog.py +++ b/glue_ar/jupyter/tests/test_dialog.py @@ -49,8 +49,9 @@ def test_default_ui(self): assert self.dialog.filetype_items == [ {"text": "glB", "value": 0}, {"text": "glTF", "value": 1}, - {"text": "USDC", "value": 2}, - {"text": "USDA", "value": 3} + {"text": "USDZ", "value": 2}, + {"text": "USDC", "value": 3}, + {"text": "USDA", "value": 4}, ] assert self.dialog.filetype_selected == 0 assert set([item["text"] for item in self.dialog.method_items]) == {"Isosurface", "Voxel"} diff --git a/glue_ar/qt/export_tool.py b/glue_ar/qt/export_tool.py index a39037d..8deb786 100644 --- a/glue_ar/qt/export_tool.py +++ b/glue_ar/qt/export_tool.py @@ -18,6 +18,7 @@ _FILETYPE_NAMES = { "gltf": "glTF", "glb": "glB", + "usdz": "USDZ", "usdc": "USDC", "usda": "USDA", } diff --git a/glue_ar/qt/tests/test_tool_scatter.py b/glue_ar/qt/tests/test_tool_scatter.py index c2d1f54..640b498 100644 --- a/glue_ar/qt/tests/test_tool_scatter.py +++ b/glue_ar/qt/tests/test_tool_scatter.py @@ -55,7 +55,7 @@ def test_toolbar(self): assert len([subtool for subtool in tool.subtools if isinstance(subtool, QtARExportTool)]) == 1 @pytest.mark.parametrize("extension,compression", - product(("glB", "glTF", "USDA", "USDC"), ("None", "Draco", "Meshoptimizer"))) + product(("glB", "glTF", "USDA", "USDC", "USDZ"), ("None", "Draco", "Meshoptimizer"))) def test_tool_export_call(self, extension, compression): auto_accept = dialog_auto_accept_with_options(filetype=extension, compression=compression) with patch("qtpy.compat.getsavefilename") as fd, \ diff --git a/glue_ar/qt/tests/test_tool_volume.py b/glue_ar/qt/tests/test_tool_volume.py index c81b772..7c27c22 100644 --- a/glue_ar/qt/tests/test_tool_volume.py +++ b/glue_ar/qt/tests/test_tool_volume.py @@ -60,7 +60,7 @@ def test_toolbar(self): assert len([subtool for subtool in tool.subtools if isinstance(subtool, QtARExportTool)]) == 1 @pytest.mark.parametrize("extension,compression", - product(("glB", "glTF", "USDA", "USDC"), ("None", "Draco", "Meshoptimizer"))) + product(("glB", "glTF", "USDA", "USDC", "USDZ"), ("None", "Draco", "Meshoptimizer"))) def test_tool_export_call(self, extension, compression): auto_accept = dialog_auto_accept_with_options(filetype=extension, compression=compression) with patch("qtpy.compat.getsavefilename") as fd, \