Powertools for syncing Sun Study with CesiumGeoreference

exts/cesium.powertools/cesium/powertools/georefhelper/georef_helper_window.py

@@ -4,6 +4,7 @@
 from .proj import epsg_to_ecef, epsg_to_wgs84, get_crs_name_from_epsg
 import math
 from .custom_fields import string_field_with_label, int_field_with_label, float_field_with_label
+from pxr import Sdf
 
 
 class CesiumGeorefHelperWindow(ui.Window):
@@ -28,7 +29,7 @@ def __del__(self):
 
     @staticmethod
     def create_window():
-        return CesiumGeorefHelperWindow(width=250, height=550)
+        return CesiumGeorefHelperWindow(width=250, height=600)
 
     def _convert_coordinates(self):
         # Get the CRS and check if it is valid, adjust UI values accordingly
@@ -86,6 +87,26 @@ def _set_georeference_prim(self):
             self._wgs84_height_model.get_value_as_float()
         )
 
+    @staticmethod
+    def set_georef_from_environment():
+        stage = omni.usd.get_context().get_stage()
+
+        environment_prim = stage.GetPrimAtPath("/Environment")
+        cesium_prim = stage.GetPrimAtPath("/CesiumGeoreference")
+
+        lat_attr = environment_prim.GetAttribute("location:latitude")
+        long_attr = environment_prim.GetAttribute("location:longitude")
+
+        if lat_attr and long_attr:
+            cesium_prim.GetAttribute("cesium:georeferenceOrigin:latitude").Set(lat_attr.Get())
+            cesium_prim.GetAttribute("cesium:georeferenceOrigin:longitude").Set(long_attr.Get())
+            cesium_prim.GetAttribute("cesium:georeferenceOrigin:height").Set(0.0)
+        else:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                "Cannot set CesiumGeoreference as environment prim does not have latitude or longitude attributes"
+            )
+
     def _build_fn(self):
         """Builds out the UI buttons and their handlers."""
 
@@ -147,7 +168,10 @@ def on_coordinate_update(event):
             float_field_with_label("Y", model=self._ecef_y_model, enabled=False)
             float_field_with_label("Z", model=self._ecef_z_model, enabled=False)
 
-            ui.Button("Set Georeference", height=20, clicked_fn=self._set_georeference_prim)
+            ui.Button("Set Georeference from EPSG", height=20, clicked_fn=self._set_georeference_prim)
+            ui.Button(
+                "Set Georeference from Environment Prim", height=20, clicked_fn=self.set_georef_from_environment
+            )
 
             # Do the first conversion
             self._convert_coordinates()

exts/cesium.powertools/cesium/powertools/powertools_window.py

@@ -3,7 +3,7 @@
 from typing import Callable, Optional, List
 from cesium.omniverse.ui import CesiumOmniverseDebugWindow
 from .georefhelper.georef_helper_window import CesiumGeorefHelperWindow
-from .utils import extend_far_plane, save_carb_settings
+from .utils import extend_far_plane, save_carb_settings, set_sunstudy_from_georef
 import os
 from functools import partial
 
@@ -44,6 +44,7 @@ def __init__(self, **kwargs):
             PowertoolsAction("Open Cesium Georeference Helper Window", CesiumGeorefHelperWindow.create_window),
             PowertoolsAction("Extend Far Plane", extend_far_plane),
             PowertoolsAction("Save Carb Settings", partial(save_carb_settings, powertools_extension_location)),
+            PowertoolsAction("Set Sun Study from Georef", set_sunstudy_from_georef),
         ]
 
         self.frame.set_build_fn(self._build_fn)

exts/cesium.powertools/cesium/powertools/utils/utils.py

@@ -1,6 +1,6 @@
 import omni.usd
 from omni.kit.viewport.utility import get_active_viewport
-from pxr import Gf, UsdGeom
+from pxr import Gf, UsdGeom, Sdf
 import json
 import carb.settings
 import os
@@ -47,3 +47,27 @@ def save_carb_settings(powertools_extension_location: str):
     carb_settings_path = os.path.join(powertools_extension_location, "carb_settings.txt")
     with open(carb_settings_path, "w") as fh:
         fh.write(json.dumps(carb.settings.get_settings().get("/"), indent=2))
+
+
+# Helper function to search for an attribute on a prim, or create it if not present
+def get_or_create_attribute(prim, name, type):
+    attribute = prim.GetAttribute(name)
+    if not attribute:
+        attribute = prim.CreateAttribute(name, type)
+    return attribute
+
+
+def set_sunstudy_from_georef():
+    stage = omni.usd.get_context().get_stage()
+
+    environment_prim = stage.GetPrimAtPath("/Environment")
+    cesium_prim = stage.GetPrimAtPath("/CesiumGeoreference")
+
+    lat_attr = get_or_create_attribute(environment_prim, "location:latitude", Sdf.ValueTypeNames.Float)
+    lat_attr.Set(cesium_prim.GetAttribute("cesium:georeferenceOrigin:latitude").Get())
+
+    long_attr = get_or_create_attribute(environment_prim, "location:longitude", Sdf.ValueTypeNames.Float)
+    long_attr.Set(cesium_prim.GetAttribute("cesium:georeferenceOrigin:longitude").Get())
+
+    north_attr = get_or_create_attribute(environment_prim, "location:north_orientation", Sdf.ValueTypeNames.Float)
+    north_attr.Set(90.0)  # Always set to 90, otherwise the sun is at the wrong angle