Skip to content

Commit

Permalink
Simplify PR (only softmax blending, no color correction)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ssheorey
ssheorey committed Jan 1, 2025
1 parent 35df6e5 commit 7b7316e
Show file tree
Hide file tree
Showing 5 changed files with 33 additions and 351 deletions.
264 changes: 25 additions & 239 deletions cpp/open3d/t/geometry/TriangleMesh.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1420,80 +1420,6 @@ core::Tensor Project(const core::Tensor &t_xyz, // contiguous {...,3}
return t_xy; // contiguous {...,2}
}

/// Estimate contrast and brightness in each color channel for this_albedo to
/// match albedo, based on the overlapping area in the texture image. Contrast
/// and brightness are estimated by matching the second and first moments
/// (variance and mean) of the pixel colors, respectively.
/// albedo and this_albedo are float images with range [0,1]
std::tuple<std::array<float, 3>, std::array<float, 3>> get_color_correction(
const core::Tensor &albedo,
const core::Tensor &this_albedo,
int image_id) {
const float MIN_PIXEL_WEIGHT = 0.001;
const float MIN_COLOR_VAR = 0.001; // stddev of 8 out of 255
const float shift = 0.5f; // compute shifted sum / sumsqr for stability.
const unsigned MIN_OVERLAP_PIXELS = 1024;
// Perform the color correction with albedo downsampled to size 256x256.
float resize_down = 256.f / albedo.GetShape(0);
std::array<float, 3> this_contrast{1.f, 1.f, 1.f},
this_brightness{0.f, 0.f, 0.f}, sum{}, this_sum{}, sumsqr{},
this_sumsqr{};
auto q_albedo = Image(albedo)
.Resize(resize_down, Image::InterpType::Linear)
.AsTensor();
auto q_this_albedo = Image(this_albedo)
.Resize(resize_down, Image::InterpType::Linear)
.AsTensor();
unsigned count = 0;
for (float *pq_albedo = q_albedo.GetDataPtr<float>(),
*pq_this_albedo = q_this_albedo.GetDataPtr<float>();
pq_albedo < q_albedo.GetDataPtr<float>() + q_albedo.NumElements();
pq_albedo += 4, pq_this_albedo += 4) {
if (pq_albedo[3] <= MIN_PIXEL_WEIGHT ||
pq_this_albedo[3] <= MIN_PIXEL_WEIGHT)
continue;
++count;
for (int c = 0; c < 3; ++c) {
float update = pq_albedo[c] - shift;
sum[c] += update;
sumsqr[c] += update * update;
update = pq_this_albedo[c] - shift;
this_sum[c] += update;
this_sumsqr[c] += update * update;
}
}
if (count <= MIN_OVERLAP_PIXELS) {
if (image_id > 0) {
utility::LogWarning(
"[ProjectImagesToAlbedo] Too few overlapping pixels "
"({}/{}) found for color correction in image {}.",
count, MIN_OVERLAP_PIXELS, image_id);
}
return std::make_tuple(this_contrast, this_brightness);
}
for (int c = 0; c < 3; ++c) {
float variance = (sumsqr[c] - sum[c] * sum[c] / count) / count,
this_variance =
(this_sumsqr[c] - this_sum[c] * this_sum[c] / count) /
count;
utility::LogDebug("count: {}, variance: {}, this_variance: {}", count,
variance, this_variance);
if (this_variance < MIN_COLOR_VAR) {
utility::LogWarning(
"[ProjectImagesToAlbedo] Overlapping part of image {} is "
"too flat for color correction.",
image_id);
return std::make_tuple(this_contrast, this_brightness);
}
this_contrast[c] = sqrt((variance + MIN_COLOR_VAR) /
(this_variance + MIN_COLOR_VAR));
sum[c] += count * shift; // get un-shifted sum for brightness.
this_sum[c] += count * shift;
this_brightness[c] = (sum[c] - this_contrast[c] * this_sum[c]) / count;
}
return std::make_tuple(this_contrast, this_brightness);
}

/// Estimate minimum sqr distance from a set of points to a set of cameras.
float get_min_cam_sqrdistance(
const core::Tensor &positions,
Expand Down Expand Up @@ -1526,33 +1452,15 @@ Image TriangleMesh::ProjectImagesToAlbedo(
const std::vector<core::Tensor> &intrinsic_matrices,
const std::vector<core::Tensor> &extrinsic_matrices,
int tex_size /*=1024*/,
bool update_material /*=true*/,
BlendingMethod blending_method /*=MAX*/
) {
const bool DEBUG = true;
bool update_material /*=true*/) {
using core::None;
using tk = core::TensorKey;
constexpr float EPS = 1e-6;
// softmax_shift is used to prevent overflow in the softmax function.
// softmax_shift is set so that max value of weighting function is exp(64),
// well within float range. (exp(89.f) is inf)
float min_sqr_distance =
blending_method & BlendingMethod::AVERAGE
? get_min_cam_sqrdistance(GetVertexPositions(),
extrinsic_matrices)
: 0.01f;
float softmax_shift = 10.f, softmax_scale = 20 * min_sqr_distance;
utility::LogInfo("softmax_shift, softmax_scale: {}, {}", softmax_shift,
softmax_scale);
if (!HasTriangleAttr("texture_uvs")) {
utility::LogError(
"TriangleMesh does not contain 'texture_uvs'. Please compute "
"it with ComputeUVAtlas() first.");
}
if ((blending_method & (BlendingMethod::MAX | BlendingMethod::AVERAGE)) ==
0) {
utility::LogError("Select one of MAX and AVERAGE BlendingMethod s.");
}
core::Tensor texture_uvs =
GetTriangleAttr("texture_uvs").To(core::Device()).Contiguous();
core::AssertTensorShape(texture_uvs, {core::None, 3, 2});
Expand All @@ -1566,30 +1474,25 @@ Image TriangleMesh::ProjectImagesToAlbedo(
images.size(), extrinsic_matrices.size(),
intrinsic_matrices.size());
}

// softmax_shift is used to prevent overflow in the softmax function.
// softmax_shift is set so that max value of weighting function is exp(64),
// well within float range. (exp(89.f) is inf)
float min_sqr_distance =
get_min_cam_sqrdistance(GetVertexPositions(), extrinsic_matrices);
float softmax_shift = 10.f, softmax_scale = 20 * min_sqr_distance;
// (u,v) -> (x,y,z) : {tex_size, tex_size, 3}
core::Tensor position_map = BakeVertexAttrTextures(
tex_size, {"positions"}, 1, 0, false)["positions"];
/* if (DEBUG) { */
/* io::WriteImage("position_map.png", */
/* Image(((position_map + 1) * 127.5).To(core::UInt8)));
*/
/* } */
core::Tensor albedo = core::Tensor::Zeros({tex_size, tex_size, 4},
core::Float32),
albedo_max = albedo; // same data
albedo.Slice(2, 3, 4).Fill(EPS); // regularize
if (blending_method &
(BlendingMethod::MAX | BlendingMethod::COLOR_CORRECTION))
albedo_max = albedo.Clone(); // Copy
core::Tensor albedo =
core::Tensor::Zeros({tex_size, tex_size, 4}, core::Float32);
albedo.Slice(2, 3, 4).Fill(EPS); // regularize
std::mutex albedo_mutex;

RaycastingScene rcs;
rcs.AddTriangles(*this);

// Simulate thread_local Tensors with a vector of Tensors. thread_local
// variables are only destructed when the TBB thread pool is finalized which
// can cause memory leaks in Python code / long running processes.
// Also TBB can schedule multiple tasks on one thread at the same time
// setup working data for each task.
size_t max_workers = tbb::this_task_arena::max_concurrency();
// Tensor copy ctor does shallow copies - OK for empty tensors.
std::vector<core::Tensor> this_albedo(max_workers,
Expand All @@ -1598,14 +1501,9 @@ Image TriangleMesh::ProjectImagesToAlbedo(
uv2xy(max_workers, core::Tensor({}, core::Float32)),
uvrays(max_workers, core::Tensor({}, core::Float32));

// Used to control order of blending projected images into the texture. This
// ensures correct computation of color matching (only neded for
// COLOR_CORRECTION).
std::condition_variable cv_next_blend_image;
size_t next_blend_image{0};
auto project_one_image = [&](size_t i, tbb::feeder<size_t> &feeder) {
size_t widx = tbb::this_task_arena::current_thread_index();
// initialize thread variables
// initialize task variables
if (!this_albedo[widx].GetShape().IsCompatible(
{tex_size, tex_size, 4})) {
this_albedo[widx] =
Expand Down Expand Up @@ -1672,10 +1570,6 @@ Image TriangleMesh::ProjectImagesToAlbedo(
result = tbb::this_task_arena::isolate(
[&rcs, &uvrays, widx]() { return rcs.CastRays(uvrays[widx]); });
auto &t_hit_uv = result["t_hit"];
/* if (DEBUG) { */
/* io::WriteImage(fmt::format("t_hit_uv_{}.png", i), */
/* Image((t_hit_uv * 255).To(core::UInt8))); */
/* } */

Project(position_map, intrinsic_matrices[i], extrinsic_matrices[i],
uv2xy[widx]); // {ts, ts, 2}
Expand All @@ -1689,12 +1583,6 @@ Image TriangleMesh::ProjectImagesToAlbedo(
}
core::Tensor uv2xy2 =
uv2xy[widx].Permute({2, 0, 1}).Contiguous(); // {2, ts, ts}
/* if (DEBUG) { */
/* io::WriteImage(fmt::format("uv2x_{}.png", i), */
/* Image((uv2xy2[0].To(core::UInt16)))); */
/* io::WriteImage(fmt::format("uv2y_{}.png", i), */
/* Image((uv2xy2[1].To(core::UInt16)))); */
/* } */

// C. Interpolate weighted image to weighted texture
// albedo[u,v] = image[ i[u,v], j[u,v] ]
Expand All @@ -1705,127 +1593,25 @@ Image TriangleMesh::ProjectImagesToAlbedo(
this_albedo[widx], /*{texsz, texsz, 4} f32*/
t::geometry::Image::InterpType::Linear);
// Weights can become negative with higher order interpolation
float wtmin{}, wtmax{};
if (DEBUG) {
io::WriteImage(fmt::format("this_albedo_{}.png", i),
Image((this_albedo[widx].Slice(2, 0, 3) * 255)
.To(core::UInt8)));
wtmin = this_albedo[widx]
.Slice(2, 3, 4)
.Min({0, 1, 2})
.Item<float>();
wtmax = this_albedo[widx]
.Slice(2, 3, 4)
.Max({0, 1, 2})
.Item<float>();
io::WriteImage(
fmt::format("image_weights_{}.png", i),
Image(weighted_image[widx].Slice(2, 3, 4).Contiguous())
.To(core::UInt8, /*copy=*/false,
/*scale=*/255.f / (wtmax - wtmin),
/*offset=*/-wtmin * 255.f / (wtmax - wtmin)));
io::WriteImage(
fmt::format("this_albedo_weight_{}.png", i),
Image(this_albedo[widx].Slice(2, 3, 4).Contiguous())
.To(core::UInt8, /*copy=*/false,
/*scale=*/255.f / (wtmax - wtmin),
/*offset=*/-wtmin * 255.f / (wtmax - wtmin)));
}
std::array<float, 3> this_contrast{1.f, 1.f, 1.f},
this_brightness{0.f, 0.f, 0.f};

std::unique_lock<std::mutex> albedo_lock{albedo_mutex};
if (blending_method & BlendingMethod::COLOR_CORRECTION) {
// Ensure images are blended in order to correctly calculate
// color correction
cv_next_blend_image.wait(albedo_lock, [&i, &next_blend_image]() {
return next_blend_image == i;
});
io::WriteImage(fmt::format("this_albedo_overlap_{}.png", i),
Image((this_albedo[widx].Slice(2, 3, 4).Ge(1e-3) &&
albedo_max.Slice(2, 3, 4).Ge(1e-3))));
// Use unweighted albedo_max to estimate COLOR_CORRECTION. softmax
// weighted albedo does not give correct results due to under / over
// saturation of low weight regions.
std::tie(this_contrast, this_brightness) =
get_color_correction(albedo_max, this_albedo[widx], i);
utility::LogDebug(
"[ProjectImagesToAlbedo] Image {}, wtmin {}, wtmax {}, "
"contrast: {}, "
"brightness {}",
i, wtmin, wtmax, this_contrast, this_brightness);
}
if (blending_method & BlendingMethod::MAX) {
utility::LogInfo("Blending image {} with method MAX", i);
// Select albedo value with max weight
for (auto p_albedo = albedo.GetDataPtr<float>(),
p_this_albedo = this_albedo[widx].GetDataPtr<float>();
p_albedo < albedo.GetDataPtr<float>() + albedo.NumElements();
p_albedo += 4, p_this_albedo += 4) {
if (p_albedo[3] < p_this_albedo[3]) {
for (auto k = 0; k < 3; ++k)
p_albedo[k] = this_contrast[k] * p_this_albedo[k] +
this_brightness[k];
p_albedo[3] = p_this_albedo[3];
}
}
} else if (blending_method & BlendingMethod::AVERAGE) {
utility::LogInfo("Blending image {} with method AVERAGE", i);
for (auto p_albedo = albedo.GetDataPtr<float>(),
p_albedo_max = albedo_max.GetDataPtr<float>(),
p_this_albedo = this_albedo[widx].GetDataPtr<float>();
p_albedo < albedo.GetDataPtr<float>() + albedo.NumElements();
p_albedo += 4, p_albedo_max += 4, p_this_albedo += 4) {
float softmax_weight =
exp(softmax_scale * p_this_albedo[3] - softmax_shift);
for (auto k = 0; k < 3; ++k)
p_albedo[k] += (this_contrast[k] * p_this_albedo[k] +
this_brightness[k]) *
softmax_weight;
p_albedo[3] += softmax_weight;
// Update albedo_max with MAX blending
if (blending_method & BlendingMethod::COLOR_CORRECTION &&
p_albedo_max[3] < p_this_albedo[3]) {
for (auto k = 0; k < 3; ++k)
p_albedo_max[k] = this_contrast[k] * p_this_albedo[k] +
this_brightness[k];
p_albedo_max[3] = p_this_albedo[3];
}
}
}
if (DEBUG) {
io::WriteImage(fmt::format("albedo_{}.png", i),
Image((albedo.Slice(2, 0, 3) / albedo.Slice(2, 3, 4))
.Contiguous())
.To(core::UInt8, true, 255));
wtmax = albedo.Slice(2, 3, 4).Max({0, 1, 2}).Item<float>();
wtmin = albedo.Slice(2, 3, 4).Min({0, 1, 2}).Item<float>();
io::WriteImage(
fmt::format("albedo_weight_{}.png", i),
Image(albedo.Slice(2, 3, 4).Contiguous())
.To(core::UInt8, true, 255. / (wtmax - wtmin)));
utility::LogDebug("albedo weight range: {}-{}", wtmax, wtmin);
}
if (blending_method & BlendingMethod::COLOR_CORRECTION) {
cv_next_blend_image.notify_all();
if (next_blend_image + max_workers < images.size()) {
feeder.add(next_blend_image + max_workers);
}
++next_blend_image;
// ^^^ released when lambda returns.
for (auto p_albedo = albedo.GetDataPtr<float>(),
p_this_albedo = this_albedo[widx].GetDataPtr<float>();
p_albedo < albedo.GetDataPtr<float>() + albedo.NumElements();
p_albedo += 4, p_this_albedo += 4) {
float softmax_weight =
exp(softmax_scale * p_this_albedo[3] - softmax_shift);
for (auto k = 0; k < 3; ++k)
p_albedo[k] += p_this_albedo[k] * softmax_weight;
p_albedo[3] += softmax_weight;
}
};

// With COLOR_CORRECTION, we should not start more than max_workers tasks to
// avoid deadlock, since images need to be processed in order.
size_t n_init_images = blending_method & BlendingMethod::COLOR_CORRECTION
? std::min(max_workers, images.size())
: images.size();
std::vector<size_t> range(n_init_images, 0);
std::vector<size_t> range(images.size(), 0);
std::iota(range.begin(), range.end(), 0);
tbb::parallel_for_each(range, project_one_image);
if (blending_method & BlendingMethod::AVERAGE) {
albedo.Slice(2, 0, 3) /= albedo.Slice(2, 3, 4);
}
albedo.Slice(2, 0, 3) /= albedo.Slice(2, 3, 4);

// Image::To uses saturate_cast
Image albedo_texture =
Expand Down
Loading

0 comments on commit 7b7316e

Please sign in to comment.