bindings working

Author: mugamma

genmetaballs/src/cuda/bindings.cu

@@ -4,13 +4,15 @@
 #include <nanobind/operators.h>
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
+#include <tuple>
 
 #include "core/blender.cuh"
 #include "core/camera.cuh"
 #include "core/confidence.cuh"
 #include "core/fmb.cuh"
 #include "core/geometry.cuh"
 #include "core/image.cuh"
+#include "core/intersector.cuh"
 #include "core/utils.cuh"
 
 namespace nb = nanobind;
@@ -24,6 +26,28 @@ void bind_image_view(nb::module_& m, const char* name);
 
 NB_MODULE(_genmetaballs_bindings, m) {
 
+    /*
+     * Confidence module bindings
+     */
+
+    nb::module_ confidence = m.def_submodule("confidence");
+    nb::class_<ZeroParameterConfidence>(confidence, "ZeroParameterConfidence")
+        .def(nb::init<>())
+        .def("get_confidence", &ZeroParameterConfidence::get_confidence, nb::arg("sumexpd"),
+             "Get the confidence value for a given sumexpd")
+        .def("__repr__",
+             [](const ZeroParameterConfidence& c) { return nb::str("ZeroParameterConfidence()"); });
+
+    nb::class_<TwoParameterConfidence>(confidence, "TwoParameterConfidence")
+        .def(nb::init<float, float>())
+        .def_ro("beta4", &TwoParameterConfidence::beta4)
+        .def_ro("beta5", &TwoParameterConfidence::beta5)
+        .def("get_confidence", &TwoParameterConfidence::get_confidence, nb::arg("sumexpd"),
+             "Get the confidence value for a given sumexpd")
+        .def("__repr__", [](const TwoParameterConfidence& c) {
+            return nb::str("TwoParameterConfidence(beta4={}, beta5={})").format(c.beta4, c.beta5);
+        });
+
     /*
      * FMB module bindings
      */
@@ -85,9 +109,9 @@ NB_MODULE(_genmetaballs_bindings, m) {
         .def("inv", &Pose::inv, "Inverse pose");
 
     nb::class_<Ray>(geometry, "Ray")
-        .def(nb::init<>())
-        .def_rw("start", &Ray::start)
-        .def_rw("direction", &Ray::direction);
+        .def(nb::init<Vec3D, Vec3D>())
+        .def_ro("start", &Ray::start)
+        .def_ro("direction", &Ray::direction);
 
     /*
      * Camera module bindings
@@ -116,26 +140,17 @@ NB_MODULE(_genmetaballs_bindings, m) {
     bind_image<MemoryLocation::DEVICE>(image, "GPUImage");
 
     /*
-     * Confidence module bindings
+     * Intersector module bindings
      */
 
-    nb::module_ confidence = m.def_submodule("confidence");
-    nb::class_<ZeroParameterConfidence>(confidence, "ZeroParameterConfidence")
-        .def(nb::init<>())
-        .def("get_confidence", &ZeroParameterConfidence::get_confidence, nb::arg("sumexpd"),
-             "Get the confidence value for a given sumexpd")
-        .def("__repr__",
-             [](const ZeroParameterConfidence& c) { return nb::str("ZeroParameterConfidence()"); });
-
-    nb::class_<TwoParameterConfidence>(confidence, "TwoParameterConfidence")
-        .def(nb::init<float, float>())
-        .def_ro("beta4", &TwoParameterConfidence::beta4)
-        .def_ro("beta5", &TwoParameterConfidence::beta5)
-        .def("get_confidence", &TwoParameterConfidence::get_confidence, nb::arg("sumexpd"),
-             "Get the confidence value for a given sumexpd")
-        .def("__repr__", [](const TwoParameterConfidence& c) {
-            return nb::str("TwoParameterConfidence(beta4={}, beta5={})").format(c.beta4, c.beta5);
-        });
+    nb::module_ intersector = m.def_submodule("intersector");
+    intersector.def(
+        "linear_intersect",
+        [](const FMB& fmb, const Ray& ray) {
+            auto [t, d] = LinearIntersector::intersect(fmb, ray);
+            return std::make_tuple(t, d);
+        },
+        "Linear intersection of ray and FMB.", nb::arg("fmb"), nb::arg("ray"));
 
     /*
      * Utils module bindings

genmetaballs/src/cuda/core/geometry.cuh

@@ -102,4 +102,6 @@ public:
 struct Ray {
     Vec3D start;
     Vec3D direction;
+
+    CUDA_CALLABLE Ray(const Vec3D _start, const Vec3D _dir) : start{_start}, direction{_dir} {}
 };

genmetaballs/src/cuda/core/intersector.cuh

@@ -1,19 +1,20 @@
 #pragma once
 
-#include <utility>
+#include <cuda/std/tuple>
 
-#include "fmb.h"
-#include "geometry.h"
+#include "fmb.cuh"
+#include "geometry.cuh"
 
 // implement equation (6) in the paper
 class LinearIntersector {
-
-    static CUDA_CALLABLE std::pair<float, float> intersect(const FMB& fmb, const Ray& ray) const {
+public:
+    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Ray& ray) {
         auto vecdiv = [](const Vec3D& u, const Vec3D& v) {
             return Vec3D{u.x / v.x, u.y / v.y, u.z / v.z};
         };
         auto rot = fmb.get_pose().get_rot();
         auto tmp = rot.inv().apply(vecdiv(rot.apply(ray.direction), fmb.get_extent()));
-        return dot(fmb.get_pose().get_tran() - ray.start, tmp) / dot(ray.direction, tmp);
+        auto t = dot(fmb.get_pose().get_tran() - ray.start, tmp) / dot(ray.direction, tmp);
+        return {t, fmb.quadratic_form(ray.start + t * ray.direction)};
     }
 };

genmetaballs/src/genmetaballs/core/__init__.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from genmetaballs._genmetaballs_bindings import fmb, geometry
+from genmetaballs._genmetaballs_bindings import fmb, geometry, intersector
 from genmetaballs._genmetaballs_bindings.blender import (
     FourParameterBlender,
     ThreeParameterBlender,
@@ -55,6 +55,7 @@ def make_image(height: int, width: int, device: DeviceType) -> CPUImage | GPUIma
     "geometry",
     "Camera",
     "Intrinsics",
+    "intersector",
     "sigmoid",
     "FourParameterBlender",
     "ThreeParameterBlender",