in pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp [48:151]
void PointsToVolumesForwardCpu(
const torch::Tensor& points_3d,
const torch::Tensor& points_features,
const torch::Tensor& volume_densities,
const torch::Tensor& volume_features,
const torch::Tensor& grid_sizes,
const torch::Tensor& mask,
const float point_weight,
const bool align_corners,
const bool splat) {
const int64_t batch_size = points_3d.size(0);
const int64_t P = points_3d.size(1);
const int64_t n_features = points_features.size(2);
// We unify the formula for the location of px in the comment above as
// ({(px + 1) * 0.5} * (grid_size_x-scale_offset)) - offset.
const int scale_offset = align_corners ? 1 : 0;
const float offset = align_corners ? 0 : 0.5;
auto points_3d_a = points_3d.accessor<float, 3>();
auto points_features_a = points_features.accessor<float, 3>();
auto volume_densities_a = volume_densities.accessor<float, 5>();
auto volume_features_a = volume_features.accessor<float, 5>();
auto grid_sizes_a = grid_sizes.accessor<int64_t, 2>();
auto mask_a = mask.accessor<float, 2>();
// For each batch element
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
auto points_3d_aa = points_3d_a[batch_idx];
auto points_features_aa = points_features_a[batch_idx];
auto volume_densities_aa = volume_densities_a[batch_idx][0];
auto volume_features_aa = volume_features_a[batch_idx];
auto grid_sizes_aa = grid_sizes_a[batch_idx];
auto mask_aa = mask_a[batch_idx];
const int64_t grid_size_x = grid_sizes_aa[2];
const int64_t grid_size_y = grid_sizes_aa[1];
const int64_t grid_size_z = grid_sizes_aa[0];
// For each point
for (int64_t point_idx = 0; point_idx < P; ++point_idx) {
// Ignore point if mask is 0
if (mask_aa[point_idx] == 0) {
continue;
}
auto point = points_3d_aa[point_idx];
auto point_features = points_features_aa[point_idx];
// Define how to increment a location in the volume by an amount. The need
// for this depends on the interpolation method:
// once per point for nearest, eight times for splat.
auto increment_location =
[&](int64_t x, int64_t y, int64_t z, float weight) {
if (x >= grid_size_x || y >= grid_size_y || z >= grid_size_z) {
return;
}
if (x < 0 || y < 0 || z < 0) {
return;
}
volume_densities_aa[z][y][x] += weight * point_weight;
for (int64_t feature_idx = 0; feature_idx < n_features;
++feature_idx) {
volume_features_aa[feature_idx][z][y][x] +=
point_features[feature_idx] * weight * point_weight;
}
};
if (!splat) {
// Increment the location nearest the point.
long x = std::lround(
(point[0] + 1) * 0.5 * (grid_size_x - scale_offset) - offset);
long y = std::lround(
(point[1] + 1) * 0.5 * (grid_size_y - scale_offset) - offset);
long z = std::lround(
(point[2] + 1) * 0.5 * (grid_size_z - scale_offset) - offset);
increment_location(x, y, z, 1);
} else {
// There are 8 locations around the point which we need to worry about.
// Their coordinates are (x or x+1, y or y+1, z or z+1).
// rx is a number between 0 and 1 for the proportion in the x direction:
// rx==0 means weight all on the lower bound, x, rx=1-eps means most
// weight on x+1. Ditto for ry and yz.
float x = 0, y = 0, z = 0;
float rx = std::modf(
(point[0] + 1) * 0.5 * (grid_size_x - scale_offset) - offset, &x);
float ry = std::modf(
(point[1] + 1) * 0.5 * (grid_size_y - scale_offset) - offset, &y);
float rz = std::modf(
(point[2] + 1) * 0.5 * (grid_size_z - scale_offset) - offset, &z);
// Define how to fractionally increment one of the 8 locations around
// the point.
auto handle_point = [&](bool up_x, bool up_y, bool up_z) {
float weight = (up_x ? rx : 1 - rx) * (up_y ? ry : 1 - ry) *
(up_z ? rz : 1 - rz);
increment_location(x + up_x, y + up_y, z + up_z, weight);
};
// and do so.
EightDirections(handle_point);
}
}
}
}