vuer 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

vuer/addons/nerf_vuer/mixins.py

@@ -1,7 +1,8 @@
 from collections import defaultdict
+from copy import deepcopy
 from dataclasses import dataclass
 from functools import wraps
-from typing import Dict, Union, List, DefaultDict, Optional, Generator
+from typing import Dict, Union, List, DefaultDict, Optional, Generator, Callable, Sequence
 
 from instant_feature.cameras.cameras import Cameras
 from instant_feature.viewer.se3 import rotation_matrix
@@ -106,7 +107,16 @@ def process_aabb(render_gen):
         render,
         **kwargs,
     ):
-        if render.get("use_aabb", None):
+
+        if settings.get("use_aabb", None):
+            # use_aabb = settings.get("use_aabb", None)
+            # aabb_min = settings.get("aabb_min", None)
+            # aabb_max = settings.get("aabb_max", None)
+            scene_box = AABB(
+                aabb_min=Vector3(**settings.get("aabb_min", None)),
+                aabb_max=Vector3(**settings.get("aabb_max", None)),
+            )
+        elif render.get("use_aabb", None):
             scene_box = AABB(
                 aabb_min=Vector3(**render.get("aabb_min", None)),
                 aabb_max=Vector3(**render.get("aabb_max", None)),
@@ -114,35 +124,32 @@ def process_aabb(render_gen):
         else:
             scene_box = None
 
-        # if settings.get("use_aabb", None):
-        #     use_aabb = settings.get("use_aabb", None)
-        #     aabb_min = settings.get("aabb_min", None)
-        #     aabb_max = settings.get("aabb_max", None)
-
         async for event in render_gen(*args, aabb=scene_box, settings=settings, **kwargs):
             yield event
 
     return wrap_gen
 
 
-def _se3(position, rotation, scale, **rest) -> Dict[str, Union[Vector3, EulerDeg, float]]:
-    return dict(
-        position=Vector3(**position),
-        rotation=EulerDeg(**rotation).to_rad(),
-        scale=scale,
-        **rest,
-    )
+def _se3(position=None, rotation=None, scale=None, **rest) -> Dict[str, Union[Vector3, EulerDeg, float]]:
+    """
+    Add handling of undefined parameters.
+    """
+    output = rest
+    if position:
+        output["position"] = Vector3(**position)
+    if rotation:
+        output["rotation"] = EulerDeg(**rotation).to_rad()
+    if isinstance(scale, dict):
+        output["scale"] = Vector3(**scale)
+    elif scale:
+        output["scale"] = scale
+    return output
 
 
 # Need more encapsulated handling. Transformations should be propagated down.
 def process_world(render_gen):
     # @wraps(render_gen)
     async def wrap_gen(*args, camera, world, settings, **kwargs):
-        # position = Vector3(**world["position"])
-        # # sets the unit to "deg" internally.
-        # rotation = EulerDeg(**world["rotation"])
-        # scale = world["scale"]
-
         async for event in render_gen(
             *args,
             parent=_se3(**world),
@@ -205,6 +212,17 @@ def _lie_action(ray_bundle, position: Vector3, rotation: Euler, scale: float, **
         ray_bundle.origins /= scale
 
 
+def _transformation(rotation: Euler, position: Vector3, scale: Union[float, Vector3], **_) -> torch.Tensor:
+    rot_mat = torch.from_numpy(rotation_matrix(*rotation))
+    transform = torch.eye(4)
+    if isinstance(scale, Vector3):
+        transform[:3, :3] = rot_mat @ torch.DoubleTensor(scale).diag()
+    else:
+        transform[:3, :3] = rot_mat * scale
+    transform[:3, 3] = torch.FloatTensor(position)
+    return transform
+
+
 def collect_rays(render_bundle):
     # @wraps(render_bundle)
     async def ray_gen(
@@ -212,24 +230,30 @@ def collect_rays(render_bundle):
         camera: Cameras,
         parent: Dict = None,
         settings: Dict = None,
-        # position: Vector3 = None,
-        # # use degree convention on the front end
-        # rotation: EulerDeg = None,
-        # # Non-scalar scale is current not supported.
-        # scale: Union[Vector3, float] = None,
         aabb: Optional[AABB] = None,
         **kwargs,
     ) -> Dict[str, torch.Tensor]:
         assert len(camera) == 1
-        # todo: to correctly handle things, we need to transform the camera instead.
-        ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=aabb)
 
-        ray_bundle = ray_bundle.to(camera.device)
+        parent2world = _transformation(**parent)
+
+        if settings and "rotation" in settings:
+            settings2parent = _transformation(**settings)
+            settings2world = parent2world @ settings2parent
+            world2setting = torch.linalg.inv(settings2world)
+
+        else:
+            world2setting = torch.linalg.inv(parent2world)
+
+        camera2world = torch.eye(4)
+        camera2world[:3, :] = camera.camera_to_worlds[0]
+        camera2setting = world2setting @ camera2world
 
-        # note: Ge: this needs to be moved up the stack.
-        _lie_action(ray_bundle, **parent)
-        # apply render specific transformations
-        _lie_action(ray_bundle, **settings)
+        camera = deepcopy(camera)
+        camera.camera_to_worlds[0] = camera2setting[:3]
+
+        ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=aabb)
+        ray_bundle = ray_bundle.to(camera.device)
 
         # timing is not accurate
         # with torch.no_grad(), logger.time("rendering", fmt=lambda s: f"{1000 * s:.1f}ms"):
@@ -242,22 +266,27 @@ def collect_rays(render_bundle):
     return ray_gen
 
 
-# from .render_nodes import rgb, alpha, chainer
-# this is not a mix-in
-
-
 def collector(
     pipe: Chainer = None,
     channels: List[str] = None,
-    # nodes=List[RenderNode],  # tuple(rgb, alpha),
+    channel_prefix: Union[str, Callable] = None,
 ):
     """
     Decorator for processing the rendered images.
+
+    Params:
+    pipe: this is the chain of operations that forms the processing flow.
+    channels: keys sent over the wire to the client. We use whitelisting to recude bandwidth.
+    layer_prefix: when this is set, the output channels are attached a prefix, to accommodate
+            conflicting keys in the pipeline.
+            Also takes in a function, to allow more dynamics generation of prefixes
     """
 
     def decorator(async_render):
         # @wraps(async_render)
+        # adding self argument to allow passing into the prefix function
         async def render_event(
+            self,
             *args,
             camera,
             pipe=pipe,
@@ -266,11 +295,8 @@ def collector(
 
             outputs = defaultdict(list)
 
-            async for output_chunk in async_render(*args, channels=channels, camera=camera, **kwargs):
+            async for output_chunk in async_render(self, *args, channels=channels, camera=camera, **kwargs):
                 signal = yield None
-                # not needed. Outer can simply delete.
-                # if signal == "TERMINATE":
-                #     return
 
                 # collect the results
                 for key, value in output_chunk.items():
@@ -283,11 +309,24 @@ def collector(
                 height=camera.height,
                 prefix="raw_",
             )
-
             # pipe = chainer(*nodes)
-            flow = pipe(**flow)
-            # filter out the non-requested channels
-            flow = {k: v for k, v in flow.items() if k in channels}
+            flow = pipe(
+                **flow,
+                render=kwargs["render"],
+                settings=kwargs["settings"],
+            )
+
+            # info: do NOT mutate channel_prefix. This is inside the closure and is shared between function calls.
+            prefix = channel_prefix
+            if callable(channel_prefix):
+                prefix = channel_prefix(self)
+
+            if prefix:
+                # select channels, and prefix with
+                flow = {prefix + k: v for k, v in flow.items() if k in channels}
+            else:
+                # filter out the non-requested channels
+                flow = {k: v for k, v in flow.items() if k in channels}
 
             yield ServerEvent(
                 etype="RENDER",
@@ -296,6 +335,10 @@ def collector(
                 data={**flow},
             )
 
+        # holds the current pipe.
+        render_event._pipeline = pipe
+        render_event._channels = channels
+
         return render_event
 
     return decorator

vuer/addons/nerf_vuer/nerf_vuer.py

@@ -1,4 +1,5 @@
 from asyncio import sleep
+from collections import defaultdict
 from copy import deepcopy
 
 import numpy as np
@@ -17,13 +18,17 @@ from vuer.schemas import (
 class RenderVuer(Vuer):
     device = "cuda:0"
 
-    def __init__(self, render: Render = None, scene: Scene = None, **args):
-        super().__init__()
+    WEBSOCKET_MAX_SIZE = 2**28
+
+    def __init__(self, render: Render = None, scene: Scene = None, **kwargs):
+        super().__init__(**kwargs)
 
         self.scene = scene
         self.render = render
 
-        self.spawn(self.on_connect, start=True)
+        self.handlers = defaultdict(dict)
+
+        self.spawn(self.on_connect)
 
     async def on_connect(self, ws_id):
         # from ml_logger import logger
@@ -40,8 +45,6 @@ class RenderVuer(Vuer):
 
             self.set @ scene
 
-        cprint("Need to allow for mulitple event handlers", color="yellow")
-
         # need to explicitly terminate the coroutine
         while ws_id in self.ws:
             await sleep(0)
@@ -51,11 +54,24 @@ class RenderVuer(Vuer):
             if not event:
                 continue
 
+            if event.etype in self.handlers:
+                handlers = self.handlers[event.etype]
+                for fn_factory in handlers.values():
+                    # todo: see if we want to add throttling here.
+                    # also pass in an event handler.
+                    # Use an arrow function to avoid exposing the server instance.
+                    my_task = self.spawn_task(fn_factory(event, lambda e: self @ e))
+                    await sleep(0.0)
+
             if event == "CAMERA_MOVE":
                 value = event.value
                 self.clear()
 
-                world = value.pop("world")
+                world = value.pop("world", None)
+                if world is None:
+                    print("initial camera movement does not contain world params.")
+                    continue
+
                 camera = value.pop("camera")
                 render_params = value.pop("render", {})
                 # initially the values are non.
@@ -88,7 +104,7 @@ class RenderVuer(Vuer):
                     async for render_response in self.render.render(
                         camera=quick_cam,
                         world=world,
-                        render=render_params,
+                        render=deepcopy(render_params),
                         # other params
                         chunk_size=8096,
                         to_cpu="features",
@@ -107,7 +123,7 @@ class RenderVuer(Vuer):
                 async for render_response in self.render.render(
                     camera=camera,
                     world=world,
-                    render=render_params,
+                    render=deepcopy(render_params),
                     # other params
                     chunk_size=8096,
                     to_cpu="features",

vuer/addons/nerf_vuer/render_components.py

@@ -1,12 +1,11 @@
 from typing import List, Generator
 
 from vuer.addons.nerf_vuer.mixins import collector, process_aabb, process_world, collect_rays, chunk_rays
-from vuer.addons.nerf_vuer.render_nodes import rgb, Chainer, alpha
-from vuer.addons.nerf_vuer.control_components import Controls
+from vuer.addons.nerf_vuer.render_nodes import Chainer, RGBA
 from vuer.events import ServerEvent
 from vuer.schemas import SceneElement
 
-from instant_feature.viewer.neko.constants.default_settings import RENDER_DEFAULT, RGB_DEFAULT
+from instant_feature.viewer.nerf_vuer.constants.default_settings import RENDER_DEFAULT, RGB_DEFAULT
 
 
 # =============== Leva Components ===============
@@ -55,12 +54,6 @@ class Render(SceneElement):
             **kwargs,
         )
 
-    # def serialize(self):
-    #     obj = super().serialize()
-    #     # if self.settings:
-    #     #     obj["settings"] = self.settings
-    #     return obj
-
     def __post_init__(self):
         self._fields = {}
         for child in self.children:
@@ -90,21 +83,30 @@ class Render(SceneElement):
 class RenderLayer(SceneElement):
     tag = "RenderLayer"
 
+    channel = "rgb"
+    alphaChannel = "alpha"
+    displacementMap = None
+    distance = 10.1
+
+    settings = RGB_DEFAULT
+
     def __init__(
         self,
-        channel="rgb",
-        alphaChannel="alpha",
+        channel=None,
+        alphaChannel=None,
         displacementMap=None,
         geometry="plane",
-        settings=RGB_DEFAULT,
+        settings=None,
+        distance=None,
         **kwargs,
     ):
         super().__init__(
-            channel=channel,
-            alphaChannel=alphaChannel,
-            displacementMap=displacementMap,
+            channel=channel or self.channel,
+            alphaChannel=alphaChannel or self.alphaChannel,
+            displacementMap=displacementMap or self.displacementMap,
             geometry=geometry,
-            settings=settings,
+            settings=settings or self.settings,
+            distance=distance or self.distance,
             **kwargs,
         )
 
@@ -112,7 +114,7 @@ class RenderLayer(SceneElement):
         self._render = _render
 
     @collector(
-        pipe=Chainer(rgb, alpha),
+        pipe=Chainer(RGBA.rgb, RGBA.alpha),
         channels=["rgb", "alpha"],
     )
     @process_aabb
@@ -120,10 +122,7 @@ class RenderLayer(SceneElement):
     @collect_rays
     @chunk_rays
     def render(self, ray_bundle, **kwargs) -> Generator[ServerEvent, None, None]:
-        import torch
-
-        with torch.no_grad():
-            return self._render(ray_bundle=ray_bundle, **kwargs)
+        return self._render(ray_bundle=ray_bundle, **kwargs)
 
 
 class Heatmap(RenderLayer):