From 476dece8e15c841d4a246e19df36b0d51597790c Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Sat, 20 Apr 2024 04:17:53 +0800
Subject: [PATCH 01/10] feat(omni-nerf-viewport): Connect to RPyC backend if
 Python version supported

---
 .../omni/nerf/viewport/extension.py           | 65 +++++++++++++++----
 1 file changed, 54 insertions(+), 11 deletions(-)

diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
index 402e63a..8bea0cb 100644
--- a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
+++ b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
@@ -1,7 +1,11 @@
+import platform
+
+import cv2
 import numpy as np
 import omni.ext
 import omni.ui as ui
 import omni.usd
+import rpyc
 from omni.kit.viewport.utility import get_active_viewport
 from pxr import Usd, UsdGeom
 
@@ -16,10 +20,17 @@ def some_public_function(x: int):
 # instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
 # on_shutdown() is called.
 class OmniNerfViewportExtension(omni.ext.IExt):
+
+    def __init__(self):
+        super().__init__()
+        self.is_python_supported: bool = platform.python_version().startswith("3.10")
+        """The Python version must match the backend version for RPyC to work."""
+
     # ext_id is current extension id. It can be used with extension manager to query additional information, like where
     # this extension is located on filesystem.
     def on_startup(self, ext_id):
-        # To see the Python print output, open the `Script Editor`.
+        # To see the Python print output in Omniverse Code, open the `Script Editor`.
+        # In Isaac Sim, see the startup console instead.
         print("[omni.nerf.viewport] omni nerf viewport startup")
         self.selected_camera_path = None
         # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/usd/stage/get-current-stage.html
@@ -45,22 +56,42 @@ def on_startup(self, ext_id):
         # TODO: Consider subscribing to update events
         # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/events.html#subscribe-to-update-events
         # Allocate memory
-        self.rgba_w, self.rgba_h = 256, 256
-        self.rgba = np.ones((self.rgba_w, self.rgba_h, 4), dtype=np.uint8) * 128
+        self.rgba_w, self.rgba_h = 640, 480
+        self.rgba = np.ones((self.rgba_h, self.rgba_w, 4), dtype=np.uint8) * 128
+        """RGBA image buffer. The shape is (H, W, 4), following the NumPy convention."""
         self.rgba[:,:,3] = 255
+        # Init RPyC connection
+        if self.is_python_supported:
+            self.init_rpyc()
         # Build UI
         self.build_ui()
 
+    def init_rpyc(self):
+        # TODO: Make the following configurable
+        host = 'localhost'
+        port = 7007
+        model_config_path = '/workspace/outputs/poster/nerfacto/DATE_TIME/config.yml'
+        model_checkpoint_path = '/workspace/outputs/poster/nerfacto/DATE_TIME/nerfstudio_models/CHECKPOINT_NAME.ckpt'
+        device = 'cuda'
+        self.rpyc_conn = rpyc.classic.connect(host, port)
+        self.rpyc_conn.execute('from nerfstudio_renderer import NerfStudioRenderQueue')
+        self.rpyc_conn.execute('from pathlib import Path')
+        self.rpyc_conn.execute('import torch')
+        self.rpyc_conn.execute(f'rq = NerfStudioRenderQueue(model_config_path=Path("{model_config_path}"), checkpoint_path="{model_checkpoint_path}", device=torch.device("{device}"))')
+
     def build_ui(self):
         """Build the UI. Should be called upon startup."""
         # Please refer to the `Omni::UI Doc` tab in Omniverse Code for efficient development.
         # Ref: https://youtu.be/j1Pwi1KRkhk
         # Ref: https://github.com/NVIDIA-Omniverse
         # Ref: https://youtu.be/dNLFpVhBrGs
-        self.ui_window = ui.Window("NeRF Viewport")
+        self.ui_window = ui.Window("NeRF Viewport", width=self.rgba_w, height=self.rgba_h)
 
         with self.ui_window.frame:
-            with ui.VStack(height=0):
+            with ui.VStack():
+                self.ui_lbl = ui.Label("(To Be Updated)")
+                state = "supported" if platform.python_version().startswith("3.10") else "NOT supported"
+                self.ui_lbl.text = f"Python {platform.python_version()} is {state}"
                 # Camera Viewport
                 # Ref: https://docs.omniverse.nvidia.com/kit/docs/omni.kit.viewport.docs/latest/overview.html#simplest-example
                 # Don't create a new viewport widget as below, since the viewport widget will often flicker.
@@ -89,8 +120,8 @@ def build_ui(self):
                 # TODO: Potentially optimize with `set_bytes_data_from_gpu`
                 self.ui_nerf_img = ui.ImageWithProvider(
                     self.ui_nerf_provider,
-                    width=self.rgba_w,
-                    height=self.rgba_h,
+                    width=ui.Percent(100),
+                    height=ui.Percent(100),
                 )
                 # TODO: Larger image size?
                 # TODO: Get viewport data and show it
@@ -98,13 +129,12 @@ def build_ui(self):
                 # TODO: Get viewport matrices and show it
                 print("Viewport Projection", self.viewport_api.projection)
                 print("Viewport Transform", self.viewport_api.transform)
-                self.ui_lbl = ui.Label("(To Be Updated)")
         self.update_ui()
 
     def update_ui(self):
         print("[omni.nerf.viewport] Updating UI")
         print(f"[omni.nerf.viewport] Selected Camera: {self.selected_camera_path}")
-        self.ui_lbl.text = f"Selected Camera: {self.selected_camera_path}"
+        # self.ui_lbl.text = f"Selected Camera: {self.selected_camera_path}"
         # Ref: https://forums.developer.nvidia.com/t/refresh-window-ui/221200
         self.ui_window.frame.rebuild()
 
@@ -137,12 +167,25 @@ def _on_stage_event(self, event):
     def _on_rendering_event(self, event):
         """Called by rendering_event_stream."""
         # No need to check event type, since there is only one event type: `NEW_FRAME`.
-        # TODO: Below color change is for testing purposes.
-        self.rgba[:,:,:3] = (self.rgba[:,:,:3] + np.ones((self.rgba_w, self.rgba_h, 3), dtype=np.uint8)) % 256
+        if self.is_python_supported:
+            # TODO: Use viewport transform.
+            camera_position = [0, 0, 0]
+            camera_rotation = [0, 0, 0]
+            self.rpyc_conn.execute(f'rq.update_camera({camera_position}, {camera_rotation})')
+            image = self.rpyc_conn.eval('rq.get_rgb_image()')
+            if image is not None:
+                image = np.array(image) # received with shape (H*, W*, 3)
+                image = cv2.resize(image, (self.rgba_w, self.rgba_h), interpolation=cv2.INTER_LINEAR) # resize to (H, W, 3)
+                self.rgba[:,:,:3] = image * 255
+        else:
+            # If python version is not supported, render the dummy image.
+            self.rgba[:,:,:3] = (self.rgba[:,:,:3] + np.ones((self.rgba_h, self.rgba_w, 3), dtype=np.uint8)) % 256
         self.ui_nerf_provider.set_bytes_data(self.rgba.flatten().tolist(), (self.rgba_w, self.rgba_h))
 
     def on_shutdown(self):
         print("[omni.nerf.viewport] omni nerf viewport shutdown")
+        if self.is_python_supported:
+            self.rpyc_conn.execute('del rq')
 
     def destroy(self):
         # Ref: https://docs.omniverse.nvidia.com/workflows/latest/extensions/object_info.html#step-3-4-use-usdcontext-to-listen-for-selection-changes

From 92e2b0ed2e0cfa101c5467b46d2bb1e458c986a4 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Sat, 20 Apr 2024 23:27:39 +0800
Subject: [PATCH 02/10] feat(omni-nerf-viewport): Update NeRF request based on
 the camera pose of the active viewport

---
 .../omni/nerf/viewport/extension.py           | 56 ++++++++++++++-----
 1 file changed, 41 insertions(+), 15 deletions(-)

diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
index 8bea0cb..592b94d 100644
--- a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
+++ b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
@@ -7,7 +7,7 @@
 import omni.usd
 import rpyc
 from omni.kit.viewport.utility import get_active_viewport
-from pxr import Usd, UsdGeom
+from pxr import Gf, Usd, UsdGeom
 
 
 # Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)`
@@ -25,6 +25,8 @@ def __init__(self):
         super().__init__()
         self.is_python_supported: bool = platform.python_version().startswith("3.10")
         """The Python version must match the backend version for RPyC to work."""
+        self.camera_position: Gf.Vec3d = None
+        self.camera_rotation: Gf.Vec3d = None
 
     # ext_id is current extension id. It can be used with extension manager to query additional information, like where
     # this extension is located on filesystem.
@@ -69,7 +71,7 @@ def on_startup(self, ext_id):
     def init_rpyc(self):
         # TODO: Make the following configurable
         host = 'localhost'
-        port = 7007
+        port = 10001
         model_config_path = '/workspace/outputs/poster/nerfacto/DATE_TIME/config.yml'
         model_checkpoint_path = '/workspace/outputs/poster/nerfacto/DATE_TIME/nerfstudio_models/CHECKPOINT_NAME.ckpt'
         device = 'cuda'
@@ -92,6 +94,7 @@ def build_ui(self):
                 self.ui_lbl = ui.Label("(To Be Updated)")
                 state = "supported" if platform.python_version().startswith("3.10") else "NOT supported"
                 self.ui_lbl.text = f"Python {platform.python_version()} is {state}"
+                self.ui_btn = ui.Button("Reset Camera", width=20, clicked_fn=self.on_btn_click)
                 # Camera Viewport
                 # Ref: https://docs.omniverse.nvidia.com/kit/docs/omni.kit.viewport.docs/latest/overview.html#simplest-example
                 # Don't create a new viewport widget as below, since the viewport widget will often flicker.
@@ -106,10 +109,7 @@ def build_ui(self):
                 # self.viewport_api = self.ui_viewport_widget.viewport_api
                 # ````
                 # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/python-snippets/viewport/change-viewport-active-camera.html
-                self.viewport_api = get_active_viewport()
-                # We chose to use Viewport instead of Isaac Sim's Camera Sensor to avoid dependency on Isaac Sim.
-                # We want the extension to work with any Omniverse app, not just Isaac Sim.
-                # Ref: https://docs.omniverse.nvidia.com/isaacsim/latest/features/sensors_simulation/isaac_sim_sensors_camera.html
+                # Instead, the viewport is obtained from the active viewport in new renderings.
 
                 # NeRF Viewport
                 # Examples on using ByteImageProvider can be found by installing Isaac Sim
@@ -124,11 +124,6 @@ def build_ui(self):
                     height=ui.Percent(100),
                 )
                 # TODO: Larger image size?
-                # TODO: Get viewport data and show it
-                # Ref: https://forums.developer.nvidia.com/t/how-can-i-grab-the-viewport-or-the-camera-rendering-in-a-python-script/238365/2
-                # TODO: Get viewport matrices and show it
-                print("Viewport Projection", self.viewport_api.projection)
-                print("Viewport Transform", self.viewport_api.transform)
         self.update_ui()
 
     def update_ui(self):
@@ -138,6 +133,18 @@ def update_ui(self):
         # Ref: https://forums.developer.nvidia.com/t/refresh-window-ui/221200
         self.ui_window.frame.rebuild()
 
+    def on_btn_click(self):
+        # TODO: Allow resetting the camera to a specific position
+        # Below doesn't seem to work
+        # stage: Usd.Stage = self.usd_context.get_stage()
+        # prim: Usd.Prim = stage.GetPrimAtPath('/OmniverseKit_Persp')
+        # # `UsdGeom.Xformable(prim).SetTranslateOp` doesn't seem to exist
+        # prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(0, 0, 0.1722))
+        # prim.GetAttribute("xformOp:rotateXYZ").Set(Gf.Vec3d(0, -152, 0))
+        # print("translateOp", prim.GetAttribute("xformOp:translate").Get())
+        # print("rotateXYZOp", prim.GetAttribute("xformOp:rotateXYZ").Get())
+        print("[omni.nerf.viewport] (TODO) Reset Camera")
+
     def _get_selected_camera_path(self):
         """Get the selected camera prim. Return None if no camera is selected or the first selected prim isn't a camera."""
         # Ref: https://docs.omniverse.nvidia.com/workflows/latest/extensions/object_info.html#step-5-get-the-selected-prims-data
@@ -168,12 +175,31 @@ def _on_rendering_event(self, event):
         """Called by rendering_event_stream."""
         # No need to check event type, since there is only one event type: `NEW_FRAME`.
         if self.is_python_supported:
-            # TODO: Use viewport transform.
-            camera_position = [0, 0, 0]
-            camera_rotation = [0, 0, 0]
-            self.rpyc_conn.execute(f'rq.update_camera({camera_position}, {camera_rotation})')
+            viewport_api = get_active_viewport()
+            # We chose to use Viewport instead of Isaac Sim's Camera Sensor to avoid dependency on Isaac Sim.
+            # We want the extension to work with any Omniverse app, not just Isaac Sim.
+            # Ref: https://docs.omniverse.nvidia.com/isaacsim/latest/features/sensors_simulation/isaac_sim_sensors_camera.html
+            camera_mat: Gf.Matrix4d = viewport_api.transform
+            # TODO: Use viewport camera projection matrix `viewport_api.projection`?
+            camera_position: Gf.Vec3d = camera_mat.ExtractTranslation()
+            camera_rotation: Gf.Vec3d = camera_mat.ExtractRotation().Decompose(*Gf.Matrix3d())
+            # Not same as below due to the potential difference in rotation matrix representation
+            # ```
+            # from scipy.spatial.transform import Rotation as R
+            # camera_rotation: Gf.Vec3d = R.from_matrix(camera_mat.ExtractRotationMatrix()).as_euler('xyz', degrees=True) # in degrees
+            # ```
+            # TODO: Consider object transform (if it is moved or rotated)
+            # No need to transform from Isaac Sim space to Nerfstudio space, since they are both in the same space.
+            # Ref: https://github.com/j3soon/coordinate-system-conventions
+            if camera_position != self.camera_position or camera_rotation != self.camera_rotation:
+                self.camera_position = camera_position
+                self.camera_rotation = camera_rotation
+                print("[omni.nerf.viewport] New camera position:", camera_position)
+                print("[omni.nerf.viewport] New camera rotation:", camera_rotation)
+                self.rpyc_conn.execute(f'rq.update_camera({list(camera_position)}, {list(np.deg2rad(camera_rotation))})')
             image = self.rpyc_conn.eval('rq.get_rgb_image()')
             if image is not None:
+                print("[omni.nerf.viewport] NeRF viewport updated")
                 image = np.array(image) # received with shape (H*, W*, 3)
                 image = cv2.resize(image, (self.rgba_w, self.rgba_h), interpolation=cv2.INTER_LINEAR) # resize to (H, W, 3)
                 self.rgba[:,:,:3] = image * 255

From c1f7808d7fd793292be5a2935187f9da81d45b4c Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Mon, 22 Apr 2024 21:55:33 +0800
Subject: [PATCH 03/10] fix: Fix rotation and FoV misalignment

---
 .../omni/nerf/viewport/extension.py           | 31 ++++++++++++-------
 .../src/nerfstudio_renderer/render_queue.py   | 29 +++++++++++++++--
 nerfstudio_renderer/tests/pygame_test.py      |  2 +-
 3 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
index 592b94d..db85ece 100644
--- a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
+++ b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
@@ -58,7 +58,7 @@ def on_startup(self, ext_id):
         # TODO: Consider subscribing to update events
         # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/events.html#subscribe-to-update-events
         # Allocate memory
-        self.rgba_w, self.rgba_h = 640, 480
+        self.rgba_w, self.rgba_h = 640, 360 # Follow default camera resolution 1280x720
         self.rgba = np.ones((self.rgba_h, self.rgba_w, 4), dtype=np.uint8) * 128
         """RGBA image buffer. The shape is (H, W, 4), following the NumPy convention."""
         self.rgba[:,:,3] = 255
@@ -90,11 +90,7 @@ def build_ui(self):
         self.ui_window = ui.Window("NeRF Viewport", width=self.rgba_w, height=self.rgba_h)
 
         with self.ui_window.frame:
-            with ui.VStack():
-                self.ui_lbl = ui.Label("(To Be Updated)")
-                state = "supported" if platform.python_version().startswith("3.10") else "NOT supported"
-                self.ui_lbl.text = f"Python {platform.python_version()} is {state}"
-                self.ui_btn = ui.Button("Reset Camera", width=20, clicked_fn=self.on_btn_click)
+            with ui.ZStack():
                 # Camera Viewport
                 # Ref: https://docs.omniverse.nvidia.com/kit/docs/omni.kit.viewport.docs/latest/overview.html#simplest-example
                 # Don't create a new viewport widget as below, since the viewport widget will often flicker.
@@ -102,9 +98,9 @@ def build_ui(self):
                 # ```
                 # from omni.kit.widget.viewport import ViewportWidget
                 # self.ui_viewport_widget = ViewportWidget(
-                #     resolution = (640, 480),
+                #     resolution = (640, 360),
                 #     width = 640,
-                #     height = 480,
+                #     height = 360,
                 # )
                 # self.viewport_api = self.ui_viewport_widget.viewport_api
                 # ````
@@ -124,6 +120,11 @@ def build_ui(self):
                     height=ui.Percent(100),
                 )
                 # TODO: Larger image size?
+                with ui.VStack(height=0):
+                    self.ui_lbl = ui.Label("(To Be Updated)")
+                    state = "supported" if platform.python_version().startswith("3.10") else "NOT supported"
+                    self.ui_lbl.text = f"Python {platform.python_version()} is {state}"
+                    self.ui_btn = ui.Button("Reset Camera", width=20, clicked_fn=self.on_btn_click)
         self.update_ui()
 
     def update_ui(self):
@@ -179,10 +180,16 @@ def _on_rendering_event(self, event):
             # We chose to use Viewport instead of Isaac Sim's Camera Sensor to avoid dependency on Isaac Sim.
             # We want the extension to work with any Omniverse app, not just Isaac Sim.
             # Ref: https://docs.omniverse.nvidia.com/isaacsim/latest/features/sensors_simulation/isaac_sim_sensors_camera.html
-            camera_mat: Gf.Matrix4d = viewport_api.transform
-            # TODO: Use viewport camera projection matrix `viewport_api.projection`?
-            camera_position: Gf.Vec3d = camera_mat.ExtractTranslation()
-            camera_rotation: Gf.Vec3d = camera_mat.ExtractRotation().Decompose(*Gf.Matrix3d())
+            camera_to_world_mat: Gf.Matrix4d = viewport_api.transform
+            camera_position: Gf.Vec3d = camera_to_world_mat.ExtractTranslation()
+            # I suspect that the `Decompose` function will extract the rotation in the order of the input axes.
+            # So for EulerXYZ, we want to first extract and remove the Z rotation, then Y, then X.
+            # Then we reverse the order to get the XYZ rotation.
+            # I haven't spend time looking into the source code to confirm this hypothesis though.
+            # Ref: https://forums.developer.nvidia.com/t/how-to-get-euler-angle-of-the-prim-through-script-with-script-editor/269704/3
+            # Ref: https://github.com/PixarAnimationStudios/OpenUSD/blob/2864f3d04f396432f22ec5d6928fc37d34bb4c90/pxr/base/gf/rotation.cpp#L108
+            camera_rotation: Gf.Vec3d = Gf.Vec3d(*reversed(camera_to_world_mat.ExtractRotation().Decompose(*reversed(Gf.Matrix3d()))))
+            # TODO: Consider using viewport camera projection matrix `viewport_api.projection`?
             # Not same as below due to the potential difference in rotation matrix representation
             # ```
             # from scipy.spatial.transform import Rotation as R
diff --git a/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py b/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
index e3df43a..687de31 100644
--- a/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
+++ b/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
@@ -43,10 +43,33 @@ def default_config():
             A default config.
         """
         # These configurations are chosen empirically, and may be subject to change.
+        # Nerfstudio camera defaults:
+        # - vertical FoV: 50 degrees
+        # Isaac Sim camera defaults:
+        # - Size: 1280x720
+        # - Focal Length: 18.14756
+        # - Horizontal Aperture: 20.955
+        # - Vertical Aperture: (Value Unused)
+        # - (Calculated) horizontal FoV = math.degrees(2 * math.atan(20.955 / (2 * 18.14756))) = 60
+        # The following vertical FoV is chosen to follow the Isaac Sim camera defaults.
+        # Some useful equations:
+        # - focal_length = width / (2 * math.tan(math.radians(fov_horizontal) / 2))
+        # - focal_length = height / (2 * math.tan(math.radians(fov_vertical) / 2))
+        # - fov_vertical = math.degrees(2 * math.atan(height / (2 * focal_length)))
+        # - fov_horizontal = math.degrees(2 * math.atan(width / (2 * focal_length)))
+        # - fov_horizontal = math.degrees(2 * math.atan(horiz_aperture / (2 * focal_length)))
+        #   Ref: https://forums.developer.nvidia.com/t/change-intrinsic-camera-parameters/180309/6
+        # - aspect_ratio = width / height
+        # - fov_vertical = math.degrees(2 * math.atan((height / width) * math.tan(math.radians(fov_horizontal) / 2)))
         return RendererCameraConfig([
-            { 'width': 90,  'height': 42,  'fov': 50 },
-            { 'width': 180, 'height': 84,  'fov': 50 },
-            { 'width': 450, 'height': 210, 'fov': 50 },
+            # fov_vertical = math.degrees(2 * math.atan((height / width) * math.tan(math.radians(fov_horizontal) / 2)))
+            # = 35.98339777135764
+            # 0.05x resolution
+            { 'width': 64,  'height': 36,  'fov': 35.98339777135764 },
+            # 0.1x resolution
+            { 'width': 128, 'height': 72,  'fov': 35.98339777135764 },
+            # 0.25x resolution
+            { 'width': 320, 'height': 180, 'fov': 35.98339777135764 },
         ])
 
     def load_config(file_path=None):
diff --git a/nerfstudio_renderer/tests/pygame_test.py b/nerfstudio_renderer/tests/pygame_test.py
index 8641723..208da19 100644
--- a/nerfstudio_renderer/tests/pygame_test.py
+++ b/nerfstudio_renderer/tests/pygame_test.py
@@ -44,7 +44,7 @@ def main(args):
     pygame.init()
 
     # Set the width and height of the window
-    width, height = 640, 480
+    width, height = 640, 360
     window_size = (width, height)
 
     # Create a Pygame window

From d01fe2e772eb94b8d7e52b30813857a38c4aac05 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Mon, 22 Apr 2024 22:53:11 +0800
Subject: [PATCH 04/10] feat(omni-nerf-viewport): Consider NeRF mesh prim
 transform (including scale), Remove selected camera detection

---
 .../omni/nerf/viewport/extension.py           | 91 ++++++++++---------
 1 file changed, 46 insertions(+), 45 deletions(-)

diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
index db85ece..0e33dea 100644
--- a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
+++ b/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
@@ -34,18 +34,10 @@ def on_startup(self, ext_id):
         # To see the Python print output in Omniverse Code, open the `Script Editor`.
         # In Isaac Sim, see the startup console instead.
         print("[omni.nerf.viewport] omni nerf viewport startup")
-        self.selected_camera_path = None
         # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/usd/stage/get-current-stage.html
         self.usd_context = omni.usd.get_context()
         # Subscribe to event streams
         # Ref: https://docs.omniverse.nvidia.com/kit/docs/kit-manual/latest/guide/event_streams.html
-        # Listen to selection changes
-        # Ref: https://docs.omniverse.nvidia.com/workflows/latest/extensions/object_info.html#step-3-4-use-usdcontext-to-listen-for-selection-changes
-        self.stage_event_stream = self.usd_context.get_stage_event_stream()
-        self.stage_event_delegate = self.stage_event_stream.create_subscription_to_pop(
-            self._on_stage_event, name="Object Info Selection Update"
-        )
-        # TODO: Subscribe to only certain event types
         # Ref: https://docs.omniverse.nvidia.com/kit/docs/kit-manual/104.0/carb.events/carb.events.IEventStream.html#carb.events.IEventStream.create_subscription_to_pop_by_type
         # Listen to rendering events. Only triggered when the viewport is rendering is updated.
         # Will not be triggered when no viewport is visible on the screen.
@@ -121,20 +113,35 @@ def build_ui(self):
                 )
                 # TODO: Larger image size?
                 with ui.VStack(height=0):
-                    self.ui_lbl = ui.Label("(To Be Updated)")
+                    self.ui_lbl_py = ui.Label("(To Be Updated)")
                     state = "supported" if platform.python_version().startswith("3.10") else "NOT supported"
-                    self.ui_lbl.text = f"Python {platform.python_version()} is {state}"
-                    self.ui_btn = ui.Button("Reset Camera", width=20, clicked_fn=self.on_btn_click)
+                    self.ui_lbl_py.text = f"Python {platform.python_version()} is {state}"
+                    # UI for setting the NeRF mesh
+                    # Ref: https://docs.omniverse.nvidia.com/workflows/latest/extensions/scatter_tool.html
+                    with ui.HStack():
+                        self.ui_lbl_mesh = ui.Label("NeRF Mesh", width=65)
+                        # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/ui/widgets/stringfield.html
+                        self._mesh_prim_model = ui.SimpleStringModel()
+                        ui.StringField(model=self._mesh_prim_model)
+                        ui.Button(
+                            " S ",
+                            width=0,
+                            height=0,
+                            clicked_fn=self._on_btn_set_click,
+                            tooltip="Get From Selection",
+                        )
+                    ui.Button("Reset Camera", width=20, clicked_fn=self.on_btn_reset_click)
         self.update_ui()
 
     def update_ui(self):
         print("[omni.nerf.viewport] Updating UI")
-        print(f"[omni.nerf.viewport] Selected Camera: {self.selected_camera_path}")
-        # self.ui_lbl.text = f"Selected Camera: {self.selected_camera_path}"
         # Ref: https://forums.developer.nvidia.com/t/refresh-window-ui/221200
         self.ui_window.frame.rebuild()
 
-    def on_btn_click(self):
+    def _on_btn_set_click(self):
+        self._mesh_prim_model.as_string = self._get_selected_prim_path()
+
+    def on_btn_reset_click(self):
         # TODO: Allow resetting the camera to a specific position
         # Below doesn't seem to work
         # stage: Usd.Stage = self.usd_context.get_stage()
@@ -146,49 +153,43 @@ def on_btn_click(self):
         # print("rotateXYZOp", prim.GetAttribute("xformOp:rotateXYZ").Get())
         print("[omni.nerf.viewport] (TODO) Reset Camera")
 
-    def _get_selected_camera_path(self):
-        """Get the selected camera prim. Return None if no camera is selected or the first selected prim isn't a camera."""
+    def _get_selected_prim_path(self):
+        """Get the selected prim. Return '' if no prim is selected."""
         # Ref: https://docs.omniverse.nvidia.com/workflows/latest/extensions/object_info.html#step-5-get-the-selected-prims-data
         selected_prim_paths = self.usd_context.get_selection().get_selected_prim_paths()
         if not selected_prim_paths:
-            return None
-        stage: Usd.Stage = self.usd_context.get_stage()
-        selected_prim = stage.GetPrimAtPath(selected_prim_paths[0])
-        assert type(selected_prim) == Usd.Prim
-        if not selected_prim.IsA(UsdGeom.Camera):
-            return None
-        return selected_prim.GetPath()
-
-    def _on_stage_event(self, event):
-        """Called by stage_event_stream. We only care about selection changes."""
-        print("[omni.nerf.viewport] on_stage_event", omni.usd.StageEventType(event.type))
-        if event.type != int(omni.usd.StageEventType.SELECTION_CHANGED):
-            return
-        selected_camera_path = self._get_selected_camera_path()
-        if self.selected_camera_path == selected_camera_path:
-            # Skip if the selected camera hasn't changed
-            print("[omni.nerf.viewport] Skip updating UI")
-            return
-        self.selected_camera_path = selected_camera_path
-        self.update_ui()
+            return ''
+        return selected_prim_paths[0]
 
     def _on_rendering_event(self, event):
         """Called by rendering_event_stream."""
         # No need to check event type, since there is only one event type: `NEW_FRAME`.
-        if self.is_python_supported:
+        if self.is_python_supported and self._mesh_prim_model.as_string != '':
             viewport_api = get_active_viewport()
             # We chose to use Viewport instead of Isaac Sim's Camera Sensor to avoid dependency on Isaac Sim.
             # We want the extension to work with any Omniverse app, not just Isaac Sim.
             # Ref: https://docs.omniverse.nvidia.com/isaacsim/latest/features/sensors_simulation/isaac_sim_sensors_camera.html
             camera_to_world_mat: Gf.Matrix4d = viewport_api.transform
-            camera_position: Gf.Vec3d = camera_to_world_mat.ExtractTranslation()
+            object_to_world_mat: Gf.Matrix4d = Gf.Matrix4d()
+            if self._mesh_prim_model.as_string != '':
+                stage: Usd.Stage = self.usd_context.get_stage()
+                selected_prim: Usd.Prim = stage.GetPrimAtPath(self._mesh_prim_model.as_string)
+                selected_xform: UsdGeom.Xformable = UsdGeom.Xformable(selected_prim)
+                object_to_world_mat = selected_xform.GetLocalTransformation()
+            # In USD, pre-multiplication is used for matrices.
+            # Ref: https://openusd.org/dev/api/usd_geom_page_front.html#UsdGeom_LinAlgBasics
+            world_to_object_mat: Gf.Matrix4d = object_to_world_mat.GetInverse()
+            camera_to_object_mat: Gf.Matrix4d = camera_to_world_mat * world_to_object_mat
+            camera_to_object_pos: Gf.Vec3d = camera_to_object_mat.ExtractTranslation()
             # I suspect that the `Decompose` function will extract the rotation in the order of the input axes.
             # So for EulerXYZ, we want to first extract and remove the Z rotation, then Y, then X.
             # Then we reverse the order to get the XYZ rotation.
             # I haven't spend time looking into the source code to confirm this hypothesis though.
             # Ref: https://forums.developer.nvidia.com/t/how-to-get-euler-angle-of-the-prim-through-script-with-script-editor/269704/3
             # Ref: https://github.com/PixarAnimationStudios/OpenUSD/blob/2864f3d04f396432f22ec5d6928fc37d34bb4c90/pxr/base/gf/rotation.cpp#L108
-            camera_rotation: Gf.Vec3d = Gf.Vec3d(*reversed(camera_to_world_mat.ExtractRotation().Decompose(*reversed(Gf.Matrix3d()))))
+            # must remove scale before rotation
+            camera_to_object_mat.Orthonormalize()
+            camera_to_object_rot: Gf.Vec3d = Gf.Vec3d(*reversed(camera_to_object_mat.ExtractRotation().Decompose(*reversed(Gf.Matrix3d()))))
             # TODO: Consider using viewport camera projection matrix `viewport_api.projection`?
             # Not same as below due to the potential difference in rotation matrix representation
             # ```
@@ -198,12 +199,12 @@ def _on_rendering_event(self, event):
             # TODO: Consider object transform (if it is moved or rotated)
             # No need to transform from Isaac Sim space to Nerfstudio space, since they are both in the same space.
             # Ref: https://github.com/j3soon/coordinate-system-conventions
-            if camera_position != self.camera_position or camera_rotation != self.camera_rotation:
-                self.camera_position = camera_position
-                self.camera_rotation = camera_rotation
-                print("[omni.nerf.viewport] New camera position:", camera_position)
-                print("[omni.nerf.viewport] New camera rotation:", camera_rotation)
-                self.rpyc_conn.execute(f'rq.update_camera({list(camera_position)}, {list(np.deg2rad(camera_rotation))})')
+            if camera_to_object_pos != self.camera_position or camera_to_object_rot != self.camera_rotation:
+                self.camera_position = camera_to_object_pos
+                self.camera_rotation = camera_to_object_rot
+                print("[omni.nerf.viewport] New camera position:", camera_to_object_pos)
+                print("[omni.nerf.viewport] New camera rotation:", camera_to_object_rot)
+                self.rpyc_conn.execute(f'rq.update_camera({list(camera_to_object_pos)}, {list(np.deg2rad(camera_to_object_rot))})')
             image = self.rpyc_conn.eval('rq.get_rgb_image()')
             if image is not None:
                 print("[omni.nerf.viewport] NeRF viewport updated")

From 04cd1012455fc74e763201bfd8d4fe4ec0881e58 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Mon, 22 Apr 2024 23:08:49 +0800
Subject: [PATCH 05/10] refactor(omni-nerf-viewport): Move extension content to
 its dedicated directory

---
 .gitignore => extension/.gitignore                  |   0
 {.vscode => extension/.vscode}/extensions.json      |   0
 {.vscode => extension/.vscode}/launch.json          |   0
 {.vscode => extension/.vscode}/settings.json        |   0
 README.md => extension/README.md                    |   0
 .../exts}/omni.nerf.viewport/config/extension.toml  |   0
 .../exts}/omni.nerf.viewport/data/icon.png          | Bin
 .../exts}/omni.nerf.viewport/data/preview.png       | Bin
 .../exts}/omni.nerf.viewport/docs/CHANGELOG.md      |   0
 .../exts}/omni.nerf.viewport/docs/README.md         |   0
 .../exts}/omni.nerf.viewport/docs/index.rst         |   0
 .../omni/nerf/viewport/__init__.py                  |   0
 .../omni/nerf/viewport/extension.py                 |   0
 .../omni/nerf/viewport/tests/__init__.py            |   0
 .../omni/nerf/viewport/tests/test_hello_world.py    |   0
 link_app.bat => extension/link_app.bat              |   0
 link_app.sh => extension/link_app.sh                |   0
 .../tools}/packman/bootstrap/configure.bat          |   0
 .../packman/bootstrap/download_file_from_url.ps1    |   0
 .../bootstrap/fetch_file_from_packman_bootstrap.cmd |   0
 .../packman/bootstrap/generate_temp_file_name.ps1   |   0
 .../packman/bootstrap/generate_temp_folder.ps1      |   0
 .../tools}/packman/bootstrap/install_package.py     |   0
 .../tools}/packman/config.packman.xml               |   0
 {tools => extension/tools}/packman/packman          |   0
 {tools => extension/tools}/packman/packman.cmd      |   0
 {tools => extension/tools}/packman/python.bat       |   0
 {tools => extension/tools}/packman/python.sh        |   0
 {tools => extension/tools}/scripts/link_app.py      |   0
 29 files changed, 0 insertions(+), 0 deletions(-)
 rename .gitignore => extension/.gitignore (100%)
 rename {.vscode => extension/.vscode}/extensions.json (100%)
 rename {.vscode => extension/.vscode}/launch.json (100%)
 rename {.vscode => extension/.vscode}/settings.json (100%)
 rename README.md => extension/README.md (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/config/extension.toml (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/data/icon.png (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/data/preview.png (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/docs/CHANGELOG.md (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/docs/README.md (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/docs/index.rst (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/omni/nerf/viewport/__init__.py (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/omni/nerf/viewport/extension.py (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/omni/nerf/viewport/tests/__init__.py (100%)
 rename {exts => extension/exts}/omni.nerf.viewport/omni/nerf/viewport/tests/test_hello_world.py (100%)
 rename link_app.bat => extension/link_app.bat (100%)
 rename link_app.sh => extension/link_app.sh (100%)
 rename {tools => extension/tools}/packman/bootstrap/configure.bat (100%)
 rename {tools => extension/tools}/packman/bootstrap/download_file_from_url.ps1 (100%)
 rename {tools => extension/tools}/packman/bootstrap/fetch_file_from_packman_bootstrap.cmd (100%)
 rename {tools => extension/tools}/packman/bootstrap/generate_temp_file_name.ps1 (100%)
 rename {tools => extension/tools}/packman/bootstrap/generate_temp_folder.ps1 (100%)
 rename {tools => extension/tools}/packman/bootstrap/install_package.py (100%)
 rename {tools => extension/tools}/packman/config.packman.xml (100%)
 rename {tools => extension/tools}/packman/packman (100%)
 rename {tools => extension/tools}/packman/packman.cmd (100%)
 rename {tools => extension/tools}/packman/python.bat (100%)
 rename {tools => extension/tools}/packman/python.sh (100%)
 rename {tools => extension/tools}/scripts/link_app.py (100%)

diff --git a/.gitignore b/extension/.gitignore
similarity index 100%
rename from .gitignore
rename to extension/.gitignore
diff --git a/.vscode/extensions.json b/extension/.vscode/extensions.json
similarity index 100%
rename from .vscode/extensions.json
rename to extension/.vscode/extensions.json
diff --git a/.vscode/launch.json b/extension/.vscode/launch.json
similarity index 100%
rename from .vscode/launch.json
rename to extension/.vscode/launch.json
diff --git a/.vscode/settings.json b/extension/.vscode/settings.json
similarity index 100%
rename from .vscode/settings.json
rename to extension/.vscode/settings.json
diff --git a/README.md b/extension/README.md
similarity index 100%
rename from README.md
rename to extension/README.md
diff --git a/exts/omni.nerf.viewport/config/extension.toml b/extension/exts/omni.nerf.viewport/config/extension.toml
similarity index 100%
rename from exts/omni.nerf.viewport/config/extension.toml
rename to extension/exts/omni.nerf.viewport/config/extension.toml
diff --git a/exts/omni.nerf.viewport/data/icon.png b/extension/exts/omni.nerf.viewport/data/icon.png
similarity index 100%
rename from exts/omni.nerf.viewport/data/icon.png
rename to extension/exts/omni.nerf.viewport/data/icon.png
diff --git a/exts/omni.nerf.viewport/data/preview.png b/extension/exts/omni.nerf.viewport/data/preview.png
similarity index 100%
rename from exts/omni.nerf.viewport/data/preview.png
rename to extension/exts/omni.nerf.viewport/data/preview.png
diff --git a/exts/omni.nerf.viewport/docs/CHANGELOG.md b/extension/exts/omni.nerf.viewport/docs/CHANGELOG.md
similarity index 100%
rename from exts/omni.nerf.viewport/docs/CHANGELOG.md
rename to extension/exts/omni.nerf.viewport/docs/CHANGELOG.md
diff --git a/exts/omni.nerf.viewport/docs/README.md b/extension/exts/omni.nerf.viewport/docs/README.md
similarity index 100%
rename from exts/omni.nerf.viewport/docs/README.md
rename to extension/exts/omni.nerf.viewport/docs/README.md
diff --git a/exts/omni.nerf.viewport/docs/index.rst b/extension/exts/omni.nerf.viewport/docs/index.rst
similarity index 100%
rename from exts/omni.nerf.viewport/docs/index.rst
rename to extension/exts/omni.nerf.viewport/docs/index.rst
diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/__init__.py b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/__init__.py
similarity index 100%
rename from exts/omni.nerf.viewport/omni/nerf/viewport/__init__.py
rename to extension/exts/omni.nerf.viewport/omni/nerf/viewport/__init__.py
diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
similarity index 100%
rename from exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
rename to extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/tests/__init__.py b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/tests/__init__.py
similarity index 100%
rename from exts/omni.nerf.viewport/omni/nerf/viewport/tests/__init__.py
rename to extension/exts/omni.nerf.viewport/omni/nerf/viewport/tests/__init__.py
diff --git a/exts/omni.nerf.viewport/omni/nerf/viewport/tests/test_hello_world.py b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/tests/test_hello_world.py
similarity index 100%
rename from exts/omni.nerf.viewport/omni/nerf/viewport/tests/test_hello_world.py
rename to extension/exts/omni.nerf.viewport/omni/nerf/viewport/tests/test_hello_world.py
diff --git a/link_app.bat b/extension/link_app.bat
similarity index 100%
rename from link_app.bat
rename to extension/link_app.bat
diff --git a/link_app.sh b/extension/link_app.sh
similarity index 100%
rename from link_app.sh
rename to extension/link_app.sh
diff --git a/tools/packman/bootstrap/configure.bat b/extension/tools/packman/bootstrap/configure.bat
similarity index 100%
rename from tools/packman/bootstrap/configure.bat
rename to extension/tools/packman/bootstrap/configure.bat
diff --git a/tools/packman/bootstrap/download_file_from_url.ps1 b/extension/tools/packman/bootstrap/download_file_from_url.ps1
similarity index 100%
rename from tools/packman/bootstrap/download_file_from_url.ps1
rename to extension/tools/packman/bootstrap/download_file_from_url.ps1
diff --git a/tools/packman/bootstrap/fetch_file_from_packman_bootstrap.cmd b/extension/tools/packman/bootstrap/fetch_file_from_packman_bootstrap.cmd
similarity index 100%
rename from tools/packman/bootstrap/fetch_file_from_packman_bootstrap.cmd
rename to extension/tools/packman/bootstrap/fetch_file_from_packman_bootstrap.cmd
diff --git a/tools/packman/bootstrap/generate_temp_file_name.ps1 b/extension/tools/packman/bootstrap/generate_temp_file_name.ps1
similarity index 100%
rename from tools/packman/bootstrap/generate_temp_file_name.ps1
rename to extension/tools/packman/bootstrap/generate_temp_file_name.ps1
diff --git a/tools/packman/bootstrap/generate_temp_folder.ps1 b/extension/tools/packman/bootstrap/generate_temp_folder.ps1
similarity index 100%
rename from tools/packman/bootstrap/generate_temp_folder.ps1
rename to extension/tools/packman/bootstrap/generate_temp_folder.ps1
diff --git a/tools/packman/bootstrap/install_package.py b/extension/tools/packman/bootstrap/install_package.py
similarity index 100%
rename from tools/packman/bootstrap/install_package.py
rename to extension/tools/packman/bootstrap/install_package.py
diff --git a/tools/packman/config.packman.xml b/extension/tools/packman/config.packman.xml
similarity index 100%
rename from tools/packman/config.packman.xml
rename to extension/tools/packman/config.packman.xml
diff --git a/tools/packman/packman b/extension/tools/packman/packman
similarity index 100%
rename from tools/packman/packman
rename to extension/tools/packman/packman
diff --git a/tools/packman/packman.cmd b/extension/tools/packman/packman.cmd
similarity index 100%
rename from tools/packman/packman.cmd
rename to extension/tools/packman/packman.cmd
diff --git a/tools/packman/python.bat b/extension/tools/packman/python.bat
similarity index 100%
rename from tools/packman/python.bat
rename to extension/tools/packman/python.bat
diff --git a/tools/packman/python.sh b/extension/tools/packman/python.sh
similarity index 100%
rename from tools/packman/python.sh
rename to extension/tools/packman/python.sh
diff --git a/tools/scripts/link_app.py b/extension/tools/scripts/link_app.py
similarity index 100%
rename from tools/scripts/link_app.py
rename to extension/tools/scripts/link_app.py

From 7bd22b9da7a316841e2da6b24722d33b5f3f2f04 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Tue, 23 Apr 2024 06:59:56 +0800
Subject: [PATCH 06/10] refactor(nerfstudio-renderer): Isolate pygame viewer to
 its dedicated directory

---
 nerfstudio_renderer/compose.yaml => compose.yaml              | 0
 {nerfstudio_renderer/tests => pygame_viewer}/.dockerignore    | 0
 {nerfstudio_renderer/tests => pygame_viewer}/Dockerfile       | 2 +-
 {nerfstudio_renderer/tests => pygame_viewer}/pygame_test.py   | 0
 {nerfstudio_renderer/tests => pygame_viewer}/requirements.txt | 0
 {nerfstudio_renderer/tests => pygame_viewer}/run.sh           | 0
 {nerfstudio_renderer/tests => pygame_viewer}/run_local.sh     | 0
 7 files changed, 1 insertion(+), 1 deletion(-)
 rename nerfstudio_renderer/compose.yaml => compose.yaml (100%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/.dockerignore (100%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/Dockerfile (91%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/pygame_test.py (100%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/requirements.txt (100%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/run.sh (100%)
 rename {nerfstudio_renderer/tests => pygame_viewer}/run_local.sh (100%)

diff --git a/nerfstudio_renderer/compose.yaml b/compose.yaml
similarity index 100%
rename from nerfstudio_renderer/compose.yaml
rename to compose.yaml
diff --git a/nerfstudio_renderer/tests/.dockerignore b/pygame_viewer/.dockerignore
similarity index 100%
rename from nerfstudio_renderer/tests/.dockerignore
rename to pygame_viewer/.dockerignore
diff --git a/nerfstudio_renderer/tests/Dockerfile b/pygame_viewer/Dockerfile
similarity index 91%
rename from nerfstudio_renderer/tests/Dockerfile
rename to pygame_viewer/Dockerfile
index 04ccb1d..5cf2fa8 100644
--- a/nerfstudio_renderer/tests/Dockerfile
+++ b/pygame_viewer/Dockerfile
@@ -6,6 +6,6 @@ RUN apt-get update \
 WORKDIR /root
 COPY requirements.txt requirements.txt
 RUN pip install -r requirements.txt
-WORKDIR /workspace
+WORKDIR /src
 
 ENTRYPOINT bash
diff --git a/nerfstudio_renderer/tests/pygame_test.py b/pygame_viewer/pygame_test.py
similarity index 100%
rename from nerfstudio_renderer/tests/pygame_test.py
rename to pygame_viewer/pygame_test.py
diff --git a/nerfstudio_renderer/tests/requirements.txt b/pygame_viewer/requirements.txt
similarity index 100%
rename from nerfstudio_renderer/tests/requirements.txt
rename to pygame_viewer/requirements.txt
diff --git a/nerfstudio_renderer/tests/run.sh b/pygame_viewer/run.sh
similarity index 100%
rename from nerfstudio_renderer/tests/run.sh
rename to pygame_viewer/run.sh
diff --git a/nerfstudio_renderer/tests/run_local.sh b/pygame_viewer/run_local.sh
similarity index 100%
rename from nerfstudio_renderer/tests/run_local.sh
rename to pygame_viewer/run_local.sh

From 750d91ee5378d7bf25b2110f1b876e4674811132 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Tue, 23 Apr 2024 09:29:10 +0800
Subject: [PATCH 07/10] refactor: Refactor docker compose setup

---
 .gitignore                        |  1 +
 compose.yaml                      | 51 +++++++++++++++++++++++++------
 extension/.dockerignore           |  1 +
 extension/Dockerfile              |  6 ++++
 nerfstudio_renderer/.dockerignore |  1 -
 nerfstudio_renderer/.gitignore    |  1 -
 nerfstudio_renderer/Dockerfile    |  4 +--
 7 files changed, 51 insertions(+), 14 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 extension/.dockerignore
 create mode 100644 extension/Dockerfile
 delete mode 100644 nerfstudio_renderer/.gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..d93aea1
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+/assets
diff --git a/compose.yaml b/compose.yaml
index 53ed912..587991c 100644
--- a/compose.yaml
+++ b/compose.yaml
@@ -3,7 +3,7 @@ services:
   nerfstudio-renderer:
     image: j3soon/nerfstudio-renderer
     build:
-      context: .
+      context: nerfstudio_renderer
       args:
         - CUDA_VERSION=11.8.0
         - CUDA_ARCHITECTURES=86
@@ -16,9 +16,8 @@ services:
       - DISPLAY=$DISPLAY
     volumes:
       - /tmp/.X11-unix:/tmp/.X11-unix
-      - ./data/outputs:/workspace/outputs:ro
-      - ./src:/workspace/src:ro
-      - ./tests:/workspace/tests:ro
+      - ./nerfstudio_renderer/src:/src:ro
+      - ./assets:/workspace:ro
       - cache:/home/user/.cache
     shm_size: '6gb'
     deploy:
@@ -28,11 +27,11 @@ services:
             - driver: nvidia
               count: 1
               capabilities: [gpu]
-  pygame-window:
-    image: j3soon/pygame-window
+  pygame-viewer:
+    image: j3soon/pygame-viewer
     build:
-      context: tests
-    container_name: pygame-window
+      context: pygame_viewer
+    container_name: pygame-viewer
     stdin_open: true
     tty: true
     network_mode: host
@@ -40,7 +39,39 @@ services:
       - DISPLAY=$DISPLAY
     volumes:
       - /tmp/.X11-unix:/tmp/.X11-unix
-      - ./tests:/workspace:ro
-    working_dir: /workspace
+      - ./pygame_viewer:/src:ro
+  isaac-sim-viewer:
+    # Ref: https://github.com/j3soon/isaac-extended?tab=readme-ov-file#docker-container-with-display
+    image: j3soon/isaac-sim-viewer
+    build:
+      context: extension
+    container_name: isaac-sim-viewer
+    entrypoint: [bash]
+    stdin_open: true
+    tty: true
+    network_mode: host
+    environment:
+      - ACCEPT_EULA=Y
+      - PRIVACY_CONSENT=Y
+      - DISPLAY=$DISPLAY
+    volumes:
+      - ~/docker/isaac-sim/cache/kit:/isaac-sim/kit/cache:rw
+      - ~/docker/isaac-sim/cache/ov:/root/.cache/ov:rw
+      - ~/docker/isaac-sim/cache/pip:/root/.cache/pip:rw
+      - ~/docker/isaac-sim/cache/glcache:/root/.cache/nvidia/GLCache:rw
+      - ~/docker/isaac-sim/cache/computecache:/root/.nv/ComputeCache:rw
+      - ~/docker/isaac-sim/logs:/root/.nvidia-omniverse/logs:rw
+      - ~/docker/isaac-sim/data:/root/.local/share/ov/data:rw
+      - ~/docker/isaac-sim/documents:/root/Documents:rw
+      - /tmp/.X11-unix:/tmp/.X11-unix
+      - ./assets:/workspace
+      - ./extension:/src
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
 volumes:
   cache:
diff --git a/extension/.dockerignore b/extension/.dockerignore
new file mode 100644
index 0000000..72e8ffc
--- /dev/null
+++ b/extension/.dockerignore
@@ -0,0 +1 @@
+*
diff --git a/extension/Dockerfile b/extension/Dockerfile
new file mode 100644
index 0000000..89024b7
--- /dev/null
+++ b/extension/Dockerfile
@@ -0,0 +1,6 @@
+FROM nvcr.io/nvidia/isaac-sim:2023.1.1
+# Ref: https://stackoverflow.com/a/53361581
+RUN apt-get update \
+    && apt-get install -y libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+RUN /isaac-sim/python.sh -m pip install opencv-python rpyc
diff --git a/nerfstudio_renderer/.dockerignore b/nerfstudio_renderer/.dockerignore
index 7a1eba3..72e8ffc 100644
--- a/nerfstudio_renderer/.dockerignore
+++ b/nerfstudio_renderer/.dockerignore
@@ -1,2 +1 @@
 *
-!src
diff --git a/nerfstudio_renderer/.gitignore b/nerfstudio_renderer/.gitignore
deleted file mode 100644
index 3af0ccb..0000000
--- a/nerfstudio_renderer/.gitignore
+++ /dev/null
@@ -1 +0,0 @@
-/data
diff --git a/nerfstudio_renderer/Dockerfile b/nerfstudio_renderer/Dockerfile
index 41220a2..bea8579 100644
--- a/nerfstudio_renderer/Dockerfile
+++ b/nerfstudio_renderer/Dockerfile
@@ -5,11 +5,11 @@ ENV SERVER_PORT=$SERVER_PORT
 
 RUN sudo pip install rpyc
 
-WORKDIR /workspace/src
+WORKDIR /src
 
 # Force using the latest version of the source code for the ease of development
 # Entrypoint: RPyC Server
-ENTRYPOINT cp -r /workspace/src ~/src \
+ENTRYPOINT cp -r /src ~/src \
            && cd ~/src \
            && sudo pip install . \
            && python3 --version \

From 4bbc536b06cd5773a9a63793da88ba4213e3f5ab Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Tue, 23 Apr 2024 09:29:17 +0800
Subject: [PATCH 08/10] docs: Add README

---
 README.md | 108 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 README.md

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..302f144
--- /dev/null
+++ b/README.md
@@ -0,0 +1,108 @@
+# Omniverse NeRF Plugin
+
+## Prerequisites
+
+- **Hardware**:
+  - CPU: x86
+  - GPU: NVIDIA RTX GPU
+  - See [this page](https://docs.omniverse.nvidia.com/isaacsim/latest/installation/requirements.html#system-requirements) for more details.
+- **Operating System**: Ubuntu 20.04/22.04.
+- **Software**:
+  - [NVIDIA Driver](https://ubuntu.com/server/docs/nvidia-drivers-installation)
+  - [Docker](https://docs.docker.com/engine/install/ubuntu/)
+  - [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
+
+## Setup
+
+```sh
+git clone https://github.com/j3soon/omni-nerf-extension
+cd omni-nerf-extension
+# TODO: Download example files
+```
+
+The following assumes that you are running the commands from the root of the repository.
+
+## Managing Containers
+
+Login to NGC and pull the image `nvcr.io/nvidia/isaac-sim:2023.1.1` by following [this guide](https://docs.omniverse.nvidia.com/isaacsim/latest/installation/install_container.html). Then build the docker images for the extension:
+
+```sh
+docker pull nvcr.io/nvidia/isaac-sim:2023.1.1
+docker compose build
+```
+
+Launch the containers:
+
+```sh
+# You might want to use `tmux` for exec-ing into the containers later
+xhost +local:docker
+docker compose up
+```
+
+Then follow the remaining sections.
+
+To remove and stop the containers, run:
+
+```sh
+docker compose down
+```
+
+### Nerfstudio Renderer
+
+Code: [`nerfstudio_renderer`](./nerfstudio_renderer)
+
+The renderer server would be listening on port `10001` upon successful startup:
+
+```
+INFO SLAVE/10001[MainThread]: server started on [0.0.0.0]:10001
+```
+
+After seeing the above logs, no additional steps are required for the renderer server.
+
+### PyGame Viewer
+
+Code: [`pygame_viewer`](./pygame_viewer)
+
+Attach to the container and run the testing script:
+
+```sh
+docker exec -it pygame-viewer /src/run.sh
+```
+
+The script may fail at the first run due to the cold start of the renderer server. If it fails, try run the script again.
+
+(TODO: Preview Video)
+
+### Isaac Sim Viewer
+
+Code: [`extension`](./extension)
+
+```sh
+docker exec -it isaac-sim-viewer bash
+# in container
+/isaac-sim/runapp.sh --ext-folder /src/exts --enable omni.nerf.viewport
+```
+
+(TODO: Preview Video x2)
+
+## Development Notes
+
+### Nerfstudio Renderer
+
+After modifying code, you need to remove and recreate the container to apply changes. This is because the container will copy and install the code upon startup.
+
+### PyGame Viewer
+
+After modifying code, you need to re-run the testing script. The docker container can be re-used since the code is mounted as a volume.
+
+### Isaac Sim Viewer
+
+Setup VSCode intellisense for the extension:
+
+```sh
+cd extension
+./link_app.sh --path "$HOME/.local/share/ov/pkg/code-2022.3.3"
+# open the `extension` folder in VSCode
+```
+
+After modifying code, you can restart Isaac Sim to apply changes. The docker container can be re-used since the code is mounted as a volume. If the change is small, it is often faster to disable and re-enable the extension in the Isaac Sim UI. This can be done through `Window > Extensions > NVIDIA > General`, search `nerf`, and then un-toggle and re-toggle the extension.

From a1b74330df5123af552210ea3753d2a5588e12ad Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Tue, 23 Apr 2024 09:29:53 +0800
Subject: [PATCH 09/10] feat: Optimize UX by rendering higher resolution

---
 .../exts/omni.nerf.viewport/omni/nerf/viewport/extension.py   | 2 +-
 nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py   | 4 ++++
 pygame_viewer/pygame_test.py                                  | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
index 0e33dea..9986dc1 100644
--- a/extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
+++ b/extension/exts/omni.nerf.viewport/omni/nerf/viewport/extension.py
@@ -50,7 +50,7 @@ def on_startup(self, ext_id):
         # TODO: Consider subscribing to update events
         # Ref: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/events.html#subscribe-to-update-events
         # Allocate memory
-        self.rgba_w, self.rgba_h = 640, 360 # Follow default camera resolution 1280x720
+        self.rgba_w, self.rgba_h = 1280, 720 # Follow default camera resolution 1280x720
         self.rgba = np.ones((self.rgba_h, self.rgba_w, 4), dtype=np.uint8) * 128
         """RGBA image buffer. The shape is (H, W, 4), following the NumPy convention."""
         self.rgba[:,:,3] = 255
diff --git a/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py b/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
index 687de31..047686b 100644
--- a/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
+++ b/nerfstudio_renderer/src/nerfstudio_renderer/render_queue.py
@@ -70,6 +70,10 @@ def default_config():
             { 'width': 128, 'height': 72,  'fov': 35.98339777135764 },
             # 0.25x resolution
             { 'width': 320, 'height': 180, 'fov': 35.98339777135764 },
+            # 0.5x resolution
+            { 'width': 640, 'height': 360, 'fov': 35.98339777135764 },
+            # 1x resolution
+            { 'width': 1280, 'height': 720, 'fov': 35.98339777135764 },
         ])
 
     def load_config(file_path=None):
diff --git a/pygame_viewer/pygame_test.py b/pygame_viewer/pygame_test.py
index 208da19..588692a 100644
--- a/pygame_viewer/pygame_test.py
+++ b/pygame_viewer/pygame_test.py
@@ -112,7 +112,7 @@ def main(args):
         else:
             rq.update_camera(camera_position, camera_rotation)
 
-        if int(time.time()) % 3 == 0:
+        if int(time.time()) % 5 == 0:
             camera_curve_time += 1.0 / 30.0
 
     if not args.rpyc:

From 982201c13948a5fd16417ba598f753ce0b83db41 Mon Sep 17 00:00:00 2001
From: Johnson Sun <j3.soon777@gmail.com>
Date: Tue, 23 Apr 2024 10:07:37 +0800
Subject: [PATCH 10/10] docs: Add asset download commands

---
 README.md | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 302f144..fffd604 100644
--- a/README.md
+++ b/README.md
@@ -17,7 +17,23 @@
 ```sh
 git clone https://github.com/j3soon/omni-nerf-extension
 cd omni-nerf-extension
-# TODO: Download example files
+```
+
+Download assets:
+
+```sh
+wget https://github.com/j3soon/omni-nerf-extension/releases/download/v0.0.1/assets.zip
+unzip assets.zip
+```
+
+Prepare assets for `nerfstudio_renderer`:
+
+```sh
+# change the DATE_TIME to the name of the placeholder
+DATE_TIME=2023-12-30_111633
+CHECKPOINT_NAME=step-000029999
+cp -r ./assets/outputs/poster/nerfacto/$DATE_TIME ./assets/outputs/poster/nerfacto/DATE_TIME
+mv ./assets/outputs/poster/nerfacto/DATE_TIME/nerfstudio_models/$CHECKPOINT_NAME.ckpt ./assets/outputs/poster/nerfacto/DATE_TIME/nerfstudio_models/CHECKPOINT_NAME.ckpt
 ```
 
 The following assumes that you are running the commands from the root of the repository.