fix: integrate metric specific flows in test viz

+ various monads/components

Author: nmvrs

moai/conf/data/test/loader/torch.yaml

@@ -6,7 +6,7 @@ shuffle: false
 num_workers: ${oc.select:experiment.workers,0}
 pin_memory: false
 drop_last: false
-prefetch_factor: 0
+prefetch_factor: null # 0
 persistent_workers: false
 pin_memory_device: ""
 # in_order: true

moai/conf/data/train/loader/torch.yaml

@@ -6,7 +6,7 @@ shuffle: true
 num_workers: ${oc.select:experiment.workers,0}
 pin_memory: false
 drop_last: false
-prefetch_factor: 0
+prefetch_factor: null # 0
 persistent_workers: false
 pin_memory_device: ""
 # in_order: true

moai/conf/model/components/tcnn/encoding_mlp.yaml

@@ -2,6 +2,7 @@
 
 _target_: moai.components.tcnn.modules.EncodingMLP
 input_dims: ???
+hidden_dims: ???
 output_dims: ???
 encoding_config: ???
 seed: ${oc.select:engine.modules.manual_seed.seed,1337}

moai/conf/model/monads/generation/tensor/linspace.yaml

@@ -0,0 +1,6 @@
+# @package model.monads.linspace
+
+_target_: moai.monads.generation.tensor.torch.LinSpace
+start: 0.0
+end: 1.0
+steps: ???

moai/conf/model/monads/math/exp.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.exp
 
-_target_: moai.monads.math.Exponential
+_target_: moai.monads.math.Exponential

moai/conf/model/monads/math/normalize.yaml

@@ -0,0 +1,6 @@
+# @package model.monads.normalize
+
+_target_: moai.monads.math.Normalize
+order: 2
+dim: -1
+epsilon: 1e-12

moai/conf/model/monads/math/sigmoid.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.sigmoid
 
-_target_: torch.nn.Sigmoid
+_target_: moai.monads.math.Sigmoid

moai/conf/model/monads/render/splats/gaussian_splat_rasterizer_slothfulxtx.yaml

@@ -0,0 +1,9 @@
+# @package model.monads.gaussian_splat_rasterizer
+
+_target_: moai.monads.render.splats.rasterize_slothfulxtx.GaussianSplatRasterizer
+prefiltered: false
+debug: false
+scale_modifier: 1.0
+width: null
+height: null
+background_color: black

moai/conf/model/monads/tensor/float32.yaml

@@ -0,0 +1,3 @@
+# @package model.monads.float32
+
+_target_: moai.monads.tensor.cast.Float32

moai/core/model.py

@@ -453,6 +453,8 @@ def test_step(
         if monitor := get_dict(
             self.monitor, f"{C._TEST_}.{C._DATASETS_}.{dataset_name}"
         ):
+            for step in toolz.get(C._FLOWS_, monitor, None) or []:
+                self.named_flows[step](batch)
             for metric in get_list(monitor, C._METRICS_):  # Metrics monitoring
                 self.named_metrics[metric](batch)
             for tensor_monitor in get_list(monitor, C._MONITORS_):

moai/monads/generation/tensor/torch.py

@@ -17,6 +17,7 @@
     "RandomLike",
     "OnesLike",
     "TemporalParams",
+    "LinSpace",
 ]
 
 
@@ -234,3 +235,17 @@ def __init__(
 
     def forward(self, void: torch.Tensor) -> torch.nn.parameter.Parameter:
         return dict(self.named_parameters())
+
+
+class LinSpace(torch.nn.Module):
+    def __init__(self, start: float, end: float, steps: int) -> None:
+        super().__init__()
+        self.register_buffer(
+            "spaced", torch.linspace(start, end, steps)[np.newaxis, :, np.newaxis]
+        )
+
+    def forward(
+        self,
+        tensor: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.spaced.expand(tensor.shape[0], -1, 1)

moai/monads/math/__init__.py

@@ -10,9 +10,10 @@
     PlusOne,
     Rad2Deg,
     Scale,
+    Sigmoid,
 )
 from moai.monads.math.dot import Dot
-from moai.monads.math.normalization import MinMaxNorm, Znorm
+from moai.monads.math.normalization import MinMaxNorm, Normalize, Znorm
 
 __all__ = [
     "Abs",
@@ -30,4 +31,6 @@
     "Rad2Deg",
     "Deg2Rad",
     "Exponential",
+    "Normalization",
+    "Sigmoid",
 ]

moai/monads/math/common.py

@@ -12,6 +12,7 @@
     "Rad2Deg",
     "Deg2Rad",
     "Exponential",
+    "Sigmoid",
 ]
 
 
@@ -104,3 +105,11 @@ def __init__(self) -> None:
 
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         return torch.exp(tensor)
+
+
+class Sigmoid(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return torch.sigmoid(tensor)

moai/monads/math/normalization.py

@@ -31,3 +31,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         #     .mul(scale.view(b, 1, 1, 1))\
         #     .add(self.min)
         return torch.addcmul(self.min, x - expand_dims(mins, x), expand_dims(scale, x))
+
+
+class Normalize(torch.nn.Module):
+    def __init__(self, order: int = 2, dim: int = -1, epsilon: float = 1e-12):
+        super().__init__()
+        self.order = order
+        self.dim = dim
+        self.eps = epsilon
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.normalize(
+            tensor, p=self.order, dim=self.dim, eps=self.eps
+        )

moai/monads/render/splats/rasterize_slothfulxtx.py

@@ -0,0 +1,111 @@
+import math
+import typing
+
+import numpy as np
+import torch
+from colour import Color
+from diff_gauss import GaussianRasterizationSettings, GaussianRasterizer
+
+__all__ = ["GaussianSplatRasterizer"]
+
+
+# NOTE: different gaussian splatting versions produce different number and order of outputs
+#       but all collide using the same package
+
+
+class GaussianSplatRasterizer(torch.nn.Module):
+    def __init__(
+        self,
+        width: typing.Optional[int] = None,
+        height: typing.Optional[int] = None,
+        background_color: str = "black",
+        prefiltered: bool = False,
+        debug: bool = False,
+        scale_modifier: float = 1.0,
+    ):
+        super().__init__()
+        color = Color(background_color).get_rgb()
+        self.register_buffer("background_color", torch.tensor(color))
+        self.width, self.height = width, height
+        self.prefiltered, self.debug = prefiltered, debug
+        self.scale_modifier = scale_modifier
+
+    def forward(
+        self,
+        view_matrix: torch.Tensor,  # [B, 4, 4]
+        view_projection_matrix: torch.Tensor,  # [B, 4, 4]
+        camera_position: torch.Tensor,  # [B, 3]
+        positions: torch.Tensor,  # [B, V, 3]
+        sh_coeffs: torch.Tensor,  # [B, SH, 3]
+        opacities: torch.Tensor,  # [B, V, 1]
+        rotations: torch.Tensor,  # [B, V, 4]
+        scales: torch.Tensor,  # [B, V, 3]
+        features: torch.Tensor,  # [B, V, K]
+        intrinsics: torch.Tensor,  # [B, 3, 3]
+        image: typing.Optional[torch.Tensor] = None,  # [B, C, H, W]
+        background_color: typing.Optional[torch.Tensor] = None,  # [B, 3]
+    ):
+        assert len(positions.shape) == 3
+        B = view_matrix.shape[0]
+        if positions.shape[0] != B:  # either many-to-many or one-to-many
+            positions = positions.expand(B, -1, -1)
+            sh_coeffs = sh_coeffs.expand(B, -1, -1, -1)
+            opacities = opacities.expand(B, -1, -1)
+            rotations = rotations.expand(B, -1, -1)
+            scales = scales.expand(B, -1, -1)
+            features = features.expand(B, -1, -1)
+        bg = background_color if background_color is not None else self.background_color
+        sh_degree = math.sqrt(sh_coeffs.shape[-2]) - 1
+        if bg.shape[0] != B:
+            bg = bg.expand(B, -1)
+        colors, radiis, depths, alphas, extras = [], [], [], [], []
+        for i in range(B):
+            W = image[i].shape[-1] if image is not None else self.width
+            H = image[i].shape[-2] if image is not None else self.height
+            settings = GaussianRasterizationSettings(
+                image_height=H,
+                image_width=W,
+                # tanfovx=2.0 * np.arctan(W / (2.0 * intrinsics[i, 0, 0].cpu().float())),
+                # tanfovy=2.0 * np.arctan(H / (2.0 * intrinsics[i, 1, 1].cpu().float())),
+                tanfovx=W / (2.0 * intrinsics[i, 0, 0].cpu().float()),
+                tanfovy=H / (2.0 * intrinsics[i, 1, 1].cpu().float()),
+                bg=bg[i],
+                scale_modifier=self.scale_modifier,
+                viewmatrix=view_matrix[i],
+                projmatrix=view_projection_matrix[i],
+                sh_degree=int(sh_degree),
+                campos=camera_position[i],
+                prefiltered=self.prefiltered,
+                debug=self.debug,
+            )
+            screenspace_points = torch.zeros_like(
+                positions[i]
+            )  # , requires_grad=True) + 0
+            screenspace_points.requires_grad_(True)
+            screenspace_points.retain_grad()
+            rasterizer = GaussianRasterizer(settings)
+            color, depth, norm, alpha, radii, feats = rasterizer(  # TODO inv_depth
+                means3D=positions[i],
+                means2D=screenspace_points,
+                opacities=opacities[i],
+                scales=scales[i],
+                rotations=rotations[i],
+                shs=sh_coeffs[i],
+                extra_attrs=features[i],
+                colors_precomp=None,
+                cov3Ds_precomp=None,
+            )
+            colors.append(color)
+            radiis.append(radii)
+            depths.append(depth)
+            alphas.append(alpha)
+            extras.append(feats)
+            # "viewspace_points": screenspace_points,
+            # "visibility_filter" : radii > 0,
+        return {
+            "color": torch.stack(colors).clamp(0, 1),
+            "radii": torch.stack(radiis),
+            "depth": torch.stack(depths),
+            "alpha": torch.stack(alphas),
+            "features": torch.stack(extras),
+        }

moai/monads/tensor/cast.py

@@ -1,6 +1,6 @@
 import torch
 
-__all__ = ["Int32"]
+__all__ = ["Int32, Float32"]
 
 
 class Int32(torch.nn.Module):
@@ -9,3 +9,11 @@ def __init__(self) -> None:
 
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.int()
+
+
+class Float32(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.float()