Merge pull request #60 from Carifio24/usdz

Allow exporting to USDZ

glue_ar/common/export.py

@@ -28,7 +28,8 @@
     "gltf": GLTFBuilder,
     "glb": GLTFBuilder,
     "usda": USDBuilder,
-    "usdc": USDBuilder
+    "usdc": USDBuilder,
+    "usdz": USDBuilder,
 }
 
 

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']

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

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,

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"}

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)

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

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"}

glue_ar/qt/export_tool.py

@@ -18,6 +18,7 @@
 _FILETYPE_NAMES = {
     "gltf": "glTF",
     "glb": "glB",
+    "usdz": "USDZ",
     "usdc": "USDC",
     "usda": "USDA",
 }

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, \

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, \