ansys-pyensight-core 0.8.11__py3-none-any.whl → 0.8.13__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ansys-pyensight-core might be problematic. Click here for more details.

@@ -1,690 +1,882 @@
1
- #
2
- # This file borrows heavily from the Omniverse Example Connector which
3
- # contains the following notice:
4
- #
5
- ###############################################################################
6
- # Copyright 2020 NVIDIA Corporation
7
- #
8
- # Permission is hereby granted, free of charge, to any person obtaining a copy of
9
- # this software and associated documentation files (the "Software"), to deal in
10
- # the Software without restriction, including without limitation the rights to
11
- # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
12
- # the Software, and to permit persons to whom the Software is furnished to do so,
13
- # subject to the following conditions:
14
- #
15
- # The above copyright notice and this permission notice shall be included in all
16
- # copies or substantial portions of the Software.
17
- #
18
- # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19
- # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
20
- # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
21
- # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
22
- # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
23
- # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
24
- #
25
- ###############################################################################
26
-
27
- import logging
28
- import math
29
- import os
30
- import shutil
31
- import tempfile
32
- from typing import Any, Dict, List, Optional
33
-
34
- from ansys.pyensight.core.utils.dsg_server import Part, UpdateHandler
35
- import png
36
- from pxr import Gf, Sdf, Usd, UsdGeom, UsdLux, UsdShade
37
-
38
-
39
- class OmniverseWrapper(object):
40
- def __init__(self, live_edit: bool = False, destination: str = "") -> None:
41
- self._cleaned_index = 0
42
- self._cleaned_names: dict = {}
43
- self._connectionStatusSubscription = None
44
- self._stage = None
45
- self._destinationPath: str = ""
46
- self._old_stages: list = []
47
- self._stagename = "dsg_scene.usd"
48
- self._live_edit: bool = live_edit
49
- if self._live_edit:
50
- self._stagename = "dsg_scene.live"
51
- # USD time slider will have 120 tick marks per second of animation time
52
- self._time_codes_per_second = 120.0
53
-
54
- if destination:
55
- self.destination = destination
56
-
57
- @property
58
- def destination(self) -> str:
59
- """The current output directory."""
60
- return self._destinationPath
61
-
62
- @destination.setter
63
- def destination(self, directory: str) -> None:
64
- self._destinationPath = directory
65
- if not self.is_valid_destination(directory):
66
- logging.warning(f"Invalid destination path: {directory}")
67
-
68
- def shutdown(self) -> None:
69
- """
70
- Shutdown the connection to Omniverse cleanly.
71
- """
72
- self._connectionStatusSubscription = None
73
-
74
- @staticmethod
75
- def is_valid_destination(path: str) -> bool:
76
- """
77
- Verify that the target path is a writeable directory.
78
-
79
- Parameters
80
- ----------
81
- path
82
- The path to check
83
-
84
- Returns
85
- -------
86
- True if the path is a writeable directory, False otherwise.
87
- """
88
- return os.access(path, os.W_OK)
89
-
90
- def stage_url(self, name: Optional[str] = None) -> str:
91
- """
92
- For a given object name, create the URL for the item.
93
- Parameters
94
- ----------
95
- name: the name of the object to generate the URL for. If None, it will be the URL for the
96
- stage name.
97
-
98
- Returns
99
- -------
100
- The URL for the object.
101
- """
102
- if name is None:
103
- name = self._stagename
104
- return os.path.join(self._destinationPath, name)
105
-
106
- def delete_old_stages(self) -> None:
107
- """
108
- Remove all the stages included in the "_old_stages" list.
109
- If a stage is in use and cannot be removed, keep its name in _old_stages
110
- to retry later.
111
- """
112
- stages_unremoved = list()
113
- while self._old_stages:
114
- stage = self._old_stages.pop()
115
- try:
116
- if os.path.isfile(stage):
117
- os.remove(stage)
118
- else:
119
- shutil.rmtree(stage, ignore_errors=True, onerror=None)
120
- except OSError:
121
- stages_unremoved.append(stage)
122
- self._old_stages = stages_unremoved
123
-
124
- def create_new_stage(self) -> None:
125
- """
126
- Create a new stage. using the current stage name.
127
- """
128
- logging.info(f"Creating Omniverse stage: {self.stage_url()}")
129
- if self._stage:
130
- self._stage.Unload()
131
- self._stage = None
132
- self.delete_old_stages()
133
- self._stage = Usd.Stage.CreateNew(self.stage_url())
134
- # record the stage in the "_old_stages" list.
135
- self._old_stages.append(self.stage_url())
136
- UsdGeom.SetStageUpAxis(self._stage, UsdGeom.Tokens.y)
137
- # in M
138
- UsdGeom.SetStageMetersPerUnit(self._stage, 1.0)
139
- logging.info(f"Created stage: {self.stage_url()}")
140
-
141
- def save_stage(self, comment: str = "") -> None:
142
- """
143
- For live connections, save the current edit and allow live processing.
144
-
145
- Presently, live connections are disabled.
146
- """
147
- self._stage.GetRootLayer().Save() # type:ignore
148
-
149
- def clear_cleaned_names(self) -> None:
150
- """
151
- Clear the list of cleaned names
152
- """
153
- self._cleaned_names = {}
154
- self._cleaned_index = 0
155
-
156
- def clean_name(self, name: str, id_name: Any = None) -> str:
157
- """Generate a valid USD name
158
-
159
- From a base (EnSight) varname, partname, etc. and the DSG id, generate
160
- a unique, valid USD name. Save the names so that if the same name
161
- comes in again, the previously computed name is returned and if the
162
- manipulation results in a conflict, the name can be made unique.
163
-
164
- Parameters
165
- ----------
166
- name:
167
- The name to generate a USD name for.
168
-
169
- id_name:
170
- The DSG id associated with the DSG name, if any.
171
-
172
- Returns
173
- -------
174
- A unique USD name.
175
- """
176
- orig_name = name
177
- # return any previously generated name
178
- if (name, id_name) in self._cleaned_names:
179
- return self._cleaned_names[(name, id_name)]
180
- # replace invalid characters. EnSight uses a number of characters that are illegal in USD names.
181
- replacements = {
182
- ord("+"): "_",
183
- ord("-"): "_",
184
- ord("."): "_",
185
- ord(":"): "_",
186
- ord("["): "_",
187
- ord("]"): "_",
188
- ord("("): "_",
189
- ord(")"): "_",
190
- ord("<"): "_",
191
- ord(">"): "_",
192
- ord("/"): "_",
193
- ord("="): "_",
194
- ord(","): "_",
195
- ord(" "): "_",
196
- ord("\\"): "_",
197
- }
198
- name = name.translate(replacements)
199
- if name[0].isdigit():
200
- name = f"_{name}"
201
- if id_name is not None:
202
- name = name + "_" + str(id_name)
203
- if name in self._cleaned_names.values():
204
- # Make the name unique
205
- while f"{name}_{self._cleaned_index}" in self._cleaned_names.values():
206
- self._cleaned_index += 1
207
- name = f"{name}_{self._cleaned_index}"
208
- # store off the cleaned name
209
- self._cleaned_names[(orig_name, id_name)] = name
210
- return name
211
-
212
- @staticmethod
213
- def decompose_matrix(values: Any) -> Any:
214
- """
215
- Decompose an array of floats (representing a 4x4 matrix) into scale, rotation and translation.
216
- Parameters
217
- ----------
218
- values:
219
- 16 values (input to Gf.Matrix4f CTOR)
220
-
221
- Returns
222
- -------
223
- (scale, rotation, translation)
224
- """
225
- # ang_convert = 180.0/math.pi
226
- ang_convert = 1.0
227
- trans_convert = 1.0
228
- m = Gf.Matrix4f(*values)
229
- m = m.GetTranspose()
230
-
231
- s = math.sqrt(m[0][0] * m[0][0] + m[0][1] * m[0][1] + m[0][2] * m[0][2])
232
- # cleanup scale
233
- m = m.RemoveScaleShear()
234
- # r = m.ExtractRotation()
235
- R = m.ExtractRotationMatrix()
236
- r = [
237
- math.atan2(R[2][1], R[2][2]) * ang_convert,
238
- math.atan2(-R[2][0], 1.0) * ang_convert,
239
- math.atan2(R[1][0], R[0][0]) * ang_convert,
240
- ]
241
- t = m.ExtractTranslation()
242
- t = [t[0] * trans_convert, t[1] * trans_convert, t[2] * trans_convert]
243
- return s, r, t
244
-
245
- def create_dsg_mesh_block(
246
- self,
247
- name,
248
- id,
249
- part_hash,
250
- parent_prim,
251
- verts,
252
- conn,
253
- normals,
254
- tcoords,
255
- matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
256
- diffuse=[1.0, 1.0, 1.0, 1.0],
257
- variable=None,
258
- timeline=[0.0, 0.0],
259
- first_timestep=False,
260
- ):
261
- # 1D texture map for variables https://graphics.pixar.com/usd/release/tut_simple_shading.html
262
- # create the part usd object
263
- partname = self.clean_name(name + part_hash.hexdigest())
264
- stage_name = "/Parts/" + partname + ".usd"
265
- part_stage_url = self.stage_url(os.path.join("Parts", partname + ".usd"))
266
- part_stage = None
267
-
268
- if not os.path.exists(part_stage_url):
269
- part_stage = Usd.Stage.CreateNew(part_stage_url)
270
- self._old_stages.append(part_stage_url)
271
- xform = UsdGeom.Xform.Define(part_stage, "/" + partname)
272
- mesh = UsdGeom.Mesh.Define(part_stage, "/" + partname + "/Mesh")
273
- # mesh.CreateDisplayColorAttr()
274
- mesh.CreateDoubleSidedAttr().Set(True)
275
- mesh.CreatePointsAttr(verts)
276
- mesh.CreateNormalsAttr(normals)
277
- mesh.CreateFaceVertexCountsAttr([3] * int(conn.size / 3))
278
- mesh.CreateFaceVertexIndicesAttr(conn)
279
- if (tcoords is not None) and variable:
280
- # USD 22.08 changed the primvar API
281
- if hasattr(mesh, "CreatePrimvar"):
282
- texCoords = mesh.CreatePrimvar(
283
- "st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.varying
284
- )
285
- else:
286
- primvarsAPI = UsdGeom.PrimvarsAPI(mesh)
287
- texCoords = primvarsAPI.CreatePrimvar(
288
- "st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.varying
289
- )
290
- texCoords.Set(tcoords)
291
- texCoords.SetInterpolation("vertex")
292
- part_prim = part_stage.GetPrimAtPath("/" + partname)
293
- part_stage.SetDefaultPrim(part_prim)
294
-
295
- # Currently, this will never happen, but it is a setup for rigid body transforms
296
- # At present, the group transforms have been cooked into the vertices so this is not needed
297
- matrixOp = xform.AddXformOp(
298
- UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
299
- )
300
- matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
301
-
302
- self.create_dsg_material(
303
- part_stage, mesh, "/" + partname, diffuse=diffuse, variable=variable
304
- )
305
-
306
- timestep_prim = self.add_timestep_group(parent_prim, timeline, first_timestep)
307
-
308
- # glue it into our stage
309
- path = timestep_prim.GetPath().AppendChild("part_ref_" + partname)
310
- part_ref = self._stage.OverridePrim(path)
311
- part_ref.GetReferences().AddReference("." + stage_name)
312
-
313
- if part_stage is not None:
314
- part_stage.GetRootLayer().Save()
315
-
316
- return part_stage_url
317
-
318
- def add_timestep_group(
319
- self, parent_prim: UsdGeom.Xform, timeline: List[float], first_timestep: bool
320
- ) -> UsdGeom.Xform:
321
- # add a layer in the group hierarchy for the timestep
322
- timestep_group_path = parent_prim.GetPath().AppendChild(
323
- self.clean_name("t" + str(timeline[0]), None)
324
- )
325
- timestep_prim = UsdGeom.Xform.Define(self._stage, timestep_group_path)
326
- visibility_attr = UsdGeom.Imageable(timestep_prim).GetVisibilityAttr()
327
- if first_timestep:
328
- visibility_attr.Set("inherited", Usd.TimeCode.EarliestTime())
329
- else:
330
- visibility_attr.Set("invisible", Usd.TimeCode.EarliestTime())
331
- visibility_attr.Set("inherited", timeline[0] * self._time_codes_per_second)
332
- # Final timestep has timeline[0]==timeline[1]. Leave final timestep visible.
333
- if timeline[0] < timeline[1]:
334
- visibility_attr.Set("invisible", timeline[1] * self._time_codes_per_second)
335
- return timestep_prim
336
-
337
- def create_dsg_points(
338
- self,
339
- name,
340
- id,
341
- part_hash,
342
- parent_prim,
343
- verts,
344
- sizes,
345
- colors,
346
- matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
347
- default_size=1.0,
348
- default_color=[1.0, 1.0, 1.0, 1.0],
349
- timeline=[0.0, 0.0],
350
- first_timestep=False,
351
- ):
352
- # create the part usd object
353
- partname = self.clean_name(name + part_hash.hexdigest())
354
- stage_name = "/Parts/" + partname + ".usd"
355
- part_stage_url = self.stage_url(os.path.join("Parts", partname + ".usd"))
356
- part_stage = None
357
-
358
- if not os.path.exists(part_stage_url):
359
- part_stage = Usd.Stage.CreateNew(part_stage_url)
360
- self._old_stages.append(part_stage_url)
361
- xform = UsdGeom.Xform.Define(part_stage, "/" + partname)
362
-
363
- points = UsdGeom.Points.Define(part_stage, "/" + partname + "/Points")
364
- # points.GetPointsAttr().Set(Vt.Vec3fArray(verts.tolist()))
365
- points.GetPointsAttr().Set(verts)
366
- if sizes is not None and sizes.size == (verts.size // 3):
367
- points.GetWidthsAttr().Set(sizes)
368
- else:
369
- points.GetWidthsAttr().Set([default_size] * (verts.size // 3))
370
-
371
- colorAttr = points.GetPrim().GetAttribute("primvars:displayColor")
372
- colorAttr.SetMetadata("interpolation", "vertex")
373
- if colors is not None and colors.size == verts.size:
374
- colorAttr.Set(colors)
375
- else:
376
- colorAttr.Set([default_color[0:3]] * (verts.size // 3))
377
-
378
- part_prim = part_stage.GetPrimAtPath("/" + partname)
379
- part_stage.SetDefaultPrim(part_prim)
380
-
381
- # Currently, this will never happen, but it is a setup for rigid body transforms
382
- # At present, the group transforms have been cooked into the vertices so this is not needed
383
- matrixOp = xform.AddXformOp(
384
- UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
385
- )
386
- matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
387
-
388
- timestep_prim = self.add_timestep_group(parent_prim, timeline, first_timestep)
389
-
390
- # glue it into our stage
391
- path = timestep_prim.GetPath().AppendChild("part_ref_" + partname)
392
- part_ref = self._stage.OverridePrim(path)
393
- part_ref.GetReferences().AddReference("." + stage_name)
394
-
395
- if part_stage is not None:
396
- part_stage.GetRootLayer().Save()
397
-
398
- return part_stage_url
399
-
400
- def create_dsg_material(
401
- self, stage, mesh, root_name, diffuse=[1.0, 1.0, 1.0, 1.0], variable=None
402
- ):
403
- # https://graphics.pixar.com/usd/release/spec_usdpreviewsurface.html
404
- material = UsdShade.Material.Define(stage, root_name + "/Material")
405
- pbrShader = UsdShade.Shader.Define(stage, root_name + "/Material/PBRShader")
406
- pbrShader.CreateIdAttr("UsdPreviewSurface")
407
- pbrShader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0)
408
- pbrShader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
409
- pbrShader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(diffuse[3])
410
- pbrShader.CreateInput("useSpecularWorkflow", Sdf.ValueTypeNames.Int).Set(1)
411
- if variable:
412
- stReader = UsdShade.Shader.Define(stage, root_name + "/Material/stReader")
413
- stReader.CreateIdAttr("UsdPrimvarReader_float2")
414
- diffuseTextureSampler = UsdShade.Shader.Define(
415
- stage, root_name + "/Material/diffuseTexture"
416
- )
417
- diffuseTextureSampler.CreateIdAttr("UsdUVTexture")
418
- name = self.clean_name(variable.name)
419
- filename = self._destinationPath + f"/Parts/Textures/palette_{name}.png"
420
- diffuseTextureSampler.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(filename)
421
- diffuseTextureSampler.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
422
- stReader.ConnectableAPI(), "result"
423
- )
424
- diffuseTextureSampler.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
425
- pbrShader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(
426
- diffuseTextureSampler.ConnectableAPI(), "rgb"
427
- )
428
- stInput = material.CreateInput("frame:stPrimvarName", Sdf.ValueTypeNames.Token)
429
- stInput.Set("st")
430
- stReader.CreateInput("varname", Sdf.ValueTypeNames.Token).ConnectToSource(stInput)
431
- else:
432
- scale = 1.0
433
- color = Gf.Vec3f(diffuse[0] * scale, diffuse[1] * scale, diffuse[2] * scale)
434
- pbrShader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(color)
435
-
436
- material.CreateSurfaceOutput().ConnectToSource(pbrShader.ConnectableAPI(), "surface")
437
- UsdShade.MaterialBindingAPI(mesh).Bind(material)
438
-
439
- return material
440
-
441
- def create_dsg_variable_textures(self, variables):
442
- with tempfile.TemporaryDirectory() as tempdir:
443
- # make folder: {tempdir}/scratch/Textures/{palette_*.png}
444
- os.makedirs(f"{tempdir}/scratch/Textures", exist_ok=True)
445
- for var in variables.values():
446
- data = bytearray(var.texture)
447
- n_pixels = int(len(data) / 4)
448
- row = []
449
- for i in range(n_pixels):
450
- row.append(data[i * 4 + 0])
451
- row.append(data[i * 4 + 1])
452
- row.append(data[i * 4 + 2])
453
- io = png.Writer(width=n_pixels, height=2, bitdepth=8, greyscale=False)
454
- rows = [row, row]
455
- name = self.clean_name(var.name)
456
- with open(f"{tempdir}/scratch/Textures/palette_{name}.png", "wb") as fp:
457
- io.write(fp, rows)
458
- uriPath = self._destinationPath + "/Parts/Textures"
459
- shutil.rmtree(uriPath, ignore_errors=True, onerror=None)
460
- shutil.copytree(f"{tempdir}/scratch/Textures", uriPath)
461
-
462
- def create_dsg_root(self):
463
- root_name = "/Root"
464
- root_prim = UsdGeom.Xform.Define(self._stage, root_name)
465
- # Define the defaultPrim as the /Root prim
466
- root_prim = self._stage.GetPrimAtPath(root_name)
467
- self._stage.SetDefaultPrim(root_prim)
468
- return root_prim
469
-
470
- def update_camera(self, camera):
471
- if camera is not None:
472
- cam_name = "/Root/Cam"
473
- cam_prim = UsdGeom.Xform.Define(self._stage, cam_name)
474
- cam_pos = Gf.Vec3d(camera.lookfrom[0], camera.lookfrom[1], camera.lookfrom[2])
475
- target_pos = Gf.Vec3d(camera.lookat[0], camera.lookat[1], camera.lookat[2])
476
- up_vec = Gf.Vec3d(camera.upvector[0], camera.upvector[1], camera.upvector[2])
477
- cam_prim = self._stage.GetPrimAtPath(cam_name)
478
- geom_cam = UsdGeom.Camera(cam_prim)
479
- if not geom_cam:
480
- geom_cam = UsdGeom.Camera.Define(self._stage, cam_name)
481
- # Set camera values
482
- # center of interest attribute unique for Kit defines the pivot for tumbling the camera
483
- # Set as an attribute on the prim
484
- coi_attr = cam_prim.GetAttribute("omni:kit:centerOfInterest")
485
- if not coi_attr.IsValid():
486
- coi_attr = cam_prim.CreateAttribute(
487
- "omni:kit:centerOfInterest", Sdf.ValueTypeNames.Vector3d
488
- )
489
- coi_attr.Set(target_pos)
490
- # get the camera
491
- cam = geom_cam.GetCamera()
492
- # LOL, not sure why is might be correct, but so far it seems to work???
493
- cam.focalLength = camera.fieldofview
494
- cam.clippingRange = Gf.Range1f(0.1, 10)
495
- look_at = Gf.Matrix4d()
496
- look_at.SetLookAt(cam_pos, target_pos, up_vec)
497
- trans_row = look_at.GetRow(3)
498
- trans_row = Gf.Vec4d(-trans_row[0], -trans_row[1], -trans_row[2], trans_row[3])
499
- look_at.SetRow(3, trans_row)
500
- # print(look_at)
501
- cam.transform = look_at
502
-
503
- # set the updated camera
504
- geom_cam.SetFromCamera(cam)
505
-
506
- def create_dsg_group(
507
- self,
508
- name: str,
509
- parent_prim,
510
- obj_type: Any = None,
511
- matrix: List[float] = [
512
- 1.0,
513
- 0.0,
514
- 0.0,
515
- 0.0,
516
- 0.0,
517
- 1.0,
518
- 0.0,
519
- 0.0,
520
- 0.0,
521
- 0.0,
522
- 1.0,
523
- 0.0,
524
- 0.0,
525
- 0.0,
526
- 0.0,
527
- 1.0,
528
- ],
529
- ):
530
- path = parent_prim.GetPath().AppendChild(self.clean_name(name))
531
- group_prim = UsdGeom.Xform.Get(self._stage, path)
532
- if not group_prim:
533
- group_prim = UsdGeom.Xform.Define(self._stage, path)
534
- # At present, the group transforms have been cooked into the vertices so this is not needed
535
- matrixOp = group_prim.AddXformOp(
536
- UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
537
- )
538
- matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
539
- logging.info(f"Created group:'{name}' {str(obj_type)}")
540
- return group_prim
541
-
542
- def uploadMaterial(self):
543
- uriPath = self._destinationPath + "/Materials"
544
- shutil.rmtree(uriPath, ignore_errors=True, onerror=None)
545
- fullpath = os.path.join(os.path.dirname(__file__), "resources", "Materials")
546
- shutil.copytree(fullpath, uriPath)
547
-
548
- # Create a dome light in the scene.
549
- def createDomeLight(self, texturePath):
550
- newLight = UsdLux.DomeLight.Define(self._stage, "/Root/DomeLight")
551
- newLight.CreateIntensityAttr(2200.0)
552
- newLight.CreateTextureFileAttr(texturePath)
553
- newLight.CreateTextureFormatAttr("latlong")
554
-
555
- # Set rotation on domelight
556
- xForm = newLight
557
- rotateOp = xForm.AddXformOp(UsdGeom.XformOp.TypeRotateZYX, UsdGeom.XformOp.PrecisionFloat)
558
- rotateOp.Set(Gf.Vec3f(270, 0, 0))
559
-
560
-
561
- class OmniverseUpdateHandler(UpdateHandler):
562
- """
563
- Implement the Omniverse glue to a DSGSession instance
564
- """
565
-
566
- def __init__(self, omni: OmniverseWrapper):
567
- super().__init__()
568
- self._omni = omni
569
- self._group_prims: Dict[int, Any] = dict()
570
- self._root_prim = None
571
- self._sent_textures = False
572
-
573
- def add_group(self, id: int, view: bool = False) -> None:
574
- super().add_group(id, view)
575
- group = self.session.groups[id]
576
- if not view:
577
- parent_prim = self._group_prims[group.parent_id]
578
- obj_type = self.get_dsg_cmd_attribute(group, "ENS_OBJ_TYPE")
579
- matrix = self.group_matrix(group)
580
- prim = self._omni.create_dsg_group(
581
- group.name, parent_prim, matrix=matrix, obj_type=obj_type
582
- )
583
- self._group_prims[id] = prim
584
- else:
585
- # Map a view command into a new Omniverse stage and populate it with materials/lights.
586
- # Create a new root stage in Omniverse
587
-
588
- # Create or update the root group/camera
589
- if not self.session.vrmode:
590
- self._omni.update_camera(camera=group)
591
-
592
- # record
593
- self._group_prims[id] = self._root_prim
594
-
595
- if self._omni._stage is not None:
596
- self._omni._stage.SetStartTimeCode(
597
- self.session.time_limits[0] * self._omni._time_codes_per_second
598
- )
599
- self._omni._stage.SetEndTimeCode(
600
- self.session.time_limits[1] * self._omni._time_codes_per_second
601
- )
602
- self._omni._stage.SetTimeCodesPerSecond(self._omni._time_codes_per_second)
603
-
604
- # Send the variable textures. Safe to do so once the first view is processed.
605
- if not self._sent_textures:
606
- self._omni.create_dsg_variable_textures(self.session.variables)
607
- self._sent_textures = True
608
-
609
- def add_variable(self, id: int) -> None:
610
- super().add_variable(id)
611
-
612
- def finalize_part(self, part: Part) -> None:
613
- # generate an Omniverse compliant mesh from the Part
614
- if part is None or part.cmd is None:
615
- return
616
- parent_prim = self._group_prims[part.cmd.parent_id]
617
- obj_id = self.session.mesh_block_count
618
- matrix = part.cmd.matrix4x4
619
- name = part.cmd.name
620
- color = [
621
- part.cmd.fill_color[0] * part.cmd.diffuse,
622
- part.cmd.fill_color[1] * part.cmd.diffuse,
623
- part.cmd.fill_color[2] * part.cmd.diffuse,
624
- part.cmd.fill_color[3],
625
- ]
626
-
627
- if part.cmd.render == part.cmd.CONNECTIVITY:
628
- command, verts, conn, normals, tcoords, var_cmd = part.nodal_surface_rep()
629
- if command is not None:
630
- # Generate the mesh block
631
- _ = self._omni.create_dsg_mesh_block(
632
- name,
633
- obj_id,
634
- part.hash,
635
- parent_prim,
636
- verts,
637
- conn,
638
- normals,
639
- tcoords,
640
- matrix=matrix,
641
- diffuse=color,
642
- variable=var_cmd,
643
- timeline=self.session.cur_timeline,
644
- first_timestep=(self.session.cur_timeline[0] == self.session.time_limits[0]),
645
- )
646
-
647
- elif part.cmd.render == part.cmd.NODES:
648
- command, verts, sizes, colors, var_cmd = part.point_rep()
649
- if command is not None:
650
- _ = self._omni.create_dsg_points(
651
- name,
652
- obj_id,
653
- part.hash,
654
- parent_prim,
655
- verts,
656
- sizes,
657
- colors,
658
- matrix=matrix,
659
- default_size=part.cmd.node_size_default,
660
- default_color=color,
661
- timeline=self.session.cur_timeline,
662
- first_timestep=(self.session.cur_timeline[0] == self.session.time_limits[0]),
663
- )
664
- super().finalize_part(part)
665
-
666
- def start_connection(self) -> None:
667
- super().start_connection()
668
-
669
- def end_connection(self) -> None:
670
- super().end_connection()
671
-
672
- def begin_update(self) -> None:
673
- super().begin_update()
674
- # restart the name tables
675
- self._omni.clear_cleaned_names()
676
- # clear the group Omni prims list
677
- self._group_prims = dict()
678
-
679
- self._omni.create_new_stage()
680
- self._root_prim = self._omni.create_dsg_root()
681
- # Create a distance and dome light in the scene
682
- self._omni.createDomeLight("./Materials/000_sky.exr")
683
- # Upload a material to the Omniverse server
684
- self._omni.uploadMaterial()
685
- self._sent_textures = False
686
-
687
- def end_update(self) -> None:
688
- super().end_update()
689
- # Stage update complete
690
- self._omni.save_stage()
1
+ #
2
+ # This file borrows heavily from the Omniverse Example Connector which
3
+ # contains the following notice:
4
+ #
5
+ ###############################################################################
6
+ # Copyright 2020 NVIDIA Corporation
7
+ #
8
+ # Permission is hereby granted, free of charge, to any person obtaining a copy of
9
+ # this software and associated documentation files (the "Software"), to deal in
10
+ # the Software without restriction, including without limitation the rights to
11
+ # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
12
+ # the Software, and to permit persons to whom the Software is furnished to do so,
13
+ # subject to the following conditions:
14
+ #
15
+ # The above copyright notice and this permission notice shall be included in all
16
+ # copies or substantial portions of the Software.
17
+ #
18
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
20
+ # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
21
+ # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
22
+ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
23
+ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
24
+ #
25
+ ###############################################################################
26
+ import logging
27
+ import math
28
+ import os
29
+ import shutil
30
+ import tempfile
31
+ from typing import Any, Dict, List, Optional
32
+
33
+ from ansys.pyensight.core.utils.dsg_server import Part, UpdateHandler
34
+ import numpy
35
+ import png
36
+ from pxr import Gf, Kind, Sdf, Usd, UsdGeom, UsdLux, UsdShade
37
+
38
+
39
+ class OmniverseWrapper(object):
40
+ def __init__(
41
+ self,
42
+ live_edit: bool = False,
43
+ destination: str = "",
44
+ line_width: float = -0.0001,
45
+ use_lines: bool = False,
46
+ ) -> None:
47
+ self._cleaned_index = 0
48
+ self._cleaned_names: dict = {}
49
+ self._connectionStatusSubscription = None
50
+ self._stage = None
51
+ self._destinationPath: str = ""
52
+ self._old_stages: list = []
53
+ self._stagename = "dsg_scene.usd"
54
+ self._live_edit: bool = live_edit
55
+ if self._live_edit:
56
+ self._stagename = "dsg_scene.live"
57
+ # USD time slider will have 120 tick marks per second of animation time
58
+ self._time_codes_per_second = 120.0
59
+
60
+ if destination:
61
+ self.destination = destination
62
+
63
+ self._line_width = line_width
64
+ self._use_lines = use_lines
65
+
66
+ @property
67
+ def destination(self) -> str:
68
+ """The current output directory."""
69
+ return self._destinationPath
70
+
71
+ @destination.setter
72
+ def destination(self, directory: str) -> None:
73
+ self._destinationPath = directory
74
+ if not self.is_valid_destination(directory):
75
+ logging.warning(f"Invalid destination path: {directory}")
76
+
77
+ @property
78
+ def line_width(self) -> float:
79
+ return self._line_width
80
+
81
+ @line_width.setter
82
+ def line_width(self, line_width: float) -> None:
83
+ self._line_width = line_width
84
+
85
+ @property
86
+ def use_lines(self) -> bool:
87
+ return self._use_lines
88
+
89
+ def shutdown(self) -> None:
90
+ """
91
+ Shutdown the connection to Omniverse cleanly.
92
+ """
93
+ self._connectionStatusSubscription = None
94
+
95
+ @staticmethod
96
+ def is_valid_destination(path: str) -> bool:
97
+ """
98
+ Verify that the target path is a writeable directory.
99
+
100
+ Parameters
101
+ ----------
102
+ path
103
+ The path to check
104
+
105
+ Returns
106
+ -------
107
+ True if the path is a writeable directory, False otherwise.
108
+ """
109
+ return os.access(path, os.W_OK)
110
+
111
+ def stage_url(self, name: Optional[str] = None) -> str:
112
+ """
113
+ For a given object name, create the URL for the item.
114
+ Parameters
115
+ ----------
116
+ name: the name of the object to generate the URL for. If None, it will be the URL for the
117
+ stage name.
118
+
119
+ Returns
120
+ -------
121
+ The URL for the object.
122
+ """
123
+ if name is None:
124
+ name = self._stagename
125
+ return os.path.join(self._destinationPath, name)
126
+
127
+ def delete_old_stages(self) -> None:
128
+ """
129
+ Remove all the stages included in the "_old_stages" list.
130
+ If a stage is in use and cannot be removed, keep its name in _old_stages
131
+ to retry later.
132
+ """
133
+ stages_unremoved = list()
134
+ while self._old_stages:
135
+ stage = self._old_stages.pop()
136
+ try:
137
+ if os.path.isfile(stage):
138
+ os.remove(stage)
139
+ else:
140
+ shutil.rmtree(stage, ignore_errors=True, onerror=None)
141
+ except OSError:
142
+ stages_unremoved.append(stage)
143
+ self._old_stages = stages_unremoved
144
+
145
+ def create_new_stage(self) -> None:
146
+ """
147
+ Create a new stage. using the current stage name.
148
+ """
149
+ logging.info(f"Creating Omniverse stage: {self.stage_url()}")
150
+ if self._stage:
151
+ self._stage.Unload()
152
+ self._stage = None
153
+ self.delete_old_stages()
154
+ self._stage = Usd.Stage.CreateNew(self.stage_url())
155
+ # record the stage in the "_old_stages" list.
156
+ self._old_stages.append(self.stage_url())
157
+ UsdGeom.SetStageUpAxis(self._stage, UsdGeom.Tokens.y)
158
+ # in M
159
+ UsdGeom.SetStageMetersPerUnit(self._stage, 1.0)
160
+ logging.info(f"Created stage: {self.stage_url()}")
161
+
162
+ def save_stage(self, comment: str = "") -> None:
163
+ """
164
+ For live connections, save the current edit and allow live processing.
165
+
166
+ Presently, live connections are disabled.
167
+ """
168
+ self._stage.GetRootLayer().Save() # type:ignore
169
+
170
+ def clear_cleaned_names(self) -> None:
171
+ """
172
+ Clear the list of cleaned names
173
+ """
174
+ self._cleaned_names = {}
175
+ self._cleaned_index = 0
176
+
177
+ def clean_name(self, name: str, id_name: Any = None) -> str:
178
+ """Generate a valid USD name
179
+
180
+ From a base (EnSight) varname, partname, etc. and the DSG id, generate
181
+ a unique, valid USD name. Save the names so that if the same name
182
+ comes in again, the previously computed name is returned and if the
183
+ manipulation results in a conflict, the name can be made unique.
184
+
185
+ Parameters
186
+ ----------
187
+ name:
188
+ The name to generate a USD name for.
189
+
190
+ id_name:
191
+ The DSG id associated with the DSG name, if any.
192
+
193
+ Returns
194
+ -------
195
+ A unique USD name.
196
+ """
197
+ orig_name = name
198
+ # return any previously generated name
199
+ if (name, id_name) in self._cleaned_names:
200
+ return self._cleaned_names[(name, id_name)]
201
+ # replace invalid characters. EnSight uses a number of characters that are illegal in USD names.
202
+ replacements = {
203
+ ord("+"): "_",
204
+ ord("-"): "_",
205
+ ord("."): "_",
206
+ ord(":"): "_",
207
+ ord("["): "_",
208
+ ord("]"): "_",
209
+ ord("("): "_",
210
+ ord(")"): "_",
211
+ ord("<"): "_",
212
+ ord(">"): "_",
213
+ ord("/"): "_",
214
+ ord("="): "_",
215
+ ord(","): "_",
216
+ ord(" "): "_",
217
+ ord("\\"): "_",
218
+ }
219
+ name = name.translate(replacements)
220
+ if name[0].isdigit():
221
+ name = f"_{name}"
222
+ if id_name is not None:
223
+ name = name + "_" + str(id_name)
224
+ if name in self._cleaned_names.values():
225
+ # Make the name unique
226
+ while f"{name}_{self._cleaned_index}" in self._cleaned_names.values():
227
+ self._cleaned_index += 1
228
+ name = f"{name}_{self._cleaned_index}"
229
+ # store off the cleaned name
230
+ self._cleaned_names[(orig_name, id_name)] = name
231
+ return name
232
+
233
+ @staticmethod
234
+ def decompose_matrix(values: Any) -> Any:
235
+ """
236
+ Decompose an array of floats (representing a 4x4 matrix) into scale, rotation and translation.
237
+ Parameters
238
+ ----------
239
+ values:
240
+ 16 values (input to Gf.Matrix4f CTOR)
241
+
242
+ Returns
243
+ -------
244
+ (scale, rotation, translation)
245
+ """
246
+ # ang_convert = 180.0/math.pi
247
+ ang_convert = 1.0
248
+ trans_convert = 1.0
249
+ m = Gf.Matrix4f(*values)
250
+ m = m.GetTranspose()
251
+
252
+ s = math.sqrt(m[0][0] * m[0][0] + m[0][1] * m[0][1] + m[0][2] * m[0][2])
253
+ # cleanup scale
254
+ m = m.RemoveScaleShear()
255
+ # r = m.ExtractRotation()
256
+ R = m.ExtractRotationMatrix()
257
+ r = [
258
+ math.atan2(R[2][1], R[2][2]) * ang_convert,
259
+ math.atan2(-R[2][0], 1.0) * ang_convert,
260
+ math.atan2(R[1][0], R[0][0]) * ang_convert,
261
+ ]
262
+ t = m.ExtractTranslation()
263
+ t = [t[0] * trans_convert, t[1] * trans_convert, t[2] * trans_convert]
264
+ return s, r, t
265
+
266
+ def create_dsg_mesh_block(
267
+ self,
268
+ name,
269
+ id,
270
+ part_hash,
271
+ parent_prim,
272
+ verts,
273
+ conn,
274
+ normals,
275
+ tcoords,
276
+ matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
277
+ diffuse=[1.0, 1.0, 1.0, 1.0],
278
+ variable=None,
279
+ timeline=[0.0, 0.0],
280
+ first_timestep=False,
281
+ mat_info={},
282
+ ):
283
+ # 1D texture map for variables https://graphics.pixar.com/usd/release/tut_simple_shading.html
284
+ # create the part usd object
285
+ partname = self.clean_name(name + part_hash.hexdigest())
286
+ stage_name = "/Parts/" + partname + ".usd"
287
+ part_stage_url = self.stage_url(os.path.join("Parts", partname + ".usd"))
288
+ part_stage = None
289
+
290
+ if not os.path.exists(part_stage_url):
291
+ part_stage = Usd.Stage.CreateNew(part_stage_url)
292
+ self._old_stages.append(part_stage_url)
293
+ xform = UsdGeom.Xform.Define(part_stage, "/" + partname)
294
+ mesh = UsdGeom.Mesh.Define(part_stage, "/" + partname + "/Mesh")
295
+ # mesh.CreateDisplayColorAttr()
296
+ mesh.CreateDoubleSidedAttr().Set(True)
297
+ mesh.CreatePointsAttr(verts)
298
+ mesh.CreateNormalsAttr(normals)
299
+ mesh.CreateFaceVertexCountsAttr([3] * (conn.size // 3))
300
+ mesh.CreateFaceVertexIndicesAttr(conn)
301
+ if (tcoords is not None) and variable:
302
+ primvarsAPI = UsdGeom.PrimvarsAPI(mesh)
303
+ texCoords = primvarsAPI.CreatePrimvar(
304
+ "st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.varying
305
+ )
306
+ texCoords.Set(tcoords)
307
+ texCoords.SetInterpolation("vertex")
308
+ part_prim = part_stage.GetPrimAtPath("/" + partname)
309
+ part_stage.SetDefaultPrim(part_prim)
310
+
311
+ # Currently, this will never happen, but it is a setup for rigid body transforms
312
+ # At present, the group transforms have been cooked into the vertices so this is not needed
313
+ matrixOp = xform.AddXformOp(
314
+ UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
315
+ )
316
+ matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
317
+
318
+ self.create_dsg_material(
319
+ part_stage,
320
+ mesh,
321
+ "/" + partname,
322
+ diffuse=diffuse,
323
+ variable=variable,
324
+ mat_info=mat_info,
325
+ )
326
+
327
+ timestep_prim = self.add_timestep_group(parent_prim, timeline, first_timestep)
328
+
329
+ # glue it into our stage
330
+ path = timestep_prim.GetPath().AppendChild("part_ref_" + partname)
331
+ part_ref = self._stage.OverridePrim(path)
332
+ part_ref.GetReferences().AddReference("." + stage_name)
333
+
334
+ if part_stage is not None:
335
+ part_stage.GetRootLayer().Save()
336
+
337
+ return part_stage_url
338
+
339
+ def add_timestep_group(
340
+ self, parent_prim: UsdGeom.Xform, timeline: List[float], first_timestep: bool
341
+ ) -> UsdGeom.Xform:
342
+ # add a layer in the group hierarchy for the timestep
343
+ timestep_group_path = parent_prim.GetPath().AppendChild(
344
+ self.clean_name("t" + str(timeline[0]), None)
345
+ )
346
+ timestep_prim = UsdGeom.Xform.Define(self._stage, timestep_group_path)
347
+ visibility_attr = UsdGeom.Imageable(timestep_prim).GetVisibilityAttr()
348
+ if first_timestep:
349
+ visibility_attr.Set("inherited", Usd.TimeCode.EarliestTime())
350
+ else:
351
+ visibility_attr.Set("invisible", Usd.TimeCode.EarliestTime())
352
+ visibility_attr.Set("inherited", timeline[0] * self._time_codes_per_second)
353
+ # Final timestep has timeline[0]==timeline[1]. Leave final timestep visible.
354
+ if timeline[0] < timeline[1]:
355
+ visibility_attr.Set("invisible", timeline[1] * self._time_codes_per_second)
356
+ return timestep_prim
357
+
358
+ def create_dsg_lines(
359
+ self,
360
+ name,
361
+ id,
362
+ part_hash,
363
+ parent_prim,
364
+ verts,
365
+ tcoords,
366
+ matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
367
+ diffuse=[1.0, 1.0, 1.0, 1.0],
368
+ variable=None,
369
+ timeline=[0.0, 0.0],
370
+ first_timestep=False,
371
+ mat_info={},
372
+ ):
373
+ # TODO: GLB extension maps to DSG PART attribute map
374
+ width = self.line_width
375
+ wireframe = width == 0.0
376
+ if width < 0.0:
377
+ tmp = verts.reshape(-1, 3)
378
+ mins = numpy.min(tmp, axis=0)
379
+ maxs = numpy.max(tmp, axis=0)
380
+ dx = maxs[0] - mins[0]
381
+ dy = maxs[1] - mins[1]
382
+ dz = maxs[2] - mins[2]
383
+ diagonal = math.sqrt(dx * dx + dy * dy + dz * dz)
384
+ width = diagonal * math.fabs(width)
385
+ self.line_width = width
386
+
387
+ # include the line width in the hash
388
+ part_hash.update(str(self.line_width).encode("utf-8"))
389
+
390
+ # 1D texture map for variables https://graphics.pixar.com/usd/release/tut_simple_shading.html
391
+ # create the part usd object
392
+ partname = self.clean_name(name + part_hash.hexdigest()) + "_l"
393
+ stage_name = "/Parts/" + partname + ".usd"
394
+ part_stage_url = self.stage_url(os.path.join("Parts", partname + ".usd"))
395
+ part_stage = None
396
+
397
+ var_cmd = variable
398
+
399
+ if not os.path.exists(part_stage_url):
400
+ part_stage = Usd.Stage.CreateNew(part_stage_url)
401
+ self._old_stages.append(part_stage_url)
402
+ xform = UsdGeom.Xform.Define(part_stage, "/" + partname)
403
+ lines = UsdGeom.BasisCurves.Define(part_stage, "/" + partname + "/Lines")
404
+ lines.CreateDoubleSidedAttr().Set(True)
405
+ lines.CreatePointsAttr(verts)
406
+ lines.CreateCurveVertexCountsAttr([2] * (verts.size // 6))
407
+ lines.CreatePurposeAttr().Set("render")
408
+ lines.CreateTypeAttr().Set("linear")
409
+ lines.CreateWidthsAttr([width])
410
+ lines.SetWidthsInterpolation("constant")
411
+ # Rounded endpoint are a primvar
412
+ primvarsAPI = UsdGeom.PrimvarsAPI(lines)
413
+ endCaps = primvarsAPI.CreatePrimvar(
414
+ "endcaps", Sdf.ValueTypeNames.Int, UsdGeom.Tokens.constant
415
+ )
416
+ endCaps.Set(2) # Rounded = 2
417
+
418
+ prim = lines.GetPrim()
419
+ prim.CreateAttribute(
420
+ "omni:scene:visualization:drawWireframe", Sdf.ValueTypeNames.Bool
421
+ ).Set(wireframe)
422
+ if (tcoords is not None) and var_cmd:
423
+ primvarsAPI = UsdGeom.PrimvarsAPI(lines)
424
+ texCoords = primvarsAPI.CreatePrimvar(
425
+ "st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.varying
426
+ )
427
+ texCoords.Set(tcoords)
428
+ texCoords.SetInterpolation("vertex")
429
+ part_prim = part_stage.GetPrimAtPath("/" + partname)
430
+ part_stage.SetDefaultPrim(part_prim)
431
+
432
+ # Currently, this will never happen, but it is a setup for rigid body transforms
433
+ # At present, the group transforms have been cooked into the vertices so this is not needed
434
+ matrixOp = xform.AddXformOp(
435
+ UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
436
+ )
437
+ matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
438
+
439
+ self.create_dsg_material(
440
+ part_stage,
441
+ lines,
442
+ "/" + partname,
443
+ diffuse=diffuse,
444
+ variable=var_cmd,
445
+ mat_info=mat_info,
446
+ )
447
+
448
+ timestep_prim = self.add_timestep_group(parent_prim, timeline, first_timestep)
449
+
450
+ # glue it into our stage
451
+ path = timestep_prim.GetPath().AppendChild("part_ref_" + partname)
452
+ part_ref = self._stage.OverridePrim(path)
453
+ part_ref.GetReferences().AddReference("." + stage_name)
454
+
455
+ if part_stage is not None:
456
+ part_stage.GetRootLayer().Save()
457
+
458
+ return part_stage_url
459
+
460
+ def create_dsg_points(
461
+ self,
462
+ name,
463
+ id,
464
+ part_hash,
465
+ parent_prim,
466
+ verts,
467
+ sizes,
468
+ colors,
469
+ matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
470
+ default_size=1.0,
471
+ default_color=[1.0, 1.0, 1.0, 1.0],
472
+ timeline=[0.0, 0.0],
473
+ first_timestep=False,
474
+ ):
475
+ # create the part usd object
476
+ partname = self.clean_name(name + part_hash.hexdigest())
477
+ stage_name = "/Parts/" + partname + ".usd"
478
+ part_stage_url = self.stage_url(os.path.join("Parts", partname + ".usd"))
479
+ part_stage = None
480
+
481
+ if not os.path.exists(part_stage_url):
482
+ part_stage = Usd.Stage.CreateNew(part_stage_url)
483
+ self._old_stages.append(part_stage_url)
484
+ xform = UsdGeom.Xform.Define(part_stage, "/" + partname)
485
+
486
+ points = UsdGeom.Points.Define(part_stage, "/" + partname + "/Points")
487
+ # points.GetPointsAttr().Set(Vt.Vec3fArray(verts.tolist()))
488
+ points.GetPointsAttr().Set(verts)
489
+ if sizes is not None and sizes.size == (verts.size // 3):
490
+ points.GetWidthsAttr().Set(sizes)
491
+ else:
492
+ points.GetWidthsAttr().Set([default_size] * (verts.size // 3))
493
+
494
+ colorAttr = points.GetPrim().GetAttribute("primvars:displayColor")
495
+ colorAttr.SetMetadata("interpolation", "vertex")
496
+ if colors is not None and colors.size == verts.size:
497
+ colorAttr.Set(colors)
498
+ else:
499
+ colorAttr.Set([default_color[0:3]] * (verts.size // 3))
500
+
501
+ part_prim = part_stage.GetPrimAtPath("/" + partname)
502
+ part_stage.SetDefaultPrim(part_prim)
503
+
504
+ # Currently, this will never happen, but it is a setup for rigid body transforms
505
+ # At present, the group transforms have been cooked into the vertices so this is not needed
506
+ matrixOp = xform.AddXformOp(
507
+ UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
508
+ )
509
+ matrixOp.Set(Gf.Matrix4d(*matrix).GetTranspose())
510
+
511
+ timestep_prim = self.add_timestep_group(parent_prim, timeline, first_timestep)
512
+
513
+ # glue it into our stage
514
+ path = timestep_prim.GetPath().AppendChild("part_ref_" + partname)
515
+ part_ref = self._stage.OverridePrim(path)
516
+ part_ref.GetReferences().AddReference("." + stage_name)
517
+
518
+ if part_stage is not None:
519
+ part_stage.GetRootLayer().Save()
520
+
521
+ return part_stage_url
522
+
523
+ def create_dsg_material(
524
+ self, stage, mesh, root_name, diffuse=[1.0, 1.0, 1.0, 1.0], variable=None, mat_info={}
525
+ ):
526
+ # https://graphics.pixar.com/usd/release/spec_usdpreviewsurface.html
527
+ # Use ior==1.0 to be more like EnSight - rays of light do not bend when passing through transparent objs
528
+ material = UsdShade.Material.Define(stage, root_name + "/Material")
529
+ pbrShader = UsdShade.Shader.Define(stage, root_name + "/Material/PBRShader")
530
+ pbrShader.CreateIdAttr("UsdPreviewSurface")
531
+ smoothness = mat_info.get("smoothness", 0.0)
532
+ pbrShader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - smoothness)
533
+ metallic = mat_info.get("metallic", 0.0)
534
+ pbrShader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(metallic)
535
+ opacity = mat_info.get("opacity", diffuse[3])
536
+ pbrShader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(opacity)
537
+ pbrShader.CreateInput("ior", Sdf.ValueTypeNames.Float).Set(1.0)
538
+ pbrShader.CreateInput("useSpecularWorkflow", Sdf.ValueTypeNames.Int).Set(1)
539
+ if variable:
540
+ stReader = UsdShade.Shader.Define(stage, root_name + "/Material/stReader")
541
+ stReader.CreateIdAttr("UsdPrimvarReader_float2")
542
+ diffuseTextureSampler = UsdShade.Shader.Define(
543
+ stage, root_name + "/Material/diffuseTexture"
544
+ )
545
+ diffuseTextureSampler.CreateIdAttr("UsdUVTexture")
546
+ name = self.clean_name(variable.name)
547
+ filename = f"./Textures/palette_{name}.png"
548
+ diffuseTextureSampler.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(filename)
549
+ diffuseTextureSampler.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
550
+ stReader.ConnectableAPI(), "result"
551
+ )
552
+ diffuseTextureSampler.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
553
+ pbrShader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(
554
+ diffuseTextureSampler.ConnectableAPI(), "rgb"
555
+ )
556
+ stInput = material.CreateInput("frame:stPrimvarName", Sdf.ValueTypeNames.Token)
557
+ stInput.Set("st")
558
+ stReader.CreateInput("varname", Sdf.ValueTypeNames.Token).ConnectToSource(stInput)
559
+ else:
560
+ # The colors are a mixture of content from the DSG PART protocol buffer
561
+ # and the JSON material block from the material_name field.
562
+ kd = 1.0
563
+ diffuse_color = [diffuse[0], diffuse[1], diffuse[2]]
564
+ ke = 1.0
565
+ emissive_color = [0.0, 0.0, 0.0]
566
+ ks = 1.0
567
+ specular_color = [0.0, 0.0, 0.0]
568
+ mat_name = mat_info.get("name", "")
569
+ if mat_name.startswith("ensight"):
570
+ diffuse_color = mat_info.get("diffuse", diffuse_color)
571
+ if mat_name != "ensight/Default":
572
+ ke = mat_info.get("ke", ke)
573
+ emissive_color = mat_info.get("emissive", emissive_color)
574
+ ks = mat_info.get("ks", ks)
575
+ specular_color = mat_info.get("specular", specular_color)
576
+ # Set the colors
577
+ color = Gf.Vec3f(diffuse_color[0] * kd, diffuse_color[1] * kd, diffuse_color[2] * kd)
578
+ pbrShader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(color)
579
+ color = Gf.Vec3f(emissive_color[0] * ke, emissive_color[1] * ke, emissive_color[2] * ke)
580
+ pbrShader.CreateInput("emissiveColor", Sdf.ValueTypeNames.Color3f).Set(color)
581
+ color = Gf.Vec3f(specular_color[0] * ks, specular_color[1] * ks, specular_color[2] * ks)
582
+ pbrShader.CreateInput("specularColor", Sdf.ValueTypeNames.Color3f).Set(color)
583
+
584
+ material.CreateSurfaceOutput().ConnectToSource(pbrShader.ConnectableAPI(), "surface")
585
+ UsdShade.MaterialBindingAPI(mesh).Bind(material)
586
+
587
+ return material
588
+
589
+ def create_dsg_variable_textures(self, variables):
590
+ with tempfile.TemporaryDirectory() as tempdir:
591
+ # make folder: {tempdir}/scratch/Textures/{palette_*.png}
592
+ os.makedirs(f"{tempdir}/scratch/Textures", exist_ok=True)
593
+ for var in variables.values():
594
+ data = bytearray(var.texture)
595
+ n_pixels = int(len(data) / 4)
596
+ row = []
597
+ for i in range(n_pixels):
598
+ row.append(data[i * 4 + 0])
599
+ row.append(data[i * 4 + 1])
600
+ row.append(data[i * 4 + 2])
601
+ io = png.Writer(width=n_pixels, height=2, bitdepth=8, greyscale=False)
602
+ rows = [row, row]
603
+ name = self.clean_name(var.name)
604
+ with open(f"{tempdir}/scratch/Textures/palette_{name}.png", "wb") as fp:
605
+ io.write(fp, rows)
606
+ uriPath = self._destinationPath + "/Parts/Textures"
607
+ shutil.rmtree(uriPath, ignore_errors=True, onerror=None)
608
+ shutil.copytree(f"{tempdir}/scratch/Textures", uriPath)
609
+
610
+ def create_dsg_root(self):
611
+ root_name = "/Root"
612
+ root_prim = UsdGeom.Xform.Define(self._stage, root_name)
613
+ # Define the defaultPrim as the /Root prim
614
+ root_prim = self._stage.GetPrimAtPath(root_name)
615
+ self._stage.SetDefaultPrim(root_prim)
616
+ return root_prim
617
+
618
+ def update_camera(self, camera):
619
+ if camera is not None:
620
+ cam_name = "/Root/Cam"
621
+ cam_prim = UsdGeom.Xform.Define(self._stage, cam_name)
622
+ cam_pos = Gf.Vec3d(camera.lookfrom[0], camera.lookfrom[1], camera.lookfrom[2])
623
+ target_pos = Gf.Vec3d(camera.lookat[0], camera.lookat[1], camera.lookat[2])
624
+ up_vec = Gf.Vec3d(camera.upvector[0], camera.upvector[1], camera.upvector[2])
625
+ cam_prim = self._stage.GetPrimAtPath(cam_name)
626
+ geom_cam = UsdGeom.Camera(cam_prim)
627
+ if not geom_cam:
628
+ geom_cam = UsdGeom.Camera.Define(self._stage, cam_name)
629
+ # Set camera values
630
+ # center of interest attribute unique for Kit defines the pivot for tumbling the camera
631
+ # Set as an attribute on the prim
632
+ coi_attr = cam_prim.GetAttribute("omni:kit:centerOfInterest")
633
+ if not coi_attr.IsValid():
634
+ coi_attr = cam_prim.CreateAttribute(
635
+ "omni:kit:centerOfInterest", Sdf.ValueTypeNames.Vector3d
636
+ )
637
+ coi_attr.Set(target_pos)
638
+ # get the camera
639
+ cam = geom_cam.GetCamera()
640
+ # LOL, not sure why is might be correct, but so far it seems to work???
641
+ cam.focalLength = camera.fieldofview
642
+ cam.clippingRange = Gf.Range1f(0.1, 10)
643
+ look_at = Gf.Matrix4d()
644
+ look_at.SetLookAt(cam_pos, target_pos, up_vec)
645
+ trans_row = look_at.GetRow(3)
646
+ trans_row = Gf.Vec4d(-trans_row[0], -trans_row[1], -trans_row[2], trans_row[3])
647
+ look_at.SetRow(3, trans_row)
648
+ cam.transform = look_at
649
+
650
+ # set the updated camera
651
+ geom_cam.SetFromCamera(cam)
652
+
653
+ def create_dsg_group(
654
+ self,
655
+ name: str,
656
+ parent_prim,
657
+ obj_type: Any = None,
658
+ matrix: List[float] = [
659
+ 1.0,
660
+ 0.0,
661
+ 0.0,
662
+ 0.0,
663
+ 0.0,
664
+ 1.0,
665
+ 0.0,
666
+ 0.0,
667
+ 0.0,
668
+ 0.0,
669
+ 1.0,
670
+ 0.0,
671
+ 0.0,
672
+ 0.0,
673
+ 0.0,
674
+ 1.0,
675
+ ],
676
+ ):
677
+ path = parent_prim.GetPath().AppendChild(self.clean_name(name))
678
+ group_prim = UsdGeom.Xform.Get(self._stage, path)
679
+ if not group_prim:
680
+ group_prim = UsdGeom.Xform.Define(self._stage, path)
681
+ # At present, the group transforms have been cooked into the vertices so this is not needed
682
+ matrix_op = group_prim.AddXformOp(
683
+ UsdGeom.XformOp.TypeTransform, UsdGeom.XformOp.PrecisionDouble
684
+ )
685
+ matrix_op.Set(Gf.Matrix4d(*matrix).GetTranspose())
686
+ # Map kinds
687
+ kind = Kind.Tokens.group
688
+ if obj_type == "ENS_CASE":
689
+ kind = Kind.Tokens.assembly
690
+ elif obj_type == "ENS_PART":
691
+ kind = Kind.Tokens.component
692
+ Usd.ModelAPI(group_prim).SetKind(kind)
693
+ logging.info(f"Created group:'{name}' {str(obj_type)}")
694
+ return group_prim
695
+
696
+ def uploadMaterial(self):
697
+ uriPath = self._destinationPath + "/Materials"
698
+ shutil.rmtree(uriPath, ignore_errors=True, onerror=None)
699
+ fullpath = os.path.join(os.path.dirname(__file__), "resources", "Materials")
700
+ shutil.copytree(fullpath, uriPath)
701
+
702
+ # Create a dome light in the scene.
703
+ def createDomeLight(self, texturePath):
704
+ newLight = UsdLux.DomeLight.Define(self._stage, "/Root/DomeLight")
705
+ newLight.CreateIntensityAttr(2200.0)
706
+ newLight.CreateTextureFileAttr(texturePath)
707
+ newLight.CreateTextureFormatAttr("latlong")
708
+
709
+ # Set rotation on domelight
710
+ xForm = newLight
711
+ rotateOp = xForm.AddXformOp(UsdGeom.XformOp.TypeRotateZYX, UsdGeom.XformOp.PrecisionFloat)
712
+ rotateOp.Set(Gf.Vec3f(270, 0, 0))
713
+
714
+
715
+ class OmniverseUpdateHandler(UpdateHandler):
716
+ """
717
+ Implement the Omniverse glue to a DSGSession instance
718
+ """
719
+
720
+ def __init__(self, omni: OmniverseWrapper):
721
+ super().__init__()
722
+ self._omni = omni
723
+ self._group_prims: Dict[int, Any] = dict()
724
+ self._root_prim = None
725
+ self._sent_textures = False
726
+
727
+ def add_group(self, id: int, view: bool = False) -> None:
728
+ super().add_group(id, view)
729
+ group = self.session.groups[id]
730
+ if not view:
731
+ parent_prim = self._group_prims[group.parent_id]
732
+ obj_type = self.get_dsg_cmd_attribute(group, "ENS_OBJ_TYPE")
733
+ matrix = self.group_matrix(group)
734
+ prim = self._omni.create_dsg_group(
735
+ group.name, parent_prim, matrix=matrix, obj_type=obj_type
736
+ )
737
+ self._group_prims[id] = prim
738
+ else:
739
+ # Map a view command into a new Omniverse stage and populate it with materials/lights.
740
+ # Create a new root stage in Omniverse
741
+
742
+ # Create or update the root group/camera
743
+ if not self.session.vrmode:
744
+ self._omni.update_camera(camera=group)
745
+
746
+ # record
747
+ self._group_prims[id] = self._root_prim
748
+
749
+ if self._omni._stage is not None:
750
+ self._omni._stage.SetStartTimeCode(
751
+ self.session.time_limits[0] * self._omni._time_codes_per_second
752
+ )
753
+ self._omni._stage.SetEndTimeCode(
754
+ self.session.time_limits[1] * self._omni._time_codes_per_second
755
+ )
756
+ self._omni._stage.SetTimeCodesPerSecond(self._omni._time_codes_per_second)
757
+
758
+ # Send the variable textures. Safe to do so once the first view is processed.
759
+ if not self._sent_textures:
760
+ self._omni.create_dsg_variable_textures(self.session.variables)
761
+ self._sent_textures = True
762
+
763
+ def add_variable(self, id: int) -> None:
764
+ super().add_variable(id)
765
+
766
+ def finalize_part(self, part: Part) -> None:
767
+ # generate an Omniverse compliant mesh from the Part
768
+ if part is None or part.cmd is None:
769
+ return
770
+ parent_prim = self._group_prims[part.cmd.parent_id]
771
+ obj_id = self.session.mesh_block_count
772
+ matrix = part.cmd.matrix4x4
773
+ name = part.cmd.name
774
+ color = [
775
+ part.cmd.fill_color[0] * part.cmd.diffuse,
776
+ part.cmd.fill_color[1] * part.cmd.diffuse,
777
+ part.cmd.fill_color[2] * part.cmd.diffuse,
778
+ part.cmd.fill_color[3],
779
+ ]
780
+
781
+ mat_info = part.material()
782
+ if part.cmd.render == part.cmd.CONNECTIVITY:
783
+ has_triangles = False
784
+ command, verts, conn, normals, tcoords, var_cmd = part.nodal_surface_rep()
785
+ if command is not None:
786
+ has_triangles = True
787
+ # Generate the mesh block
788
+ _ = self._omni.create_dsg_mesh_block(
789
+ name,
790
+ obj_id,
791
+ part.hash,
792
+ parent_prim,
793
+ verts,
794
+ conn,
795
+ normals,
796
+ tcoords,
797
+ matrix=matrix,
798
+ diffuse=color,
799
+ variable=var_cmd,
800
+ timeline=self.session.cur_timeline,
801
+ first_timestep=(self.session.cur_timeline[0] == self.session.time_limits[0]),
802
+ mat_info=mat_info,
803
+ )
804
+ if self._omni.use_lines:
805
+ command, verts, tcoords, var_cmd = part.line_rep()
806
+ if command is not None:
807
+ # If there are no triangle (ideally if these are not hidden line
808
+ # edges), then use the base color for the part. If there are
809
+ # triangles, then assume these are hidden line edges and use the
810
+ # line_color.
811
+ line_color = color
812
+ if has_triangles:
813
+ line_color = [
814
+ part.cmd.line_color[0] * part.cmd.diffuse,
815
+ part.cmd.line_color[1] * part.cmd.diffuse,
816
+ part.cmd.line_color[2] * part.cmd.diffuse,
817
+ part.cmd.line_color[3],
818
+ ]
819
+ # TODO: texture coordinates on lines are current invalid in OV
820
+ var_cmd = None
821
+ tcoords = None
822
+ # Generate the lines
823
+ _ = self._omni.create_dsg_lines(
824
+ name,
825
+ obj_id,
826
+ part.hash,
827
+ parent_prim,
828
+ verts,
829
+ tcoords,
830
+ matrix=matrix,
831
+ diffuse=line_color,
832
+ variable=var_cmd,
833
+ timeline=self.session.cur_timeline,
834
+ first_timestep=(
835
+ self.session.cur_timeline[0] == self.session.time_limits[0]
836
+ ),
837
+ )
838
+
839
+ elif part.cmd.render == part.cmd.NODES:
840
+ command, verts, sizes, colors, var_cmd = part.point_rep()
841
+ if command is not None:
842
+ _ = self._omni.create_dsg_points(
843
+ name,
844
+ obj_id,
845
+ part.hash,
846
+ parent_prim,
847
+ verts,
848
+ sizes,
849
+ colors,
850
+ matrix=matrix,
851
+ default_size=part.cmd.node_size_default,
852
+ default_color=color,
853
+ timeline=self.session.cur_timeline,
854
+ first_timestep=(self.session.cur_timeline[0] == self.session.time_limits[0]),
855
+ )
856
+ super().finalize_part(part)
857
+
858
+ def start_connection(self) -> None:
859
+ super().start_connection()
860
+
861
+ def end_connection(self) -> None:
862
+ super().end_connection()
863
+
864
+ def begin_update(self) -> None:
865
+ super().begin_update()
866
+ # restart the name tables
867
+ self._omni.clear_cleaned_names()
868
+ # clear the group Omni prims list
869
+ self._group_prims = dict()
870
+
871
+ self._omni.create_new_stage()
872
+ self._root_prim = self._omni.create_dsg_root()
873
+ # Create a distance and dome light in the scene
874
+ self._omni.createDomeLight("./Materials/000_sky.exr")
875
+ # Upload a material to the Omniverse server
876
+ self._omni.uploadMaterial()
877
+ self._sent_textures = False
878
+
879
+ def end_update(self) -> None:
880
+ super().end_update()
881
+ # Stage update complete
882
+ self._omni.save_stage()