Smoke test to make sure that forward function runs

Author: horizon-blue

genmetaballs/src/cuda/bindings.cu

@@ -43,7 +43,7 @@ NB_MODULE(_genmetaballs_bindings, m) {
              [](const ZeroParameterConfidence& c) { return nb::str("ZeroParameterConfidence()"); });
 
     nb::class_<TwoParameterConfidence>(confidence, "TwoParameterConfidence")
-        .def(nb::init<float, float>())
+        .def(nb::init<float, float>(), nb::arg("beta4"), nb::arg("beta5"))
         .def_ro("beta4", &TwoParameterConfidence::beta4)
         .def_ro("beta5", &TwoParameterConfidence::beta5)
         .def("get_confidence", &TwoParameterConfidence::get_confidence, nb::arg("sumexpd"),
@@ -81,7 +81,7 @@ NB_MODULE(_genmetaballs_bindings, m) {
     // Getter/Intersector/Blender/Confidence at runtime
     forward.def("render_fmbs",
                 &render_fmbs<AllGetter<MemoryLocation::DEVICE>, LinearIntersector,
-                             FourParameterBlender, TwoParameterConfidence>,
+                             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"));
@@ -178,7 +178,8 @@ NB_MODULE(_genmetaballs_bindings, m) {
     // blender submodule
     nb::module_ blender = m.def_submodule("blender");
     nb::class_<FourParameterBlender>(blender, "FourParameterBlender")
-        .def(nb::init<float, float, float, float>())
+        .def(nb::init<float, float, float, float>(), nb::arg("beta1"), nb::arg("beta2"),
+             nb::arg("beta3"), nb::arg("eta"))
         .def_ro("beta1", &FourParameterBlender::beta1)
         .def_ro("beta2", &FourParameterBlender::beta2)
         .def_ro("beta3", &FourParameterBlender::beta3)
@@ -191,7 +192,7 @@ NB_MODULE(_genmetaballs_bindings, m) {
         });
 
     nb::class_<ThreeParameterBlender>(blender, "ThreeParameterBlender")
-        .def(nb::init<float, float, float>())
+        .def(nb::init<float, float, float>(), nb::arg("beta1"), nb::arg("beta2"), nb::arg("eta"))
         .def_ro("beta1", &ThreeParameterBlender::beta1)
         .def_ro("beta2", &ThreeParameterBlender::beta2)
         .def_ro("eta", &ThreeParameterBlender::eta)

genmetaballs/src/genmetaballs/core/__init__.py

@@ -11,6 +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.utils import CPUFloatArray2D, GPUFloatArray2D, sigmoid
 
@@ -65,16 +66,17 @@ def make_fmb_scene(size: int, device: DeviceType) -> CPUFMBScene | GPUFMBScene:
 
 __all__ = [
     "array2d_float",
-    "ZeroParameterConfidence",
-    "TwoParameterConfidence",
     "fmb",
     "geometry",
-    "Camera",
-    "Intrinsics",
     "intersector",
+    "make_fmb_scene",
+    "make_image",
+    "render_fmbs",
     "sigmoid",
+    "Camera",
     "FourParameterBlender",
+    "Intrinsics",
     "ThreeParameterBlender",
-    "make_image",
-    "make_fmb_scene",
+    "TwoParameterConfidence",
+    "ZeroParameterConfidence",
 ]

tests/python_tests/test_forward.py

@@ -0,0 +1,28 @@
+from jax.scipy.spatial.transform import Rotation
+
+from genmetaballs.core import (
+    Intrinsics,
+    ThreeParameterBlender,
+    TwoParameterConfidence,
+    geometry,
+    make_fmb_scene,
+    make_image,
+    render_fmbs,
+)
+
+
+def test_render_fmbs_smoke() -> None:
+    """checks that render_fmbs can be called without errors"""
+    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")
+
+    render_fmbs(scene, blender, confidence, camera, pose, image.as_view())