diff --git a/include/sdf_tools/sdf.hpp b/include/sdf_tools/sdf.hpp index 0294a27..03b8d02 100644 --- a/include/sdf_tools/sdf.hpp +++ b/include/sdf_tools/sdf.hpp @@ -335,6 +335,7 @@ class SignedDistanceField : public VoxelGrid::VoxelGrid ///////////////////////////////////////////////////////////////////////// using GradientFunction = std::function(int64_t,int64_t,int64_t,bool)>; + using HessianFunction = std::function>(int64_t,int64_t,int64_t,bool)>; /** the return vector is a flat array, which will be reshape to be * X/Y/Z/gradient when saved as numpy **/ @@ -356,7 +357,26 @@ class SignedDistanceField : public VoxelGrid::VoxelGrid } return gradient_grid; } + inline VoxelGrid>> GetFullHessian(const HessianFunction& hessian_function, + const bool enable_edge_gradients = false) const + { + std::vector init_val(3, oob_value_); + VoxelGrid>> hessian_grid{origin_transform_, GetResolution(), GetNumXCells(), GetNumYCells(), + GetNumZCells(), std::vector>(3, init_val)}; + for (auto x_idx{0l}; x_idx < GetNumXCells(); ++x_idx) + { + for (auto y_idx{0l}; y_idx < GetNumYCells(); ++y_idx) + { + for (auto z_idx{0l}; z_idx < GetNumZCells(); ++z_idx) + { + auto const hessian = hessian_function(x_idx, y_idx, z_idx, enable_edge_gradients); + hessian_grid.SetValue(x_idx, y_idx, z_idx, hessian); + } + } + } + return hessian_grid; + } inline std::vector GetGradient( const double x, const double y, const double z, @@ -402,6 +422,182 @@ class SignedDistanceField : public VoxelGrid::VoxelGrid return GetGradient(index.x, index.y, index.z, enable_edge_gradients); } + inline std::vector> GetHessian( + const GRID_INDEX& index, + const bool enable_edge_gradients = false) const + { + return GetHessian(index.x, index.y, index.z, enable_edge_gradients); + } + + inline std::vector> GetHessian( + const int64_t x_index, const int64_t y_index, const int64_t z_index, + const bool enable_edge_gradients = false) const + { + const std::vector > grid_aligned_hessian + = GetGridAlignedHessian(x_index, y_index, z_index, + enable_edge_gradients); + + std::vector> hessian; + + for (const auto grid_aligned_hessian_row: grid_aligned_hessian) { + if (grid_aligned_hessian_row.size() == 3) { + const Eigen::Quaterniond grid_rotation(origin_transform_.rotation()); + // Derived from EigenHelpers::RotateVector, but without extra copies + const Eigen::Quaterniond temp(0.0, + grid_aligned_hessian_row[0], + grid_aligned_hessian_row[1], + grid_aligned_hessian_row[2]); + const Eigen::Quaterniond result = grid_rotation * (temp * grid_rotation.inverse()); + + std::vector tmp{result.x(), result.y(), result.z()}; + hessian.push_back(tmp); + } else { + return std::vector>(); + } + } + return hessian; + } + + inline std::vector> GetGridAlignedHessian( + const int64_t x_index, const int64_t y_index, const int64_t z_index, + const bool enable_edge_gradients = false) const + { + // Make sure the index is inside bounds + if (IndexInBounds(x_index, y_index, z_index)) + { + // See if the index we're trying to query is one cell in from the edge + if ((x_index > 0) && (y_index > 0) && (z_index > 0) + && (x_index < (GetNumXCells() - 1)) + && (y_index < (GetNumYCells() - 1)) + && (z_index < (GetNumZCells() - 1))) + { + const double inv_resolution = 1.0 / GetResolution(); + + const double hxx = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index + 1, y_index, z_index).first + + GetImmutable(x_index - 1, y_index, z_index).first) + * inv_resolution * inv_resolution; + const double hyy = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index, y_index+1, z_index).first + + GetImmutable(x_index, y_index-1, z_index).first) + * inv_resolution * inv_resolution; + const double hzz = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index, y_index, z_index+1).first + + GetImmutable(x_index, y_index, z_index-1).first) + * inv_resolution * inv_resolution; + + const double hxy = (GetImmutable(x_index + 1, y_index + 1, z_index).first + - GetImmutable(x_index + 1, y_index - 1, z_index).first + - GetImmutable(x_index - 1, y_index + 1, z_index).first + + GetImmutable(x_index - 1, y_index - 1, z_index).first) + * 0.25 * inv_resolution * inv_resolution; + const double hxz = (GetImmutable(x_index + 1, y_index, z_index+1).first + - GetImmutable(x_index + 1, y_index, z_index - 1).first + - GetImmutable(x_index - 1, y_index, z_index + 1).first + + GetImmutable(x_index - 1, y_index, z_index - 1).first) + * 0.25 * inv_resolution * inv_resolution; + const double hyz = (GetImmutable(x_index, y_index + 1, z_index + 1).first + - GetImmutable(x_index, y_index + 1, z_index - 1).first + - GetImmutable(x_index, y_index - 1, z_index + 1).first + + GetImmutable(x_index, y_index - 1, z_index - 1).first) + * 0.25 * inv_resolution * inv_resolution; + + std::vector h1{hxx, hxy, hxz}; + std::vector h2{hxy, hyy, hyz}; + std::vector h3{hxz, hyz, hzz}; + return std::vector>{h1, h2, h3}; + } + // If we're on the edge, handle it specially + // since if the SDF is build with a virtual border, these cells will + // get zero gradient from this approach! + else if (enable_edge_gradients) + { + // Get the "best" indices we can use + const int64_t low_x_index = std::max((int64_t)0, x_index - 1); + const int64_t high_x_index = std::min(GetNumXCells() - 1, x_index + 1); + const int64_t low_y_index = std::max((int64_t)0, y_index - 1); + const int64_t high_y_index = std::min(GetNumYCells() - 1, y_index + 1); + const int64_t low_z_index = std::max((int64_t)0, z_index - 1); + const int64_t high_z_index = std::min(GetNumZCells() - 1, z_index + 1); + // Compute the axis increments + const double x_increment + = (high_x_index - low_x_index); + const double y_increment + = (high_y_index - low_y_index); + const double z_increment + = (high_z_index - low_z_index); + const double inv_resolution = 1.0 / GetResolution(); + + // Compute the hessians elements by default these are zero + double hxx = 0.0; + double hyy = 0.0; + double hzz = 0.0; + double hxy = 0.0; + double hxz = 0.0; + double hyz = 0.0; + + // Only if (x+1, x, x-1) are all within the bounds do we compute the hessian + if (x_increment > 1) + { + hxx = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index + 1, y_index, z_index).first + + GetImmutable(x_index - 1, y_index, z_index).first) + * inv_resolution * inv_resolution; + if (y_increment > 1){ + hxy = (GetImmutable(x_index + 1, y_index + 1, z_index).first + - GetImmutable(x_index + 1, y_index - 1, z_index).first + - GetImmutable(x_index - 1, y_index + 1, z_index).first + + GetImmutable(x_index - 1, y_index - 1, z_index).first) + * 0.25 * inv_resolution * inv_resolution; + } + if (z_increment > 1){ + hxz = (GetImmutable(x_index + 1, y_index, z_index+1).first + - GetImmutable(x_index + 1, y_index, z_index - 1).first + - GetImmutable(x_index - 1, y_index, z_index + 1).first + + GetImmutable(x_index - 1, y_index, z_index - 1).first) + * 0.25 * inv_resolution * inv_resolution; + } + } + if (y_increment > 1) + { + hyy = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index, y_index+1, z_index).first + + GetImmutable(x_index, y_index-1, z_index).first) + * inv_resolution * inv_resolution; + if (z_increment > 1){ + hyz = (GetImmutable(x_index, y_index + 1, z_index + 1).first + - GetImmutable(x_index, y_index + 1, z_index - 1).first + - GetImmutable(x_index, y_index - 1, z_index + 1).first + + GetImmutable(x_index, y_index - 1, z_index - 1).first) + * 0.25 * inv_resolution * inv_resolution; + } + } + if (z_increment > 1) + { + hzz = (-2 * GetImmutable(x_index, y_index, z_index).first + + GetImmutable(x_index, y_index, z_index+1).first + + GetImmutable(x_index, y_index, z_index-1).first) + * inv_resolution * inv_resolution; + } + // Assemble and return the computed hessian + std::vector h1{hxx, hxy, hxz}; + std::vector h2{hxy, hyy, hyz}; + std::vector h3{hxz, hyz, hzz}; + return std::vector>{h1, h2, h3}; // h1, h2, h3 are rows + } + // Edge gradient disabled, return no hessian + else + { + return std::vector>(); + } + } + // If we're out of bounds, return no hessian + else + { + return std::vector>(); + } + } + inline std::vector GetGradient( const int64_t x_index, const int64_t y_index, const int64_t z_index, const bool enable_edge_gradients = false) const diff --git a/src/sdf_tools/bindings.cpp b/src/sdf_tools/bindings.cpp index b5ed40a..1f8f8b0 100644 --- a/src/sdf_tools/bindings.cpp +++ b/src/sdf_tools/bindings.cpp @@ -31,6 +31,10 @@ PYBIND11_MODULE(pysdf_tools, m) { std::vector(SignedDistanceField::* get_gradient_1)(int64_t, int64_t, int64_t, bool) const = &SignedDistanceField::GetGradient; + std::vector>(SignedDistanceField::* + get_hessian)(int64_t, int64_t, int64_t, bool) const = + &SignedDistanceField::GetHessian; + std::pair(SignedDistanceField::* sdf_get_value_by_coordinates)(double, double, double) const = &SignedDistanceField::GetImmutable; @@ -40,9 +44,13 @@ PYBIND11_MODULE(pysdf_tools, m) { .def(py::init<>()) .def("GetRawData", &SignedDistanceField::GetImmutableRawData, "Please don't mutate this") .def("GetFullGradient", &SignedDistanceField::GetFullGradient) + .def("GetFullHessian", &SignedDistanceField::GetFullHessian) .def("GetGradient", get_gradient_1, "get the gradient based on index", py::arg("x_index"), py::arg("y_index"), py::arg("z_index"), py::arg("enable_edge_gradients") = false) + .def("GetHessian", get_hessian, "get the hessian based on index", py::arg("x_index"), + py::arg("y_index"), + py::arg("z_index"), py::arg("enable_edge_gradients") = false) .def("GetMessageRepresentation", &SignedDistanceField::GetMessageRepresentation) .def("LoadFromMessageRepresentation", &SignedDistanceField::LoadFromMessageRepresentation) .def("SaveToFile", &SignedDistanceField::SaveToFile) @@ -111,4 +119,28 @@ PYBIND11_MODULE(pysdf_tools, m) { .def("DeserializeSelf", &VoxelGridVecd::DeserializeSelf, "deserialize", py::arg("buffer"), py::arg("current"), py::arg("value_deserializer")); + + using VoxelGridMatd = VoxelGrid::VoxelGrid>>; + std::pair> const &, bool>(VoxelGridMatd::* + grad_get_value_by_coordinates_mat)(double, double, double) const = &VoxelGridMatd::GetImmutable; + + std::pair> const &, bool>(VoxelGridMatd::* + grad_get_value_by_index_mat)(int64_t, int64_t, int64_t) const = &VoxelGridMatd::GetImmutable; + + py::class_(m, "VoxelGridMat") + .def(py::init<>()) + .def("GetRawData", &VoxelGridMatd::GetImmutableRawData, "Please don't mutate this") + .def("GetNumXCells", &VoxelGridMatd::GetNumXCells) + .def("GetNumYCells", &VoxelGridMatd::GetNumYCells) + .def("GetNumZCells", &VoxelGridMatd::GetNumZCells) + .def("GetValueByCoordinates", grad_get_value_by_coordinates_mat, "Please don't mutate this", py::arg("x"), + py::arg("y"), + py::arg("z")) + .def("GetValueByIndex", grad_get_value_by_index_mat, "Please don't mutate this", py::arg("x_index"), + py::arg("y_index"), + py::arg("z_index")) + .def("SerializeSelf", &VoxelGridMatd::SerializeSelf) + .def("DeserializeSelf", &VoxelGridMatd::DeserializeSelf, "deserialize", py::arg("buffer"), + py::arg("current"), py::arg("value_deserializer")); + } diff --git a/src/sdf_tools/utils_3d.py b/src/sdf_tools/utils_3d.py index 447e341..472e09c 100644 --- a/src/sdf_tools/utils_3d.py +++ b/src/sdf_tools/utils_3d.py @@ -2,12 +2,14 @@ import pysdf_tools -def compute_sdf_and_gradient(env, res, origin_point): +def compute_sdf_and_gradient(env, res, origin_point, compute_hessian=False): """ :param env: [w,h,c] :param res: float in meters :param origin_point: [3], [x,y,z] - :return: a tuple (sdf, sdf_gradient) as numpy arrays + :param compute_hessian: bool + :return: a tuple (sdf, sdf_gradient) as numpy arrays if compute_hessian is False, + else (sdf, sdf_gradient, sdf_hessian) """ origin_transform = pysdf_tools.Isometry3d([ @@ -44,6 +46,8 @@ def gradient_function(x_index, y_index, z_index, enable_edge_gradients=False): return sdf_voxelgrid.GetGradient(x_index, y_index, z_index, enable_edge_gradients) gradient_voxelgrid = sdf_voxelgrid.GetFullGradient(gradient_function, True) + + # reshape the gradient gradient = gradient_voxelgrid.GetRawData() np_gradient = np.array(gradient) d = np_gradient.shape[1] @@ -53,8 +57,34 @@ def gradient_function(x_index, y_index, z_index, enable_edge_gradients=False): d] np_gradient = np_gradient.reshape(grad_shape) np_gradient = np.transpose(np_gradient, [1, 0, 2, 3]).astype(np.float32) + # For same reason as column major vs row major, we need to swap the components of the gradients themselves + np_gradient[:, :, :, [0, 1]] = np_gradient[:, :, :, [1, 0]] np_gradient = np_gradient.astype(np.float32) + if compute_hessian: + print('computing hessian') + def hessian_function(x_index, y_index, z_index, enable_edge_gradients=False): + return sdf_voxelgrid.GetHessian(x_index, y_index, z_index, enable_edge_gradients) + + hessian_voxelgrid = sdf_voxelgrid.GetFullHessian(hessian_function, True) + hessian = hessian_voxelgrid.GetRawData() + np_hessian = np.array(hessian) + hessian_shape = [gradient_voxelgrid.GetNumXCells(), + gradient_voxelgrid.GetNumYCells(), + gradient_voxelgrid.GetNumZCells(), + d, d] + # reshape hessian + np_hessian = np_hessian.reshape(hessian_shape) + # as above, because of column major / row major issue need to swap elements of hessian + i, j, k = np.random.randint(low=0, high=32, size=(3,)) + print(np_hessian[i, j, k]) + np_hessian[:, :, :, :, [0, 1]] = np_hessian[:, :, :, :, [1, 0]] + np_hessian[:, :, :, [0, 1]] = np_hessian[:, :, :, [1, 0]] + print(np_hessian[i, j, k]) + + np_hessian.astype(np.float32) + + return np_sdf, np_gradient, np_hessian # # DEBUGGING # import matplotlib.pyplot as plt # plt.figure()