Parametrized forward rendering smoke test to check all variants of renderer

Author: horizon-blue

genmetaballs/src/cuda/bindings.cu

@@ -27,6 +27,8 @@ template <MemoryLocation location>
 void bind_image_view(nb::module_& m, const char* name);
 template <MemoryLocation location>
 void bind_fmb_scene(nb::module_& m, const char* name);
+template <typename Blender, typename Confidence>
+void bind_render_fmbs(nb::module_& m, const char* name);
 
 NB_MODULE(_genmetaballs_bindings, m) {
 
@@ -77,14 +79,14 @@ NB_MODULE(_genmetaballs_bindings, m) {
      * Forward (rendering) module bindings
      */
     nb::module_ forward = m.def_submodule("forward", "Forward rendering of FMBs");
-    // TODO: turn this into a template function to allow different combinations of
-    // Getter/Intersector/Blender/Confidence at runtime
-    forward.def("render_fmbs",
-                &render_fmbs<AllGetter<MemoryLocation::DEVICE>, LinearIntersector,
-                             ThreeParameterBlender, TwoParameterConfidence>,
-                "Render FMBs using FourParameterBlender and TwoParameterConfidence",
-                nb::arg("fmbs"), nb::arg("blender"), nb::arg("confidence"), nb::arg("intr"),
-                nb::arg("extr"), nb::arg("img"));
+    bind_render_fmbs<FourParameterBlender, ZeroParameterConfidence>(
+        forward, "render_fmbs_four_param_zero_confidence");
+    bind_render_fmbs<ThreeParameterBlender, TwoParameterConfidence>(
+        forward, "render_fmbs_three_param_two_confidence");
+    bind_render_fmbs<ThreeParameterBlender, ZeroParameterConfidence>(
+        forward, "render_fmbs_three_param_zero_confidence");
+    bind_render_fmbs<FourParameterBlender, TwoParameterConfidence>(
+        forward, "render_fmbs_four_param_two_confidence");
 
     /*
      * Geometry module bindings
@@ -277,3 +279,12 @@ void bind_fmb_scene(nb::module_& m, const char* name) {
             return nb::str("{}(size={})").format(name, scene.size());
         });
 }
+
+template <typename Blender, typename Confidence>
+void bind_render_fmbs(nb::module_& m, const char* name) {
+    m.def(name,
+          &render_fmbs<AllGetter<MemoryLocation::DEVICE>, LinearIntersector, Blender, Confidence>,
+          "Render the given FMB scene into the provided image view", nb::arg("fmbs"),
+          nb::arg("blender"), nb::arg("confidence"), nb::arg("intr"), nb::arg("extr"),
+          nb::arg("img"));
+}

genmetaballs/src/genmetaballs/core/__init__.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from genmetaballs._genmetaballs_bindings import fmb, geometry, intersector
+from genmetaballs._genmetaballs_bindings import fmb, forward, geometry, intersector
 from genmetaballs._genmetaballs_bindings.blender import (
     FourParameterBlender,
     ThreeParameterBlender,
@@ -11,8 +11,7 @@
     ZeroParameterConfidence,
 )
 from genmetaballs._genmetaballs_bindings.fmb import CPUFMBScene, GPUFMBScene
-from genmetaballs._genmetaballs_bindings.forward import render_fmbs
-from genmetaballs._genmetaballs_bindings.image import CPUImage, GPUImage
+from genmetaballs._genmetaballs_bindings.image import CPUImage, GPUImage, GPUImageView
 from genmetaballs._genmetaballs_bindings.utils import CPUFloatArray2D, GPUFloatArray2D, sigmoid
 
 type DeviceType = Literal["cpu", "gpu"]
@@ -64,6 +63,41 @@ def make_fmb_scene(size: int, device: DeviceType) -> CPUFMBScene | GPUFMBScene:
         raise ValueError(f"Unsupported device type: {device}")
 
 
+def render_fmbs(
+    fmbs: GPUFMBScene,
+    blender: FourParameterBlender | ThreeParameterBlender,
+    confidence: TwoParameterConfidence | ZeroParameterConfidence,
+    intr: Intrinsics,
+    extr: geometry.Pose,
+    img: GPUImage | None = None,
+) -> GPUImage:
+    """Render the given FMB scene into the provided image view.
+
+    If no image is provided, a new image with the dimensions specified by the intrinsics
+    will be created.
+    """
+    if img is None:
+        img = make_image(intr.height, intr.width, device="gpu")
+
+    if isinstance(blender, FourParameterBlender):
+        if isinstance(confidence, ZeroParameterConfidence):
+            render_func = forward.render_fmbs_four_param_zero_confidence
+        elif isinstance(confidence, TwoParameterConfidence):
+            render_func = forward.render_fmbs_four_param_two_confidence
+    elif isinstance(blender, ThreeParameterBlender):
+        if isinstance(confidence, ZeroParameterConfidence):
+            render_func = forward.render_fmbs_three_param_zero_confidence
+        elif isinstance(confidence, TwoParameterConfidence):
+            render_func = forward.render_fmbs_three_param_two_confidence
+    else:
+        raise TypeError("Unsupported blender and confidence combination.")
+
+    render_func(fmbs, blender, confidence, intr, extr, img.as_view())
+    return img
+
+    forward.render_fmbs(fmbs, blender, confidence, intr, extr, img)
+
+
 __all__ = [
     "array2d_float",
     "fmb",

tests/python_tests/test_forward.py

@@ -1,28 +1,42 @@
+import jax.numpy as jnp
+import pytest
 from jax.scipy.spatial.transform import Rotation
 
 from genmetaballs.core import (
+    FourParameterBlender,
     Intrinsics,
     ThreeParameterBlender,
     TwoParameterConfidence,
+    ZeroParameterConfidence,
     geometry,
     make_fmb_scene,
-    make_image,
     render_fmbs,
 )
 
 
-def test_render_fmbs_smoke() -> None:
-    """checks that render_fmbs can be called without errors"""
+@pytest.mark.parametrize(
+    "confidence", [TwoParameterConfidence(beta4=0.5, beta5=-1.0), ZeroParameterConfidence()]
+)
+@pytest.mark.parametrize(
+    "blender",
+    [
+        ThreeParameterBlender(beta1=1.0, beta2=0.5, eta=2.0),
+        FourParameterBlender(beta1=1.0, beta2=0.5, beta3=0.3, eta=2.0),
+    ],
+)
+def test_render_fmbs_smoke(blender, confidence) -> None:
+    """checks that render_fmbs can be called without errors with combinations of blender and confidence."""
     scene = make_fmb_scene(20, device="gpu")
-    blender = ThreeParameterBlender(beta1=1.0, beta2=0.5, eta=2.0)
-    confidence = TwoParameterConfidence(beta4=0.5, beta5=-1.0)
     camera = Intrinsics(fx=100.0, fy=100.0, cx=50.0, cy=50.0, width=100, height=120)
     rotation = Rotation.identity()
     pose = geometry.Pose.from_components(
         rot=geometry.Rotation.from_quat(*rotation.as_quat()),
         tran=geometry.Vec3D(0.0, 0.0, 5.0),
     )
-    # allocate output buffer
-    image = make_image(camera.height, camera.width, "gpu")
+    # Note that it's also possible to allocate the buffer ahead of time and pass it in,
+    # which could be useful if we want to reuse the same image buffer for multiple renders.
+    image = render_fmbs(scene, blender, confidence, camera, pose)
+    img_view = image.as_view()
 
-    render_fmbs(scene, blender, confidence, camera, pose, image.as_view())
+    assert img_view.num_rows == camera.height
+    assert img_view.num_cols == camera.width