From b672c2b7c3327bbc8809104d5bf166ddf2b9a1ce Mon Sep 17 00:00:00 2001 From: Alan Morris Date: Fri, 27 Feb 2026 10:32:34 -0700 Subject: [PATCH] Parallelize mesh loading with TBB and add progress reporting Surface construction and geodesic remeshing during optimize setup are now done in parallel using tbb::parallel_for, with SW_PROGRESS updates showing "Loading meshes" or "Loading and remeshing meshes" with a count and progress bar. Pre-built Surface overloads added to MeshDomain, Sampler, and Optimize to support the three-phase approach (collect, parallel compute, sequential register). --- Libs/Optimize/Domain/MeshDomain.cpp | 10 ++ Libs/Optimize/Domain/MeshDomain.h | 2 + Libs/Optimize/Optimize.cpp | 12 ++ Libs/Optimize/Optimize.h | 2 + Libs/Optimize/OptimizeParameters.cpp | 225 ++++++++++++++++++++++----- Libs/Optimize/Sampler.cpp | 12 ++ Libs/Optimize/Sampler.h | 2 + 7 files changed, 224 insertions(+), 41 deletions(-) diff --git a/Libs/Optimize/Domain/MeshDomain.cpp b/Libs/Optimize/Domain/MeshDomain.cpp index 20ada028a9..9836067d8f 100644 --- a/Libs/Optimize/Domain/MeshDomain.cpp +++ b/Libs/Optimize/Domain/MeshDomain.cpp @@ -97,6 +97,16 @@ void MeshDomain::SetMesh(std::shared_ptr mesh, double geodesic_remesh_p } } +//------------------------------------------------------------------- +void MeshDomain::SetMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area) { + m_FixedDomain = false; + surface_ = surface; + geodesics_mesh_ = geodesics_surface; + sw_mesh_ = sw_mesh; + surface_area_ = surface_area; +} + //------------------------------------------------------------------- void MeshDomain::InvalidateParticlePosition(int idx) const { this->surface_->invalidate_particle(idx); } diff --git a/Libs/Optimize/Domain/MeshDomain.h b/Libs/Optimize/Domain/MeshDomain.h index 0da464aa12..49d8107845 100644 --- a/Libs/Optimize/Domain/MeshDomain.h +++ b/Libs/Optimize/Domain/MeshDomain.h @@ -73,6 +73,8 @@ class MeshDomain : public ParticleDomain { } void SetMesh(std::shared_ptr mesh_, double geodesic_remesh_percent); + void SetMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area); std::shared_ptr GetSWMesh() const { return sw_mesh_; } diff --git a/Libs/Optimize/Optimize.cpp b/Libs/Optimize/Optimize.cpp index d8832d97ff..ae84daba25 100644 --- a/Libs/Optimize/Optimize.cpp +++ b/Libs/Optimize/Optimize.cpp @@ -1757,6 +1757,18 @@ void Optimize::AddMesh(vtkSmartPointer poly_data) { this->m_spacing = 0.5; } +//--------------------------------------------------------------------------- +void Optimize::AddMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area) { + if (!surface) { + m_sampler->AddMesh(nullptr); + } else { + m_sampler->AddMesh(surface, geodesics_surface, sw_mesh, surface_area); + } + this->m_num_shapes++; + this->m_spacing = 0.5; +} + //--------------------------------------------------------------------------- void Optimize::AddContour(vtkSmartPointer poly_data) { m_sampler->AddContour(poly_data); diff --git a/Libs/Optimize/Optimize.h b/Libs/Optimize/Optimize.h index 6aee39c0b1..4f661a7825 100644 --- a/Libs/Optimize/Optimize.h +++ b/Libs/Optimize/Optimize.h @@ -214,6 +214,8 @@ class Optimize { //! Set the shape input images void AddImage(ImageType::Pointer image, std::string name = ""); void AddMesh(vtkSmartPointer poly_data); + void AddMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area); void AddContour(vtkSmartPointer poly_data); //! Set the shape filenames (TODO: details) diff --git a/Libs/Optimize/OptimizeParameters.cpp b/Libs/Optimize/OptimizeParameters.cpp index 678b89f57f..d6e6f318ff 100644 --- a/Libs/Optimize/OptimizeParameters.cpp +++ b/Libs/Optimize/OptimizeParameters.cpp @@ -6,10 +6,15 @@ #include #include +#include #include #include #include +#include +#include +#include +#include "Libs/Optimize/Domain/Surface.h" #include "Optimize.h" using namespace shapeworks; @@ -683,9 +688,47 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) { } } + // === Phase A: Collect work items (sequential, fast I/O) === + struct DomainWorkItem { + std::string filename; + DomainType domain_type; + std::shared_ptr subject; + std::vector> transforms; + int domain_index_in_subject; // index of this domain within the subject's files + int global_domain_index; // sequential domain index across all subjects + int subject_index; // sequential subject index (non-excluded) + bool is_fixed; + + // For mesh domains — prepared poly_data from Phase A + vtkSmartPointer poly_data; + bool is_contour_hack = false; // true if mesh cells have 2 points (contour detection hack) + + // For image domains + std::optional image; + + // Results from Phase B (parallel pre-compute) + std::shared_ptr surface; + std::shared_ptr geodesics_surface; + std::shared_ptr sw_mesh; + double surface_area = 0.0; + }; + + std::vector work_items; std::vector filenames; + // Track which work items belong to each subject for callback and particle filename assignment + struct SubjectInfo { + std::shared_ptr subject; + std::vector work_item_indices; // indices into work_items + }; + std::vector subject_infos; + int count = 0; domain_count = 0; + bool geodesics_enabled = get_use_geodesic_distance(); + double geodesic_remesh_percent = get_geodesic_remesh_percent(); + int geodesic_cache_multiplier = get_geodesic_cache_multiplier(); + bool use_geodesics_to_landmarks = get_use_geodesics_to_landmarks(); + for (auto s : subjects) { if (abort_load_) { return false; @@ -700,8 +743,9 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) { throw std::invalid_argument("No groomed inputs for optimization"); } auto transforms = s->get_groomed_transforms(); - std::vector local_particle_filenames; - std::vector world_particle_filenames; + + SubjectInfo si; + si.subject = s; for (int i = 0; i < files.size(); i++) { auto filename = files[i]; @@ -713,68 +757,164 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) { auto domain_type = project_->get_groomed_domain_types()[i]; filenames.push_back(filename); + DomainWorkItem item; + item.filename = filename; + item.domain_type = domain_type; + item.subject = s; + item.transforms = transforms; + item.domain_index_in_subject = i; + item.global_domain_index = domain_count; + item.subject_index = count; + item.is_fixed = s->is_fixed(); + if (domain_type == DomainType::Mesh) { Mesh mesh = MeshUtils::threadSafeReadMesh(filename.c_str()); if (domain_count < constraints.size()) { Constraints constraint = constraints[domain_count]; constraint.clipMesh(mesh); - auto poly_data = mesh.getVTKMesh(); - if (poly_data->GetNumberOfCells() == 0) { + auto pd = mesh.getVTKMesh(); + if (pd->GetNumberOfCells() == 0) { throw std::invalid_argument("Mesh has zero cells after constraint clipping: " + filename); } } - if (get_use_geodesics_to_landmarks()) { - auto filenames = s->get_landmarks_filenames(); - if (filenames.empty()) { - throw std::invalid_argument("Geodesic distance from landmarks is enabled, but subject has no landmark files"); + if (use_geodesics_to_landmarks) { + auto landmark_filenames = s->get_landmarks_filenames(); + if (landmark_filenames.empty()) { + throw std::invalid_argument( + "Geodesic distance from landmarks is enabled, but subject has no landmark files"); } Eigen::VectorXd points; - if (!ParticleSystemEvaluation::read_particle_file(filenames[0], points)) { - SW_ERROR("Unable to read landmark file: {}", filenames[0]); + if (!ParticleSystemEvaluation::read_particle_file(landmark_filenames[0], points)) { + SW_ERROR("Unable to read landmark file: {}", landmark_filenames[0]); } - // convert points to landmarks std::vector landmarks; - for (int i = 0; i < points.size() / 3; ++i) { + for (int j = 0; j < points.size() / 3; ++j) { Point3 p; - p[0] = points(3 * i); - p[1] = points(3 * i + 1); - p[2] = points(3 * i + 2); + p[0] = points(3 * j); + p[1] = points(3 * j + 1); + p[2] = points(3 * j + 2); landmarks.push_back(p); } mesh.computeLandmarkGeodesics(landmarks); } - auto poly_data = mesh.getVTKMesh(); - - if (poly_data) { - if (poly_data->GetNumberOfCells() == 0) { - throw std::invalid_argument("Error, mesh had zero cells: " + filename); - } - // TODO This is a HACK for detecting contours - if (poly_data->GetCell(0)->GetNumberOfPoints() == 2) { - optimize->AddContour(poly_data); - } else { - optimize->AddMesh(poly_data); - } - } else { + auto pd = mesh.getVTKMesh(); + if (!pd) { throw std::invalid_argument("Error loading mesh: " + filename); } + if (pd->GetNumberOfCells() == 0) { + throw std::invalid_argument("Error, mesh had zero cells: " + filename); + } + // TODO This is a HACK for detecting contours + item.is_contour_hack = (pd->GetCell(0)->GetNumberOfPoints() == 2); + item.poly_data = pd; } else if (domain_type == DomainType::Contour) { Mesh mesh = MeshUtils::threadSafeReadMesh(filename.c_str()); - auto poly_data = mesh.getVTKMesh(); - if (poly_data) { - optimize->AddContour(poly_data); - } else { + auto pd = mesh.getVTKMesh(); + if (!pd) { throw std::invalid_argument("Error loading contour: " + filename); } + item.poly_data = pd; } else { - Image image(filename); - if (s->is_fixed()) { - optimize->AddImage(nullptr, filename); + // Image domain — read now + item.image.emplace(filename); + } + + si.work_item_indices.push_back(work_items.size()); + work_items.push_back(std::move(item)); + domain_count++; + } + + subject_infos.push_back(std::move(si)); + count++; + } + + // === Phase B: Parallel pre-compute Surface objects for mesh domains (expensive) === + // Collect indices of mesh work items that need Surface construction + std::vector mesh_work_indices; + for (size_t idx = 0; idx < work_items.size(); idx++) { + auto& item = work_items[idx]; + if (item.domain_type == DomainType::Mesh && !item.is_contour_hack && item.poly_data) { + mesh_work_indices.push_back(idx); + } + } + + if (!mesh_work_indices.empty()) { + std::atomic completed{0}; + int total = static_cast(mesh_work_indices.size()); + + double effective_remesh_percent = geodesics_enabled ? geodesic_remesh_percent : 100.0; + bool remeshing = effective_remesh_percent < 100.0; + std::string loading_message = remeshing ? "Loading and remeshing meshes" : "Loading meshes"; + SW_PROGRESS(0, "{}: 0 / {}", loading_message, total); + + tbb::parallel_for(tbb::blocked_range(0, mesh_work_indices.size()), + [&](const tbb::blocked_range& r) { + for (size_t i = r.begin(); i < r.end(); ++i) { + auto& item = work_items[mesh_work_indices[i]]; + + // Construct the main Surface (triangulation, cleaning, normals, cell locator, geodesics) + auto surface = std::make_shared( + item.poly_data, geodesics_enabled, + static_cast(geodesic_cache_multiplier)); + + // Create the Mesh wrapper and compute surface area + auto sw_mesh = std::make_shared(surface->get_polydata()); + double surface_area = sw_mesh->getSurfaceArea(); + + // Create remeshed Surface for geodesics if needed + std::shared_ptr geodesics_surface; + if (effective_remesh_percent >= 100.0) { + geodesics_surface = surface; + } else { + auto poly_copy = surface->get_polydata(); + Mesh mesh_copy(poly_copy); + mesh_copy.remeshPercent(effective_remesh_percent, 1.0); + geodesics_surface = + std::make_shared(mesh_copy.getVTKMesh()); + } + + item.surface = surface; + item.geodesics_surface = geodesics_surface; + item.sw_mesh = sw_mesh; + item.surface_area = surface_area; + + int done = ++completed; + double progress = static_cast(done) / static_cast(total) * 100.0; + SW_PROGRESS(progress, "{}: {} / {}", loading_message, done, total); + } + }); + } + + // === Phase C: Sequential registration (fast) === + domain_count = 0; + int subject_count = 0; + for (auto& si : subject_infos) { + if (abort_load_) { + return false; + } + + std::vector local_particle_filenames; + std::vector world_particle_filenames; + + for (size_t wi_idx : si.work_item_indices) { + auto& item = work_items[wi_idx]; + + if (item.domain_type == DomainType::Mesh && !item.is_contour_hack && item.poly_data) { + // Use pre-built Surface data + optimize->AddMesh(item.surface, item.geodesics_surface, item.sw_mesh, item.surface_area); + } else if (item.domain_type == DomainType::Mesh && item.is_contour_hack) { + optimize->AddContour(item.poly_data); + } else if (item.domain_type == DomainType::Contour) { + optimize->AddContour(item.poly_data); + } else { + // Image domain + if (item.is_fixed) { + optimize->AddImage(nullptr, item.filename); } else { - optimize->AddImage(image, filename); + optimize->AddImage(*item.image, item.filename); } } @@ -782,6 +922,9 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) { TransformType prefix_transform; prefix_transform.set_identity(); + int i = item.domain_index_in_subject; + auto& transforms = item.transforms; + if (i < transforms.size() && transforms[i].size() >= 12) { prefix_transform[0][3] = transforms[i][9]; prefix_transform[1][3] = transforms[i][10]; @@ -801,19 +944,19 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) { domain_count++; - auto name = StringUtils::getBaseFilenameWithoutExtension(filename); + auto name = StringUtils::getBaseFilenameWithoutExtension(item.filename); auto extension = get_particle_format(); auto prefix = get_output_prefix(); local_particle_filenames.push_back(prefix + name + "_local." + extension); world_particle_filenames.push_back(prefix + name + "_world." + extension); } - s->set_local_particle_filenames(local_particle_filenames); - s->set_world_particle_filenames(world_particle_filenames); + si.subject->set_local_particle_filenames(local_particle_filenames); + si.subject->set_world_particle_filenames(world_particle_filenames); - count++; + subject_count++; if (load_callback_) { - load_callback_(count); + load_callback_(subject_count); } } diff --git a/Libs/Optimize/Sampler.cpp b/Libs/Optimize/Sampler.cpp index 63ed8b86a7..332bc82775 100644 --- a/Libs/Optimize/Sampler.cpp +++ b/Libs/Optimize/Sampler.cpp @@ -284,6 +284,18 @@ void Sampler::AddMesh(std::shared_ptr mesh, double geodesic m_DomainList.push_back(domain); } +void Sampler::AddMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area) { + auto domain = std::make_shared(); + + if (surface) { + this->m_Spacing = 1; + domain->SetMesh(surface, geodesics_surface, sw_mesh, surface_area); + this->m_meshes.push_back(surface->get_polydata()); + } + m_DomainList.push_back(domain); +} + void Sampler::AddContour(vtkSmartPointer poly_data) { auto domain = std::make_shared(); diff --git a/Libs/Optimize/Sampler.h b/Libs/Optimize/Sampler.h index 8fb408396e..b89d13b6e0 100644 --- a/Libs/Optimize/Sampler.h +++ b/Libs/Optimize/Sampler.h @@ -99,6 +99,8 @@ class Sampler { } void AddMesh(std::shared_ptr mesh, double geodesic_remesh_percent = 100); + void AddMesh(std::shared_ptr surface, std::shared_ptr geodesics_surface, + std::shared_ptr sw_mesh, double surface_area); void AddContour(vtkSmartPointer poly_data);