Merge branch 'master' into mugamma/MET-32-fmb

Author: mugamma

CMakeLists.txt

@@ -14,13 +14,16 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 ################
 
 add_library(genmetaballs_core
+  genmetaballs/src/cuda/core/camera.cu
+  genmetaballs/src/cuda/core/camera.cuh
   genmetaballs/src/cuda/core/utils.cu
   genmetaballs/src/cuda/core/utils.cuh
   genmetaballs/src/cuda/core/fmb.cuh
   genmetaballs/src/cuda/core/fmb.cu
   genmetaballs/src/cuda/core/geometry.cuh
   genmetaballs/src/cuda/core/geometry.cu
   genmetaballs/src/cuda/core/confidence.cuh
+  genmetaballs/src/cuda/core/image.cuh
 )
 
 # Set include directories for the core library

examples/run_fmb.py

@@ -33,7 +33,7 @@
 from tqdm import tqdm
 
 # Import local utilities
-from genmetaballs.fmb.utils import DegradeLR, image_grid
+from genmetaballs.fmb.utils import DegradeLR, get_camera_rays, image_grid
 
 CURRENT_DIR = Path(__file__).parent
 PROJECT_ROOT = CURRENT_DIR.parent
@@ -475,14 +475,12 @@ def main():
 
     # Setup camera rays
     height, width = image_size
-    K = np.array([[focal_length, 0, cx], [0, focal_length, cy], [0, 0, 1]])
     pixel_list = (
         (np.array(np.meshgrid(np.arange(width), height - np.arange(height) - 1, [0]))[:, :, :, 0])
         .reshape((3, -1))
         .T
     )
-    camera_rays = (pixel_list - K[:, 2]) / np.diag(K)
-    camera_rays[:, -1] = -1
+    camera_rays = get_camera_rays(focal_length, focal_length, cx, cy, pixel_list)
     cameras_list = []
     for tran, quat in zip(trans, rand_quats, strict=False):
         R = transforms3d.quaternions.quat2mat(quat)

genmetaballs/src/cuda/bindings.cu

@@ -5,15 +5,22 @@
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
+#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/utils.cuh"
 
 namespace nb = nanobind;
 
 template <typename T, MemoryLocation location>
 void bind_array2d(nb::module_& m, const char* name);
+template <MemoryLocation location>
+void bind_image(nb::module_& m, const char* name);
+template <MemoryLocation location>
+void bind_image_view(nb::module_& m, const char* name);
 
 NB_MODULE(_genmetaballs_bindings, m) {
 
@@ -82,18 +89,53 @@ NB_MODULE(_genmetaballs_bindings, m) {
         .def_rw("start", &Ray::start)
         .def_rw("direction", &Ray::direction);
 
+    /*
+     * Camera module bindings
+     */
+    nb::module_ camera = m.def_submodule("camera", "Camera intrinsics and extrinsics");
+    nb::class_<Intrinsics>(camera, "Intrinsics")
+        .def(nb::init<uint32_t, uint32_t, float, float, float, float>(), nb::arg("height"),
+             nb::arg("width"), nb::arg("fx"), nb::arg("fy"), nb::arg("cx"), nb::arg("cy"))
+        .def_ro("height", &Intrinsics::height)
+        .def_ro("width", &Intrinsics::width)
+        .def_ro("fx", &Intrinsics::fx)
+        .def_ro("fy", &Intrinsics::fy)
+        .def_ro("cx", &Intrinsics::cx)
+        .def_ro("cy", &Intrinsics::cy)
+        .def("get_ray_direction", &Intrinsics::get_ray_direction,
+             "Get the direction of the ray going through pixel (px, py) in camera frame",
+             nb::arg("px"), nb::arg("py"));
+
+    /*
+     * Image module bindings
+     */
+    nb::module_ image = m.def_submodule("image", "Image data structure for GenMetaballs");
+    bind_image_view<MemoryLocation::HOST>(image, "CPUImageView");
+    bind_image<MemoryLocation::HOST>(image, "CPUImage");
+    bind_image_view<MemoryLocation::DEVICE>(image, "GPUImageView");
+    bind_image<MemoryLocation::DEVICE>(image, "GPUImage");
+
     /*
      * Confidence module bindings
      */
 
     nb::module_ confidence = m.def_submodule("confidence");
     nb::class_<ZeroParameterConfidence>(confidence, "ZeroParameterConfidence")
         .def(nb::init<>())
-        .def("get_confidence", &ZeroParameterConfidence::get_confidence);
+        .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("get_confidence", &TwoParameterConfidence::get_confidence);
+        .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);
+        });
 
     /*
      * Utils module bindings
@@ -102,6 +144,32 @@ NB_MODULE(_genmetaballs_bindings, m) {
     nb::module_ utils = m.def_submodule("utils");
     utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))");
 
+    // blender submodule
+    nb::module_ blender = m.def_submodule("blender");
+    nb::class_<FourParameterBlender>(blender, "FourParameterBlender")
+        .def(nb::init<float, float, float, float>())
+        .def_ro("beta1", &FourParameterBlender::beta1)
+        .def_ro("beta2", &FourParameterBlender::beta2)
+        .def_ro("beta3", &FourParameterBlender::beta3)
+        .def_ro("eta", &FourParameterBlender::eta)
+        .def("blend", &FourParameterBlender::blend, nb::arg("t"), nb::arg("d"),
+             "Blend two values with (t,d)")
+        .def("__repr__", [](const FourParameterBlender& b) {
+            return nb::str("FourParameterBlender(beta1={}, beta2={}, beta3={}, eta={})")
+                .format(b.beta1, b.beta2, b.beta3, b.eta);
+        });
+
+    nb::class_<ThreeParameterBlender>(blender, "ThreeParameterBlender")
+        .def(nb::init<float, float, float>())
+        .def_ro("beta1", &ThreeParameterBlender::beta1)
+        .def_ro("beta2", &ThreeParameterBlender::beta2)
+        .def_ro("eta", &ThreeParameterBlender::eta)
+        .def("blend", &ThreeParameterBlender::blend, nb::arg("t"), nb::arg("d"),
+             "Blend two values with (t,d)")
+        .def("__repr__", [](const ThreeParameterBlender& b) {
+            return nb::str("ThreeParameterBlender(beta1={}, beta2={}, eta={})")
+                .format(b.beta1, b.beta2, b.eta);
+        });
     bind_array2d<float, MemoryLocation::HOST>(utils, "CPUFloatArray2D");
     bind_array2d<float, MemoryLocation::DEVICE>(utils, "GPUFloatArray2D");
 
@@ -137,3 +205,30 @@ void bind_array2d(nb::module_& m, const char* name) {
         .def_prop_ro("ndim", &Array2D<T, location>::ndim)
         .def_prop_ro("size", &Array2D<T, location>::size);
 }
+
+template <MemoryLocation location>
+void bind_image_view(nb::module_& m, const char* name) {
+    nb::class_<ImageView<location>>(m, name)
+        .def(nb::init<const Array2D<float, location>&, const Array2D<float, location>&>(),
+             nb::arg("confidence"), nb::arg("depth"))
+        .def_prop_ro("confidence", [](const ImageView<location>& view) { return view.confidence; })
+        .def_prop_ro("depth", [](const ImageView<location>& view) { return view.depth; })
+        .def_prop_ro("num_rows", &ImageView<location>::num_rows)
+        .def_prop_ro("num_cols", &ImageView<location>::num_cols)
+        .def("__repr__", [=](const ImageView<location>& view) {
+            return nb::str("{}(height={}, width={})")
+                .format(name, view.num_rows(), view.num_cols());
+        });
+}
+
+template <MemoryLocation location>
+void bind_image(nb::module_& m, const char* name) {
+    nb::class_<Image<location>>(m, name)
+        .def(nb::init<uint32_t, uint32_t>(), nb::arg("height"), nb::arg("width"))
+        .def_prop_ro("num_rows", &Image<location>::num_rows)
+        .def_prop_ro("num_cols", &Image<location>::num_cols)
+        .def("as_view", &Image<location>::as_view, "Get a view of the image data as ImageView")
+        .def("__repr__", [=](const Image<location>& img) {
+            return nb::str("{}(height={}, width={})").format(name, img.num_rows(), img.num_cols());
+        });
+}

genmetaballs/src/cuda/core/blender.cuh

@@ -1,14 +1,27 @@
 #pragma once
 
-#include "fmb.h"
-#include "geometry.h"
+#include <cmath>
+#include <cuda_runtime.h>
+
+#include "utils.cuh"
+
+struct FourParameterBlender {
+    float beta1;
+    float beta2;
+    float beta3;
+    float eta;
+
+    CUDA_CALLABLE __forceinline__ float blend(float t, float d) const {
+        return expf((beta1 * d * sigmoid((beta3 / eta) * t)) - ((beta2 / eta) * t));
+    }
+};
 
 struct ThreeParameterBlender {
     float beta1;
     float beta2;
     float eta;
 
-    CUDA_CALLABLE __forceinline__ // TODO inline?
-        float
-        blend(float t, float d, const FMB& fmb, const Ray& ray) const;
+    CUDA_CALLABLE __forceinline__ float blend(float t, float d) const {
+        return expf((beta1 * d) - ((beta2 / eta) * t));
+    }
 };

genmetaballs/src/cuda/core/camera.cu

@@ -0,0 +1,13 @@
+#include <cstdint>
+#include <cuda/std/ranges>
+#include <cuda_runtime.h>
+
+#include "camera.cuh"
+#include "geometry.cuh"
+#include "utils.cuh"
+
+CUDA_CALLABLE Vec3D Intrinsics::get_ray_direction(uint32_t px, uint32_t py) const {
+    auto x = (static_cast<float>(px) - cx) / fx;
+    auto y = (static_cast<float>(py) - cy) / fy;
+    return Vec3D{x, y, -1.0f};
+}

genmetaballs/src/cuda/core/camera.cuh

@@ -1,14 +1,36 @@
 #pragma once
 
 #include <cstdint>
+#include <cuda_runtime.h>
+
+#include "geometry.cuh"
+#include "utils.cuh"
 
 struct Intrinsics {
-    uint32_t height;
-    uint32_t width;
+    uint32_t height; // in x direction
+    uint32_t width;  // in y direction
     float fx;
     float fy;
     float cx;
     float cy;
-    float near;
-    float far;
+
+    // Returns the direction of the ray going through pixel (px, py) in camera frame.
+    // For efficiency, this function does not check if the pixel is within bounds.
+    CUDA_CALLABLE Vec3D get_ray_direction(uint32_t px, uint32_t py) const;
+
+    // Returns a 2D array of ray directions in camera frame in the specified pixel range
+    // and store them in the provided buffer. By default, the full image is used
+    template <MemoryLocation location>
+    CUDA_CALLABLE Array2D<Vec3D, location>& get_ray_directions(Array2D<Vec3D, location>& buffer,
+                                                               uint32_t px_start = 0,
+                                                               uint32_t px_end = UINT32_MAX,
+                                                               uint32_t py_start = 0,
+                                                               uint32_t py_end = UINT32_MAX) const {
+        for (auto i = max(0, px_start); i < min(height, px_end); ++i) {
+            for (auto j = max(0, py_start); j < min(width, py_end); ++j) {
+                buffer[i][j] = get_ray_direction(j, i);
+            }
+        }
+        return buffer;
+    }
 };

genmetaballs/src/cuda/core/fmb.cuh

@@ -107,4 +107,10 @@ public:
     CUDA_CALLABLE ConstIterator end() const {
         return ConstIterator(fmbs_ + size_, log_weights_ + size_);
     }
+    CUDA_CALLABLE const FMB& get_fmb(uint32_t idx) const {
+        return fmbs_[idx];
+    }
+    CUDA_CALLABLE size_t size() const {
+        return size_;
+    }
 };

genmetaballs/src/cuda/core/forward.cu

@@ -47,7 +47,7 @@ __global__ render_kernel(const Getter fmb_getter, const Blender blender,
 template <class Getter, class Intersector, class Blender, class Confidence>
 void render_fmbs(const FMBs& fmbs, const Intrinsics& intr, const Pose& extr) {
     // initialize the fmb_getter
-    typename Getter::Getter fmb_getter(fmbs, intr, extr);
+    typename Getter::Getter fmb_getter(fmbs, extr);
     auto kernel = render_kernel<Getter, Intersector, Blender, Confidence>;
     kernel<<<NUM_BLOCKS, THREADS_PER_BLOCK>>>(fmb_getter, fmbs, intr, extr);
 }

genmetaballs/src/cuda/core/geometry.cuh

@@ -72,11 +72,11 @@ public:
         return {rot, tran};
     }
 
-    CUDA_CALLABLE Rotation get_rot() const {
+    CUDA_CALLABLE const Rotation& get_rot() const {
         return rot_;
     }
 
-    CUDA_CALLABLE Vec3D get_tran() const {
+    CUDA_CALLABLE const Vec3D& get_tran() const {
         return tran_;
     }
 

genmetaballs/src/cuda/core/getter.cuh

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cmath>
+#include <cuda_runtime.h>
+
+#include "camera.cuh"
+#include "fmb.cuh"
+#include "geometry.cuh"
+#include "utils.cuh"
+
+// This is the dummy version of getter, where all FMBs are relevant to any ray
+template <MemoryLocation location>
+struct AllGetter {
+    FMBScene<location>& scene;
+    Pose& extr; // Current assumption: rays are in camera frame
+
+    CUDA_CALLABLE AllGetter(FMBScene<location>& scene, Pose& extr)
+        : scene(scene), extr(extr) {}
+
+    // It does not bother using the ray, because it simply returns all FMBs
+    CUDA_CALLABLE FMBScene<location>& get_metaballs(const Ray& ray) const {
+        return scene;
+    }
+};