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

vuer/addons/nerf_vuer/render_nodes.py

@@ -3,6 +3,10 @@ RenderComponents are view components that respond to camera movements in the fro
 that returns rendered images
 """
 import base64
+import json
+from abc import abstractmethod
+from collections import deque
+from copy import deepcopy
 from io import BytesIO
 
 import numpy as np
@@ -12,15 +16,26 @@ from PIL import Image
 from torch import Tensor
 from torchtyping import TensorType
 
-from instant_feature.viewer.render_utils import get_camera_from_three
+from vuer.events import ClientEvent
 
 
-class Chainer:
+class Chainer(list):
     def __init__(self, *fns):
-        self.fns = fns
+        """Chainer is a function that chains multiple functions together
+
+        Example:
+
+            I have functions fn_1 and fn_2, both are "render nodes" for processing the rendered output from the nerf.
+            I can use Chainer to chain them together:
+
+            chained = Chainer(fn_1, fn_2)
+            chained(**pipeline) gives fn_1(fn_2(**pipeline))
+
+        """
+        super().__init__(fns)
 
     def thunk(self, **pipe):
-        for fn in self.fns:
+        for fn in self:
             try:
                 pipe = fn(**pipe)
             except Exception as e:
@@ -32,54 +47,344 @@ class Chainer:
         return self.thunk(**pipe)
 
 
-def pose(height, width, **pipeline):
-    camera_poses = get_camera_from_three([camera], width, height)
-    return {"camera_poses": camera_poses, **pipeline}
+class TrajectoryCache(deque):
+    """
+    A cache to capture the camera trajectories.
+    """
+
+    def __init__(self, **kwargs):
+        self.selections = deque(**kwargs)
+        super().__init__(**kwargs)
+
+    async def click_handler(self, event: ClientEvent, _):
+        camera_pose = deepcopy(self[-1])
+        self.selections.append(camera_pose)
+        print(camera_pose)
+
+    async def cache_camera(self, event: ClientEvent, _):
+        # do NOT mutate this object
+        camera_pose = deepcopy(event.value["camera"])
+        self.append(camera_pose)
+
+
+class Singleton(deque):
+    def __init__(self, sequence=None):
+        super().__init__(sequence or [], maxlen=1)
+        self.__post_init__()
+
+    @abstractmethod
+    def __post_init__(self):
+        pass
+
+
+class RGBA(Singleton):
+    @staticmethod
+    def rgb(raw_rgb: TensorType["hwc"], **pipeline):
+        # only needed for PNG images
+        # image_c = torch.cat([raw_rgb, raw_accumulation], dim=-1)
+        # encoded = b64jpg(image_c)
+        encoded = b64jpg(raw_rgb)
+        return {"rgb": encoded, "raw_rgb": raw_rgb, **pipeline}
+
+    @staticmethod
+    def alpha(raw_accumulation: Tensor, alpha_threshold=None, **pipeline):
+        if alpha_threshold is not None:
+            raw_accumulation[raw_accumulation < alpha_threshold] = 0
+
+        return {
+            "alpha": b64jpg(raw_accumulation),
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    @staticmethod
+    def depth(raw_depth: TensorType["hwc"], raw_accumulation, settings, **pipeline):
+        # con
+        cmap = _get_colormap(**settings)
+
+        alphaThreshold = settings.get("alphaThreshold", None)
+
+        if alphaThreshold is None:
+            mask = None
+        else:
+            mask = (raw_accumulation > alphaThreshold).squeeze().cpu()
+
+        depthmap = cmap(raw_depth.squeeze().cpu().numpy(), mask)[:, :, :3]
+        depthmap = torch.Tensor(depthmap).to(raw_depth.device)
+
+        return {
+            "depth": b64png(depthmap),
+            "raw_depth": raw_depth,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    def cache(self, raw_rgb, raw_accumulation, **pipeline):
+        rgba = torch.cat([raw_rgb, 255 * raw_accumulation], dim=-1).clip(0, 255).cpu().numpy().astype(np.uint8)
+        self.append(rgba)
+
+        return {
+            "raw_rgb": raw_rgb,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    def cache_depth(self, raw_depth, raw_accumulation, **pipeline):
+        depth_a = torch.cat([raw_depth, 255 * raw_accumulation], dim=-1).clip(0, 255).cpu().numpy().astype(np.uint8)
+        self.append(depth_a)
+
+        return {
+            "raw_depth": raw_depth,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+
+def _get_colormap(colormap, invert=False, useClip=True, clip: tuple = None, gain=1.0, offset=0.0, normalize=False, **_):
+    """
+    https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
+    get color map from matplotlib
+    returns color_map function with signature (x, mask=None),
+    where mask is the mask-in for the colormap.
+
+    """
+    import matplotlib.cm as cm
+
+    cmap = cm.get_cmap(colormap + "_r" if invert else colormap)
+
+    def map_color(x, mask=None):
+
+        if normalize:
+            if mask is None or mask.sum() == 0:
+                min, max = x.min(), x.max()
+            else:
+                min, max = x[mask].min(), x[mask].max()
 
+            x -= min
+            x /= max - min + 1e-6
+            x[x < 0] = 0
 
-def rgb(raw_rgb: TensorType["hwc"], raw_accumulation, **pipeline):
-    image_c = torch.cat([raw_rgb, raw_accumulation], dim=-1)
-    encoded = b64jpg(image_c)
-    return {"rgb": encoded, "raw_rgb": raw_rgb, "raw_accumulation": raw_accumulation, **pipeline}
+        if useClip and clip is not None:
+            x = x.clip(*clip)
 
+        if offset is not None:
+            x -= offset
 
-def alpha(raw_accumulation: Tensor, alpha_threshold=None, **pipeline):
-    if alpha_threshold is not None:
-        raw_accumulation[raw_accumulation < alpha_threshold] = 0
+        if gain is not None:
+            x *= gain
 
-    return {
-        "alpha": b64jpg(raw_accumulation),
-        "raw_accumulation": raw_accumulation,
-        **pipeline,
-    }
+        return cmap(x)
 
+    return map_color
 
-def monochrome(image, colormap, normalize, clip, gain, alpha_np, **pipeline):
-    # assert image.shape[-1] == 1 and colormap is not None, "Invalid colormap for depth"
-    # cmap = get_colormap(colormap, normalize=normalize, clip=clip, gain=gain)
-    # # Need to ignore nans and infs
-    # mask = ~torch.isnan(image) & ~torch.isinf(image)
-    # mask = mask.squeeze().cpu().numpy()
-    # image_c = cmap(image.cpu().numpy(), mask=mask)[:, :, 0, :]
-    # # Set alphas
-    # image_c[:, :, 3] = alpha_np.squeeze()
-    # # Set unmasked pixels to 0 RGBA
-    # image_c[~mask] = 0
-    # return image_c
-    pass
 
+def nan_prune(t):
+    """Drop all rows containing any nan:"""
+    t = t[~torch.any(t.isnan(), dim=-1)]
+    return t
 
-def features(raw_rgb, raw_accumulation, **pipeline):
-    # Apply PCA projection to non-nan features
-    # not_nan = ~torch.isnan(image).any(dim=-1)
-    # rgb = torch.zeros((*image.shape[:2], 3)).to(Worker.device)
-    # rgb[not_nan], info["pca_proj"], info["pca_min"], info["pca_max"] = apply_pca_colormap_return_proj(
-    #     image[not_nan], proj_V=info["pca_proj"], low_rank_min=info["pca_min"], low_rank_max=info["pca_max"]
-    # )
-    # image = torch.cat([rgb, alpha], dim=-1)
-    # image = image.cpu().numpy()
-    # return image
-    return {"raw_rgb": raw_rgb, "raw_accumulation": raw_accumulation, **pipeline}
+
+class PCA:
+    proj = None
+
+    def clear(self):
+        self.proj = None
+
+    def __call__(self, raw_features, dim=3, center=True):
+        """no batch dimension."""
+        *shape, c = raw_features.shape
+        feat_flat = raw_features.reshape(-1, c)
+
+        if self.proj is None:
+            print("Computing the PCA")
+            feat_nan_free = nan_prune(feat_flat)
+            # we can not use this u because it is potentially has the nans removed.
+            u, diag, self.proj = torch.pca_lowrank(feat_nan_free, q=dim, center=center)
+        else:
+            print("Using cached PCA")
+
+        u = raw_features @ self.proj
+        return u
+
+
+class FeatA(Singleton):
+    """Cache for Feature and Alpha Channel. Contains two caches, self, and low_res_cache.
+    The low-res cache is for the low-res feature map.
+
+    Self is the flat cache.
+    """
+
+    def __post_init__(self):
+        self.pca = PCA()
+        self.low_res_cache = deque(maxlen=1)
+        self.features_cache = deque(maxlen=1)
+
+    def features_pca(self, raw_features, raw_accumulation, alpha=None, **pipeline):
+        # print("feat_pca.shape:", raw_features.shape)
+        feat_pca = self.pca(raw_features, dim=3)
+        feat_pca_flat = feat_pca.reshape(-1, 3)
+
+        # these should be configurations
+        low = torch.nanquantile(feat_pca_flat, 0.02, dim=0)
+        top = torch.nanquantile(feat_pca_flat, 0.98, dim=0)
+
+        feat_pca_normalized = 0.02 + 0.96 * (feat_pca - low) / (top - low)
+        feat_pca_normalized.clip_(0, 1)
+
+        if alpha is None:
+            alpha = b64jpg(raw_accumulation)
+
+        return {
+            "features": b64jpg(feat_pca_normalized),
+            "features_pca": feat_pca_normalized,
+            "alpha": alpha,
+            "raw_features": raw_features,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    def probe_feature(self, raw_features, **pipeline):
+        return {"raw_features": raw_features, **pipeline}
+
+    def features_heatmap(self, raw_features, raw_accumulation, settings, **pipeline):
+
+        print("raw_features.shape:", raw_features.shape)
+        print(settings)
+
+        x, y = settings.pop("xy", [None, None])
+        size = settings.pop("size", None)
+        n = max(1, int(size))
+
+        # remove clip so that lerf_text_map does not complain about hashability in the lru_cache.
+        clip = settings.pop("clip", None)
+
+        h, w = raw_features.shape[:2]
+        i, j = int(x * h), int(y * h)
+        feat_probe = raw_features[i : i + n, j : j + n, :].to(raw_accumulation.device)
+        feat_probe = feat_probe.mean(dim=[0, 1])
+        feat_probe /= feat_probe.norm(dim=-1, keepdim=True)
+
+        raw_features_cuda = raw_features.to(raw_accumulation.device)
+        raw_features_cuda /= raw_features_cuda.norm(dim=-1, keepdim=True)
+        raw_features_cuda.shape
+        # heatmap
+        heatmap = raw_features_cuda @ feat_probe
+        del raw_features_cuda
+        mask = ~torch.isnan(heatmap)
+
+        cmap = _get_colormap(clip=clip, **settings)
+        heatmap_rgb = cmap(heatmap.cpu(), mask.cpu())[:, :, :3]
+        heatmap_rgb = torch.FloatTensor(heatmap_rgb).to(raw_accumulation.device)
+
+        # mask_float = mask.float()[..., None] * raw_accumulation
+        mask_float = raw_accumulation
+        print("raw_accumulation", raw_accumulation.shape)
+
+        return {
+            "heatmap": b64jpg(heatmap_rgb),
+            "heatmap_mask": b64jpg(mask_float),
+            "heatmap_rgb": heatmap_rgb,
+            "raw_heatmap_mask": mask_float,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    def cache_heatmap(self, heatmap_rgb, raw_heatmap_mask, **pipeline):
+        heat_alpha = torch.cat([heatmap_rgb, 255 * raw_heatmap_mask], dim=-1).clip(0, 255).cpu().numpy().astype(np.uint8)
+        self.append(heat_alpha)
+        # remove those from the pipeline
+        return {**pipeline}
+
+    def cache_pca(self, features_pca, raw_accumulation, **pipeline):
+        fpca_cuda = features_pca.to(raw_accumulation.device)
+        feat_alpha = torch.cat([fpca_cuda, 255 * raw_accumulation], dim=-1).clip(0, 255).cpu().numpy().astype(np.uint8)
+        self.append(feat_alpha)
+        return {"raw_accumulation": raw_accumulation, **pipeline}
+
+    def cache_features(self, raw_features, **pipeline):
+        """this cache is very expensive."""
+        self.features_cache.append(raw_features)
+        return {"raw_features": raw_features, **pipeline}
+
+    def cache_low_res(self, raw_features, **pipeline):
+        from einops import rearrange
+
+        h, w = raw_features.shape[:2]
+        size = h * w - torch.isnan(raw_features).any(dim=-1).sum()
+        ratio = float(120 / size) ** 0.5
+        feat = rearrange(raw_features, "h w c -> 1 c h w")
+        try:
+            feat = torch.nn.functional.upsample(feat, scale_factor=ratio, mode="bilinear")
+            feat = rearrange(feat, "1 c h w -> h w c")
+            feat_flat = feat[~torch.isnan(feat).any(dim=-1)]
+            self.low_res_cache.append(feat_flat)
+            return {
+                # note: remove [ raw_features, and features_pca ] from the flow to conserve memory
+                **pipeline,
+            }
+        except Exception as e:
+            print("upsample failed", e)
+            return pipeline
+
+    def reset_pca(self):
+        self.pca.clear()
+
+
+class CLIPQueryMap(Singleton):
+    """Cache for Feature and Alpha Channel. Contains two caches, self, and low_res_cache.
+    The low-res cache is for the low-res feature map.
+
+    Self is the flat cache."""
+
+    def set_query_vector(self, query_vector):
+        self.query_vector = query_vector
+
+    def text_heatmap(self, raw_features, raw_accumulation, settings, **pipeline):
+        from instant_feature.viewer.nerf_vuer.clip.clip_heatmap import get_lerf_text_map
+
+        print("raw_features.shape:", raw_features.shape)
+        print(settings)
+
+        # remove clip so that lerf_text_map does not complain about hashability in the lru_cache.
+        clip = settings.pop("clip", None)
+
+        text_map = get_lerf_text_map(
+            **settings,
+            device=raw_accumulation.device,
+        )
+
+        raw_features_cuda = raw_features.to(raw_accumulation.device)
+        heatmap, mask = text_map(raw_features_cuda)
+        del raw_features_cuda
+
+        heatmap_normalized = heatmap
+        # heatmap_normalized = deepcopy(raw_features)
+        # heatmap_normalized.clip_(0, 1)
+        print("heatmap_normalized.shape:", heatmap_normalized.shape)
+
+        # final_mask = raw_accumulation * mask.float()[..., None]
+
+        cmap = _get_colormap(clip=clip, **settings)
+        heatmap_rgb = cmap(heatmap_normalized.cpu(), mask.cpu())[:, :, :3]
+        heatmap_rgb = torch.FloatTensor(heatmap_rgb).to(raw_accumulation.device)
+
+        mask_float = mask.float()[..., None]
+
+        return {
+            "heatmap": b64jpg(heatmap_rgb),
+            "heatmap_mask": b64jpg(mask_float),
+            "raw_heatmap": heatmap_rgb,
+            "raw_heatmap_mask": mask_float,
+            "raw_accumulation": raw_accumulation,
+            **pipeline,
+        }
+
+    def cache(self, raw_heatmap, raw_heatmap_mask, **pipeline):
+        print("cache")
+        feat_alpha = torch.cat([raw_heatmap, 255 * raw_heatmap_mask], dim=-1).clip(0, 255).cpu().numpy().astype(np.uint8)
+        print("done caching")
+        self.append(feat_alpha)
+        return {**pipeline}
 
 
 def b64jpg(image: Tensor, quality: int = 90):
@@ -93,8 +398,11 @@ def b64jpg(image: Tensor, quality: int = 90):
         image = image[:, :, :3]
 
     image *= 255
-    image_np = image.cpu().numpy().astype(np.uint8)
-    image_np.max()
+
+    if isinstance(image, np.ndarray):
+        image_np = image.astype(np.uint8)
+    else:
+        image_np = image.cpu().numpy().astype(np.uint8)
 
     C = image_np.shape[-1]
     if C == 1:
@@ -119,7 +427,7 @@ def b64png(image: Tensor):
     rgb_pil = Image.fromarray(image_np)
     rgb_pil.save(buff, format="PNG")
     img64 = base64.b64encode(buff.getbuffer().tobytes()).decode("utf-8")
-    return img64
+    return "data:image/png;base64," + img64
 
 
 def b64png_depth(depth):

vuer/base.py

@@ -15,13 +15,13 @@ async def default_handler(request, ws):
         print(msg)
 
 
-async def websocket_handler(request, handler):
-    print('New connection!!!')
+async def websocket_handler(request, handler, **ws_kwargs):
+    print("New connection!!!")
 
-    ws = web.WebSocketResponse()
+    ws = web.WebSocketResponse(**ws_kwargs)
     await ws.prepare(request)
 
-    print('Socket stored')
+    print("Socket stored")
 
     try:
         await handler(request, ws)
@@ -30,18 +30,18 @@ async def websocket_handler(request, handler):
         print("Connection reset")
 
     except CancelledError as exp:
-        print(f'WebSocket Canceled')
+        print(f"WebSocket Canceled")
 
     except Exception as exp:
-        print(f'Error:\n{exp}\n{traceback.print_exc()}')
+        print(f"Error:\n{exp}\n{traceback.print_exc()}")
 
     finally:
         await ws.close()
-        print('WebSocket connection closed')
+        print("WebSocket connection closed")
 
 
 async def handle_file_request(request, root):
-    filename = request.match_info['filename']
+    filename = request.match_info["filename"]
     filepath = Path(root) / filename
 
     if not filepath.is_file():
@@ -64,20 +64,26 @@ class Server:
 
         default = aiohttp_cors.ResourceOptions(
             allow_credentials=True,
-            expose_headers="*", allow_methods="*",
+            expose_headers="*",
+            allow_headers="*",
+            allow_methods="*",
         )
-        cors_config = {k: default for k in self.cors.split(',')}
+        cors_config = {k: default for k in self.cors.split(",")}
 
         self.cors_context = aiohttp_cors.setup(self.app, defaults=cors_config)
 
-    def _route(self, path: str, handler: callable, method: str = "GET", ):
+    def _route(
+        self,
+        path: str,
+        handler: callable,
+        method: str = "GET",
+    ):
         route = self.app.router.add_resource(path).add_route(method, handler)
         self.cors_context.add(route)
 
     def _socket(self, path: str, handler: callable):
-        ws_handler = partial(websocket_handler, handler=handler)
-        route = self.app.router.add_resource(path).add_route("GET", ws_handler)
-        self.cors_context.add(route)
+        ws_handler = partial(websocket_handler, handler=handler, max_msg_size=self.WEBSOCKET_MAX_SIZE)
+        self._route(path, ws_handler, method="GET")
 
     def _add_task(self, fn):
         loop = asyncio.get_event_loop()
@@ -85,7 +91,7 @@ class Server:
 
     def _static(self, path, root):
         _fn = partial(handle_file_request, root=root)
-        self.app.add_routes([web.get(f'/{path}/{{filename}}', _fn)])
+        self._route(f"{path}/{{filename}}", _fn, method="GET")
 
     def run(self):
         async def init_server():
@@ -94,7 +100,7 @@ class Server:
             site = web.TCPSite(runner, self.host, self.port)
             await site.start()
 
-            print(f'Serving on http://{self.host}:{self.port}')
+            print(f"Serving on http://{self.host}:{self.port}")
 
         event_loop = asyncio.get_event_loop()
 
@@ -102,8 +108,8 @@ class Server:
         event_loop.run_forever()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     app = Server()
-    app._route('', websocket_handler)
-    app._static('static', handle_file_request, root='.')
+    app._route("", websocket_handler)
+    app._static("/static", handle_file_request, root=".")
     app.run()

vuer/events.py

@@ -1,4 +1,5 @@
 from vuer.schemas import Element
+from vuer.serdes import serializer
 
 
 class Event:
@@ -54,27 +55,6 @@ NULL = NullEvent()
 #         instance = cls.__new__(cls)
 #         return cls.__init__(instance, data)
 
-from typing import Sequence
-
-
-def serializer(data):
-    if hasattr(data, "serialize"):
-        return data.serialize()
-
-    if isinstance(data, str):
-        # return Text(data)
-        return data
-
-    # this could be dangerous.
-    if isinstance(data, Sequence):
-        return [serializer(d) for d in data]
-
-    # this could be dangerous
-    if isinstance(data, dict):
-        return {k: serializer(v) for k, v in data.items()}
-
-    NotImplementedError(f"Cannot serialize {data}")
-
 
 class ServerEvent(Event):  # , metaclass=Meta):
     def __init__(self, data, etype=None, **kwargs):