366 lines
14 KiB
C
366 lines
14 KiB
C
|
// ----------------------------------------------------------------------------
|
||
|
|
// - Open3D: www.open3d.org -
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
// Copyright (c) 2018-2023 www.open3d.org
|
||
|
|
// SPDX-License-Identifier: MIT
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include <unordered_map>
|
||
|
|
|
||
|
|
#include "open3d/core/CUDAUtils.h"
|
||
|
|
#include "open3d/core/Tensor.h"
|
||
|
|
#include "open3d/core/TensorCheck.h"
|
||
|
|
#include "open3d/utility/Helper.h"
|
||
|
|
#include "open3d/utility/Logging.h"
|
||
|
|
#include "open3d/utility/Timer.h"
|
||
|
|
|
||
|
|
namespace open3d {
|
||
|
|
namespace t {
|
||
|
|
namespace geometry {
|
||
|
|
namespace kernel {
|
||
|
|
|
||
|
|
/// Helper class for converting coordinates/indices between 3D/3D, 3D/2D, 2D/3D.
|
||
|
|
class TransformIndexer {
|
||
|
|
public:
|
||
|
|
/// intrinsic: simple pinhole camera matrix, stored in fx, fy, cx, cy
|
||
|
|
/// extrinsic: world to camera transform, stored in a 3x4 matrix
|
||
|
|
TransformIndexer(const core::Tensor& intrinsics,
|
||
|
|
const core::Tensor& extrinsics,
|
||
|
|
float scale = 1.0f) {
|
||
|
|
core::AssertTensorShape(intrinsics, {3, 3});
|
||
|
|
core::AssertTensorDtype(intrinsics, core::Float64);
|
||
|
|
core::AssertTensorDevice(intrinsics, core::Device("CPU:0"));
|
||
|
|
if (!intrinsics.IsContiguous()) {
|
||
|
|
utility::LogError("Intrinsics is not contiguous");
|
||
|
|
}
|
||
|
|
|
||
|
|
core::AssertTensorShape(extrinsics, {4, 4});
|
||
|
|
core::AssertTensorDtype(extrinsics, core::Float64);
|
||
|
|
core::AssertTensorDevice(extrinsics, core::Device("CPU:0"));
|
||
|
|
if (!extrinsics.IsContiguous()) {
|
||
|
|
utility::LogError("Extrinsics is not contiguous");
|
||
|
|
}
|
||
|
|
|
||
|
|
const double* intrinsic_ptr = intrinsics.GetDataPtr<double>();
|
||
|
|
const double* extrinsic_ptr = extrinsics.GetDataPtr<double>();
|
||
|
|
for (int i = 0; i < 3; ++i) {
|
||
|
|
for (int j = 0; j < 4; ++j) {
|
||
|
|
extrinsic_[i][j] = extrinsic_ptr[i * 4 + j];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fx_ = intrinsic_ptr[0 * 3 + 0];
|
||
|
|
fy_ = intrinsic_ptr[1 * 3 + 1];
|
||
|
|
cx_ = intrinsic_ptr[0 * 3 + 2];
|
||
|
|
cy_ = intrinsic_ptr[1 * 3 + 2];
|
||
|
|
scale_ = scale;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Transform a 3D coordinate in camera coordinate to world coordinate
|
||
|
|
OPEN3D_HOST_DEVICE void RigidTransform(float x_in,
|
||
|
|
float y_in,
|
||
|
|
float z_in,
|
||
|
|
float* x_out,
|
||
|
|
float* y_out,
|
||
|
|
float* z_out) const {
|
||
|
|
x_in *= scale_;
|
||
|
|
y_in *= scale_;
|
||
|
|
z_in *= scale_;
|
||
|
|
|
||
|
|
*x_out = x_in * extrinsic_[0][0] + y_in * extrinsic_[0][1] +
|
||
|
|
z_in * extrinsic_[0][2] + extrinsic_[0][3];
|
||
|
|
*y_out = x_in * extrinsic_[1][0] + y_in * extrinsic_[1][1] +
|
||
|
|
z_in * extrinsic_[1][2] + extrinsic_[1][3];
|
||
|
|
*z_out = x_in * extrinsic_[2][0] + y_in * extrinsic_[2][1] +
|
||
|
|
z_in * extrinsic_[2][2] + extrinsic_[2][3];
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Transform a 3D coordinate in camera coordinate to world coordinate
|
||
|
|
OPEN3D_HOST_DEVICE void Rotate(float x_in,
|
||
|
|
float y_in,
|
||
|
|
float z_in,
|
||
|
|
float* x_out,
|
||
|
|
float* y_out,
|
||
|
|
float* z_out) const {
|
||
|
|
x_in *= scale_;
|
||
|
|
y_in *= scale_;
|
||
|
|
z_in *= scale_;
|
||
|
|
|
||
|
|
*x_out = x_in * extrinsic_[0][0] + y_in * extrinsic_[0][1] +
|
||
|
|
z_in * extrinsic_[0][2];
|
||
|
|
*y_out = x_in * extrinsic_[1][0] + y_in * extrinsic_[1][1] +
|
||
|
|
z_in * extrinsic_[1][2];
|
||
|
|
*z_out = x_in * extrinsic_[2][0] + y_in * extrinsic_[2][1] +
|
||
|
|
z_in * extrinsic_[2][2];
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Project a 3D coordinate in camera coordinate to a 2D uv coordinate
|
||
|
|
OPEN3D_HOST_DEVICE void Project(float x_in,
|
||
|
|
float y_in,
|
||
|
|
float z_in,
|
||
|
|
float* u_out,
|
||
|
|
float* v_out) const {
|
||
|
|
float inv_z = 1.0f / z_in;
|
||
|
|
*u_out = fx_ * x_in * inv_z + cx_;
|
||
|
|
*v_out = fy_ * y_in * inv_z + cy_;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Unproject a 2D uv coordinate with depth to 3D in camera coordinate
|
||
|
|
OPEN3D_HOST_DEVICE void Unproject(float u_in,
|
||
|
|
float v_in,
|
||
|
|
float d_in,
|
||
|
|
float* x_out,
|
||
|
|
float* y_out,
|
||
|
|
float* z_out) const {
|
||
|
|
*x_out = (u_in - cx_) * d_in / fx_;
|
||
|
|
*y_out = (v_in - cy_) * d_in / fy_;
|
||
|
|
*z_out = d_in;
|
||
|
|
}
|
||
|
|
|
||
|
|
OPEN3D_HOST_DEVICE void GetFocalLength(float* fx, float* fy) const {
|
||
|
|
*fx = fx_;
|
||
|
|
*fy = fy_;
|
||
|
|
}
|
||
|
|
|
||
|
|
OPEN3D_HOST_DEVICE void GetCameraPosition(float* x,
|
||
|
|
float* y,
|
||
|
|
float* z) const {
|
||
|
|
*x = extrinsic_[0][3];
|
||
|
|
*y = extrinsic_[1][3];
|
||
|
|
*z = extrinsic_[2][3];
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
float extrinsic_[3][4];
|
||
|
|
|
||
|
|
float fx_;
|
||
|
|
float fy_;
|
||
|
|
float cx_;
|
||
|
|
float cy_;
|
||
|
|
|
||
|
|
float scale_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// Convert between ND coordinates and their corresponding linear offsets.
|
||
|
|
/// Input ndarray tensor must be contiguous.
|
||
|
|
/// Internal shape conversions:
|
||
|
|
/// 1D: index (x), [channel (c)]
|
||
|
|
/// 2D: height (y), weight (x), [channel (c)]
|
||
|
|
/// 3D: depth (z), height (y), width (x), [channel (c)]
|
||
|
|
/// 4D: time (t), depth (z), height (y), width (x), [channel (c)]
|
||
|
|
/// External indexing order:
|
||
|
|
/// 1D: x
|
||
|
|
/// 2D: x, y
|
||
|
|
/// 3D: x, y, z
|
||
|
|
/// 4D: x, y, z, t
|
||
|
|
const int64_t MAX_RESOLUTION_DIMS = 4;
|
||
|
|
|
||
|
|
template <typename index_t>
|
||
|
|
class TArrayIndexer {
|
||
|
|
public:
|
||
|
|
TArrayIndexer() : ptr_(nullptr), element_byte_size_(0), active_dims_(0) {
|
||
|
|
for (index_t i = 0; i < MAX_RESOLUTION_DIMS; ++i) {
|
||
|
|
shape_[i] = 0;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
TArrayIndexer(const core::Tensor& ndarray, index_t active_dims) {
|
||
|
|
if (!ndarray.IsContiguous()) {
|
||
|
|
utility::LogError(
|
||
|
|
"Only support contiguous tensors for general operations.");
|
||
|
|
}
|
||
|
|
|
||
|
|
core::SizeVector shape = ndarray.GetShape();
|
||
|
|
index_t n = ndarray.NumDims();
|
||
|
|
if (active_dims > MAX_RESOLUTION_DIMS || active_dims > n) {
|
||
|
|
utility::LogError(
|
||
|
|
"Tensor shape too large, only <= {} and <= {} array dim is "
|
||
|
|
"supported, but received {}.",
|
||
|
|
MAX_RESOLUTION_DIMS, n, active_dims);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Leading dimensions are coordinates
|
||
|
|
active_dims_ = active_dims;
|
||
|
|
for (index_t i = 0; i < active_dims_; ++i) {
|
||
|
|
shape_[i] = shape[i];
|
||
|
|
}
|
||
|
|
// Trailing dimensions are channels
|
||
|
|
element_byte_size_ = ndarray.GetDtype().ByteSize();
|
||
|
|
for (index_t i = active_dims_; i < n; ++i) {
|
||
|
|
element_byte_size_ *= shape[i];
|
||
|
|
}
|
||
|
|
|
||
|
|
// Fill-in rest to make compiler happy, not actually used.
|
||
|
|
for (index_t i = active_dims_; i < MAX_RESOLUTION_DIMS; ++i) {
|
||
|
|
shape_[i] = 0;
|
||
|
|
}
|
||
|
|
ptr_ = const_cast<void*>(ndarray.GetDataPtr());
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Only used for simple shapes
|
||
|
|
TArrayIndexer(const core::SizeVector& shape) {
|
||
|
|
index_t n = static_cast<index_t>(shape.size());
|
||
|
|
if (n > MAX_RESOLUTION_DIMS) {
|
||
|
|
utility::LogError(
|
||
|
|
"SizeVector too large, only <= {} is supported, but "
|
||
|
|
"received {}.",
|
||
|
|
MAX_RESOLUTION_DIMS, n);
|
||
|
|
}
|
||
|
|
active_dims_ = n;
|
||
|
|
for (index_t i = 0; i < active_dims_; ++i) {
|
||
|
|
shape_[i] = shape[i];
|
||
|
|
}
|
||
|
|
|
||
|
|
// Fill-in rest to make compiler happy, not actually used.
|
||
|
|
for (index_t i = active_dims_; i < MAX_RESOLUTION_DIMS; ++i) {
|
||
|
|
shape_[i] = 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reserved
|
||
|
|
element_byte_size_ = 0;
|
||
|
|
ptr_ = nullptr;
|
||
|
|
}
|
||
|
|
|
||
|
|
OPEN3D_HOST_DEVICE index_t ElementByteSize() { return element_byte_size_; }
|
||
|
|
|
||
|
|
OPEN3D_HOST_DEVICE index_t NumElements() {
|
||
|
|
index_t num_elems = 1;
|
||
|
|
for (index_t i = 0; i < active_dims_; ++i) {
|
||
|
|
num_elems *= shape_[i];
|
||
|
|
}
|
||
|
|
return num_elems;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 2D coordinate => workload
|
||
|
|
inline OPEN3D_HOST_DEVICE void CoordToWorkload(index_t x_in,
|
||
|
|
index_t y_in,
|
||
|
|
index_t* workload) const {
|
||
|
|
*workload = y_in * shape_[1] + x_in;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 3D coordinate => workload
|
||
|
|
inline OPEN3D_HOST_DEVICE void CoordToWorkload(index_t x_in,
|
||
|
|
index_t y_in,
|
||
|
|
index_t z_in,
|
||
|
|
index_t* workload) const {
|
||
|
|
*workload = (z_in * shape_[1] + y_in) * shape_[2] + x_in;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 4D coordinate => workload
|
||
|
|
inline OPEN3D_HOST_DEVICE void CoordToWorkload(index_t x_in,
|
||
|
|
index_t y_in,
|
||
|
|
index_t z_in,
|
||
|
|
index_t t_in,
|
||
|
|
index_t* workload) const {
|
||
|
|
*workload = ((t_in * shape_[1] + z_in) * shape_[2] + y_in) * shape_[3] +
|
||
|
|
x_in;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Workload => 2D coordinate
|
||
|
|
inline OPEN3D_HOST_DEVICE void WorkloadToCoord(index_t workload,
|
||
|
|
index_t* x_out,
|
||
|
|
index_t* y_out) const {
|
||
|
|
*x_out = workload % shape_[1];
|
||
|
|
*y_out = workload / shape_[1];
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Workload => 3D coordinate
|
||
|
|
inline OPEN3D_HOST_DEVICE void WorkloadToCoord(index_t workload,
|
||
|
|
index_t* x_out,
|
||
|
|
index_t* y_out,
|
||
|
|
index_t* z_out) const {
|
||
|
|
*x_out = workload % shape_[2];
|
||
|
|
workload = (workload - *x_out) / shape_[2];
|
||
|
|
*y_out = workload % shape_[1];
|
||
|
|
*z_out = workload / shape_[1];
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Workload => 4D coordinate
|
||
|
|
inline OPEN3D_HOST_DEVICE void WorkloadToCoord(index_t workload,
|
||
|
|
index_t* x_out,
|
||
|
|
index_t* y_out,
|
||
|
|
index_t* z_out,
|
||
|
|
index_t* t_out) const {
|
||
|
|
*x_out = workload % shape_[3];
|
||
|
|
workload = (workload - *x_out) / shape_[3];
|
||
|
|
*y_out = workload % shape_[2];
|
||
|
|
workload = (workload - *y_out) / shape_[2];
|
||
|
|
*z_out = workload % shape_[1];
|
||
|
|
*t_out = workload / shape_[1];
|
||
|
|
}
|
||
|
|
|
||
|
|
inline OPEN3D_HOST_DEVICE bool InBoundary(float x, float y) const {
|
||
|
|
return y >= 0 && x >= 0 && y <= shape_[0] - 1.0f &&
|
||
|
|
x <= shape_[1] - 1.0f;
|
||
|
|
}
|
||
|
|
inline OPEN3D_HOST_DEVICE bool InBoundary(float x, float y, float z) const {
|
||
|
|
return z >= 0 && y >= 0 && x >= 0 && z <= shape_[0] - 1.0f &&
|
||
|
|
y <= shape_[1] - 1.0f && x <= shape_[2] - 1.0f;
|
||
|
|
}
|
||
|
|
inline OPEN3D_HOST_DEVICE bool InBoundary(float x,
|
||
|
|
float y,
|
||
|
|
float z,
|
||
|
|
float t) const {
|
||
|
|
return t >= 0 && z >= 0 && y >= 0 && x >= 0 && t <= shape_[0] - 1.0f &&
|
||
|
|
z <= shape_[1] - 1.0f && y <= shape_[2] - 1.0f &&
|
||
|
|
x <= shape_[3] - 1.0f;
|
||
|
|
}
|
||
|
|
|
||
|
|
inline OPEN3D_HOST_DEVICE index_t GetShape(int i) const {
|
||
|
|
return shape_[i];
|
||
|
|
}
|
||
|
|
|
||
|
|
inline OPEN3D_HOST_DEVICE void* GetDataPtr() const { return ptr_; }
|
||
|
|
|
||
|
|
template <typename T>
|
||
|
|
inline OPEN3D_HOST_DEVICE T* GetDataPtr(index_t x) const {
|
||
|
|
return static_cast<T*>(static_cast<void*>(static_cast<uint8_t*>(ptr_) +
|
||
|
|
x * element_byte_size_));
|
||
|
|
}
|
||
|
|
|
||
|
|
template <typename T>
|
||
|
|
inline OPEN3D_HOST_DEVICE T* GetDataPtr(index_t x, index_t y) const {
|
||
|
|
index_t workload;
|
||
|
|
CoordToWorkload(x, y, &workload);
|
||
|
|
return static_cast<T*>(static_cast<void*>(
|
||
|
|
static_cast<uint8_t*>(ptr_) + workload * element_byte_size_));
|
||
|
|
}
|
||
|
|
|
||
|
|
template <typename T>
|
||
|
|
inline OPEN3D_HOST_DEVICE T* GetDataPtr(index_t x,
|
||
|
|
index_t y,
|
||
|
|
index_t z) const {
|
||
|
|
index_t workload;
|
||
|
|
CoordToWorkload(x, y, z, &workload);
|
||
|
|
return static_cast<T*>(static_cast<void*>(
|
||
|
|
static_cast<uint8_t*>(ptr_) + workload * element_byte_size_));
|
||
|
|
}
|
||
|
|
|
||
|
|
template <typename T>
|
||
|
|
inline OPEN3D_HOST_DEVICE T* GetDataPtr(index_t x,
|
||
|
|
index_t y,
|
||
|
|
index_t z,
|
||
|
|
index_t t) const {
|
||
|
|
index_t workload;
|
||
|
|
CoordToWorkload(x, y, z, t, &workload);
|
||
|
|
return static_cast<T*>(static_cast<void*>(
|
||
|
|
static_cast<uint8_t*>(ptr_) + workload * element_byte_size_));
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
void* ptr_;
|
||
|
|
index_t element_byte_size_;
|
||
|
|
index_t active_dims_;
|
||
|
|
|
||
|
|
index_t shape_[MAX_RESOLUTION_DIMS];
|
||
|
|
};
|
||
|
|
|
||
|
|
using NDArrayIndexer = TArrayIndexer<int64_t>;
|
||
|
|
|
||
|
|
} // namespace kernel
|
||
|
|
} // namespace geometry
|
||
|
|
} // namespace t
|
||
|
|
} // namespace open3d
|