Add sampling gradient scaling based on surface area normalization

Libs/Groom/Groom.cpp

@@ -378,6 +378,14 @@ bool Groom::run_mesh_pipeline(Mesh& mesh, GroomParameters params, const std::str
   }
 
   if (params.get_remesh()) {
+    auto poly_data = mesh.getVTKMesh();
+    if (poly_data->GetNumberOfCells() == 0 || poly_data->GetCell(0)->GetNumberOfPoints() == 2) {
+      SW_DEBUG("Number of cells: {}", poly_data->GetNumberOfCells());
+      if (poly_data->GetNumberOfCells() > 0) {
+        SW_DEBUG("Number of points in first cell: {}", poly_data->GetCell(0)->GetNumberOfPoints());
+      }
+      throw std::runtime_error("malformed mesh, mesh should be triangular");
+    }
     int total_vertices = mesh.getVTKMesh()->GetNumberOfPoints();
     int num_vertices = params.get_remesh_num_vertices();
     if (params.get_remesh_percent_mode()) {

Libs/Mesh/Mesh.cpp

@@ -33,6 +33,7 @@
 #include <vtkIncrementalPointLocator.h>
 #include <vtkKdTreePointLocator.h>
 #include <vtkLoopSubdivisionFilter.h>
+#include <vtkMassProperties.h>
 #include <vtkNew.h>
 #include <vtkOBJReader.h>
 #include <vtkOBJWriter.h>
@@ -1198,6 +1199,13 @@ Point3 Mesh::centerOfMass() const {
   return center;
 }
 
+double Mesh::getSurfaceArea() const {
+  auto mass_props = vtkSmartPointer<vtkMassProperties>::New();
+  mass_props->SetInputData(this->poly_data_);
+  mass_props->Update();
+  return mass_props->GetSurfaceArea();
+}
+
 Point3 Mesh::getPoint(int id) const {
   if (this->numPoints() < id) {
     throw std::invalid_argument("mesh has fewer indices than requested");

Libs/Mesh/Mesh.h

@@ -206,6 +206,9 @@ class Mesh {
   /// center of mass of mesh
   Point3 centerOfMass() const;
 
+  /// surface area of mesh
+  double getSurfaceArea() const;
+
   /// number of points
   int numPoints() const { return poly_data_->GetNumberOfPoints(); }
 

Libs/Optimize/Domain/ContourDomain.h

@@ -85,7 +85,10 @@ class ContourDomain : public ParticleDomain {
     return out;
   }
 
-  double GetSurfaceArea() const override { throw std::runtime_error("Contours do not have area"); }
+  double GetSurfaceArea() const override {
+    // TODO: Implement something analogous for scaling purposes
+    return 1.0;
+  }
 
   void DeleteImages() override {
     // TODO what?
@@ -100,7 +103,6 @@ class ContourDomain : public ParticleDomain {
                                   const VectorDoubleType& global_direction, double epsilon) const override;
 
  private:
-
   double ComputeLineCoordinate(const double pt[3], int line) const;
 
   // Return the number of lines that consist of i-th point
@@ -136,7 +138,6 @@ class ContourDomain : public ParticleDomain {
   void ComputeAvgEdgeLength();
 
   int GetLineForPoint(const double pt[3], int idx, double& closest_distance, double closest_pt[3]) const;
-
 };
 
 }  // namespace shapeworks

Libs/Optimize/Domain/ImageDomain.h

@@ -111,8 +111,7 @@ class ImageDomain : public ParticleRegionDomain {
     this->UpdateSurfaceArea(I);
   }
 
-  inline double GetSurfaceArea() const override {
-    throw std::runtime_error("Surface area is not computed currently.");
+  double GetSurfaceArea() const override {
     return m_SurfaceArea;
   }
 
@@ -137,7 +136,8 @@ class ImageDomain : public ParticleRegionDomain {
     } else {
       std::ostringstream message;
       message << "Domain " << m_DomainID << ": " << m_DomainName << " : Distance transform queried for a Point, " << p
-              << ", outside the given image domain. Consider increasing the padding in grooming or the narrow band optimization parameter";
+              << ", outside the given image domain. Consider increasing the padding in grooming or the narrow band "
+                 "optimization parameter";
       throw std::runtime_error(message.str());
     }
   }
@@ -179,7 +179,7 @@ class ImageDomain : public ParticleRegionDomain {
   openvdb::FloatGrid::Ptr GetVDBImage() const { return m_VDBImage; }
 
   ImageDomain() {}
-  virtual ~ImageDomain(){};
+  virtual ~ImageDomain() {};
 
   void PrintSelf(std::ostream& os, itk::Indent indent) const {
     ParticleRegionDomain::PrintSelf(os, indent);
@@ -241,20 +241,8 @@ class ImageDomain : public ParticleRegionDomain {
   }
 
   void UpdateSurfaceArea(ImageType* I) {
-    // TODO: This code has been copied from Optimize.cpp. It does not work
-    /*
-    typename itk::ImageToVTKImageFilter < ImageType > ::Pointer itk2vtkConnector;
-    itk2vtkConnector = itk::ImageToVTKImageFilter < ImageType > ::New();
-    itk2vtkConnector->SetInput(I);
-    vtkSmartPointer < vtkContourFilter > ls = vtkSmartPointer < vtkContourFilter > ::New();
-    ls->SetInputData(itk2vtkConnector->GetOutput());
-    ls->SetValue(0, 0.0);
-    ls->Update();
-    vtkSmartPointer < vtkMassProperties > mp = vtkSmartPointer < vtkMassProperties > ::New();
-    mp->SetInputData(ls->GetOutput());
-    mp->Update();
-    m_SurfaceArea = mp->GetSurfaceArea();
-    */
+    Image image(I);
+    m_SurfaceArea = image.toMesh(0).getSurfaceArea();
   }
 };
 

Libs/Optimize/Domain/MeshDomain.cpp

@@ -85,7 +85,9 @@ void MeshDomain::SetMesh(std::shared_ptr<Surface> mesh, double geodesic_remesh_p
   surface_ = mesh;
   sw_mesh_ = std::make_shared<Mesh>(surface_->get_polydata());
 
-  if (geodesic_remesh_percent >= 100.0) { // no remeshing
+  surface_area_ = sw_mesh_->getSurfaceArea();
+
+  if (geodesic_remesh_percent >= 100.0) {  // no remeshing
     geodesics_mesh_ = surface_;
   } else {
     auto poly_data = surface_->get_polydata();

Libs/Optimize/Domain/MeshDomain.h

@@ -45,8 +45,7 @@ class MeshDomain : public ParticleDomain {
   PointType GetValidLocationNear(PointType p) const override;
 
   double GetSurfaceArea() const override {
-    // TODO return actual surface area
-    return 0;
+    return surface_area_;
   }
 
   double GetMaxDiameter() const override;
@@ -86,6 +85,7 @@ class MeshDomain : public ParticleDomain {
   std::shared_ptr<Surface> geodesics_mesh_;
   std::shared_ptr<Mesh> sw_mesh_;
   PointType zero_crossing_point_;
+  double surface_area_ = 0.0;
 };
 
 }  // namespace shapeworks

Libs/Optimize/Function/SamplingFunction.cpp

@@ -1,6 +1,8 @@
 
 #include "SamplingFunction.h"
 
+#include <set>
+
 #include "Libs/Common/Logging.h"
 #include "Libs/Optimize/Domain/DomainType.h"
 #include "vnl/vnl_vector_fixed.h"
@@ -85,6 +87,9 @@ std::shared_ptr<VectorFunction> SamplingFunction::clone() {
   copy->m_avgKappa = m_avgKappa;
   copy->m_IsSharedBoundaryEnabled = m_IsSharedBoundaryEnabled;
   copy->m_SharedBoundaryWeight = m_SharedBoundaryWeight;
+  copy->m_SamplingScale = m_SamplingScale;
+  copy->m_SamplingAutoScale = m_SamplingAutoScale;
+  copy->m_SamplingScaleValue = m_SamplingScaleValue;
   copy->m_CurrentSigma = m_CurrentSigma;
   copy->m_CurrentNeighborhood = m_CurrentNeighborhood;
 
@@ -244,6 +249,37 @@ SamplingFunction::VectorType SamplingFunction::evaluate(unsigned int idx, unsign
   energy = (A * sigma2inv) / m_avgKappa;
 
   gradE = gradE / m_avgKappa;
+
+  // Apply sampling scale if enabled
+  if (m_SamplingScale) {
+    double scale_factor = 1.0;
+
+    if (m_SamplingAutoScale) {
+      // Get surface area from the domain
+      double surface_area = system->GetDomain(d)->GetSurfaceArea();
+
+      // Reference surface area for scale factor of 1.0
+      // This constant will be tuned based on experiments
+      constexpr double reference_surface_area = 12250.0;
+
+      // Scale factor is proportional to surface area
+      scale_factor = surface_area / reference_surface_area;
+
+      // Log once per domain using static set
+      static std::set<int> logged_domains;
+      if (logged_domains.find(d) == logged_domains.end()) {
+        logged_domains.insert(d);
+        SW_DEBUG("SamplingFunction: Auto scale for domain " + std::to_string(d) + ", surface_area = " +
+                 std::to_string(surface_area) + ", scale_factor = " + std::to_string(scale_factor));
+      }
+    }
+
+    // multiply by scaling value (whether auto is on or off)
+    scale_factor *= m_SamplingScaleValue;
+
+    gradE = gradE * scale_factor;
+  }
+
   return gradE;
 }
 

Libs/Optimize/Function/SamplingFunction.h

@@ -53,6 +53,15 @@ class SamplingFunction : public VectorFunction {
   void SetSharedBoundaryEnabled(bool enabled) { m_IsSharedBoundaryEnabled = enabled; }
   bool GetSharedBoundaryEnabled() const { return m_IsSharedBoundaryEnabled; }
 
+  void SetSamplingScale(bool enabled) { m_SamplingScale = enabled; }
+  bool GetSamplingScale() const { return m_SamplingScale; }
+
+  void SetSamplingAutoScale(bool auto_scale) { m_SamplingAutoScale = auto_scale; }
+  bool GetSamplingAutoScale() const { return m_SamplingAutoScale; }
+
+  void SetSamplingScaleValue(double scale_value) { m_SamplingScaleValue = scale_value; }
+  double GetSamplingScaleValue() const { return m_SamplingScaleValue; }
+
   /**Access the cache of sigma values for each particle position.  This cache
      is populated by registering this object as an observer of the correct
      particle system (see SetParticleSystem).*/
@@ -105,6 +114,9 @@ class SamplingFunction : public VectorFunction {
   double m_avgKappa{0};
   bool m_IsSharedBoundaryEnabled{false};
   double m_SharedBoundaryWeight{1.0};
+  bool m_SamplingScale{true};
+  bool m_SamplingAutoScale{true};
+  double m_SamplingScaleValue{1.0};
   double m_CurrentSigma{0.0};
   float m_MaxMoveFactor{0};
   double m_MinimumNeighborhoodRadius{0};

Libs/Optimize/Optimize.cpp

@@ -435,6 +435,10 @@ void Optimize::InitializeSampler() {
 
   m_sampler->SetCorrespondenceOn();
 
+  m_sampler->SetSamplingScale(m_sampling_scale);
+  m_sampler->SetSamplingAutoScale(m_sampling_auto_scale);
+  m_sampler->SetSamplingScaleValue(m_sampling_scale_value);
+
   m_sampler->SetAdaptivityMode();
   m_sampler->GetEnsembleEntropyFunction()->SetRecomputeCovarianceInterval(m_recompute_regularization_interval);
   m_sampler->GetDisentangledEnsembleEntropyFunction()->SetRecomputeCovarianceInterval(
@@ -1708,7 +1712,11 @@ void Optimize::AddImage(ImageType::Pointer image, std::string name) {
   m_sampler->AddImage(image, this->GetNarrowBand(), name);
   this->m_num_shapes++;
   if (image) {
-    this->m_spacing = image->GetSpacing()[0] * 5;
+    double new_spacing = image->GetSpacing()[0] * 5;
+    if (m_spacing == 0 || new_spacing < this->m_spacing) {
+      // pick the smallest spacing
+      m_spacing = new_spacing;
+    }
   }
 }
 
@@ -2019,6 +2027,24 @@ void Optimize::SetSharedBoundaryEnabled(bool enabled) { m_sampler->SetSharedBoun
 //---------------------------------------------------------------------------
 void Optimize::SetSharedBoundaryWeight(double weight) { m_sampler->SetSharedBoundaryWeight(weight); }
 
+//---------------------------------------------------------------------------
+void Optimize::SetSamplingScale(bool enabled) { m_sampling_scale = enabled; }
+
+//---------------------------------------------------------------------------
+bool Optimize::GetSamplingScale() { return m_sampling_scale; }
+
+//---------------------------------------------------------------------------
+void Optimize::SetSamplingAutoScale(bool auto_scale) { m_sampling_auto_scale = auto_scale; }
+
+//---------------------------------------------------------------------------
+bool Optimize::GetSamplingAutoScale() { return m_sampling_auto_scale; }
+
+//---------------------------------------------------------------------------
+void Optimize::SetSamplingScaleValue(double scale_value) { m_sampling_scale_value = scale_value; }
+
+//---------------------------------------------------------------------------
+double Optimize::GetSamplingScaleValue() { return m_sampling_scale_value; }
+
 //---------------------------------------------------------------------------
 void Optimize::SetEarlyStoppingConfig(EarlyStoppingConfig config) { m_sampler->SetEarlyStoppingConfig(config); }
 

Libs/Optimize/Optimize.h

@@ -240,6 +240,14 @@ class Optimize {
   void SetSharedBoundaryEnabled(bool enabled);
   void SetSharedBoundaryWeight(double weight);
 
+  //! Sampling scale settings
+  void SetSamplingScale(bool enabled);
+  bool GetSamplingScale();
+  void SetSamplingAutoScale(bool auto_scale);
+  bool GetSamplingAutoScale();
+  void SetSamplingScaleValue(double scale_value);
+  double GetSamplingScaleValue();
+
   //! Early Stopping params
   void SetEarlyStoppingConfig(EarlyStoppingConfig config);
 
@@ -416,6 +424,11 @@ class Optimize {
   size_t m_geodesic_cache_size_multiplier = 0;  // 0 => MeshWrapper will use a heuristic to determine cache size
   double m_geodesic_remesh_percent = 100.0;     // 100% by default (e.g. no remeshing)
 
+  // Sampling scale parameters
+  bool m_sampling_scale = true;                 // enabled by default
+  bool m_sampling_auto_scale = true;            // auto scale enabled by default
+  double m_sampling_scale_value = 1.0;          // default scale value
+
   // m_spacing is used to scale the random update vector for particle splitting.
   double m_spacing = 0;
 

Libs/Optimize/OptimizeParameters.cpp

@@ -53,6 +53,9 @@ const std::string particle_format = "particle_format";
 const std::string geodesic_remesh_percent = "geodesic_remesh_percent";
 const std::string shared_boundary = "shared_boundary";
 const std::string shared_boundary_weight = "shared_boundary_weight";
+const std::string sampling_scale = "sampling_scale";
+const std::string sampling_auto_scale = "sampling_auto_scale";
+const std::string sampling_scale_value = "sampling_scale_value";
 const std::string early_stopping_threshold = "early_stopping_threshold";
 const std::string early_stopping_frequency = "early_stopping_frequency";
 const std::string early_stopping_window = "early_stopping_window";
@@ -105,6 +108,9 @@ OptimizeParameters::OptimizeParameters(ProjectHandle project) {
                                          Keys::geodesic_remesh_percent,
                                          Keys::shared_boundary,
                                          Keys::shared_boundary_weight,
+                                         Keys::sampling_scale,
+                                         Keys::sampling_auto_scale,
+                                         Keys::sampling_scale_value,
                                          Keys::early_stopping_threshold,
                                          Keys::early_stopping_frequency,
                                          Keys::early_stopping_window,
@@ -451,6 +457,9 @@ bool OptimizeParameters::set_up_optimize(Optimize* optimize) {
   optimize->set_particle_format(get_particle_format());
   optimize->SetSharedBoundaryEnabled(get_shared_boundary());
   optimize->SetSharedBoundaryWeight(get_shared_boundary_weight());
+  optimize->SetSamplingScale(get_sampling_scale());
+  optimize->SetSamplingAutoScale(get_sampling_auto_scale());
+  optimize->SetSamplingScaleValue(get_sampling_scale_value());
   optimize->SetEarlyStoppingConfig(get_early_stopping_config());
 
   std::vector<bool> use_normals;
@@ -882,6 +891,24 @@ double OptimizeParameters::get_shared_boundary_weight() { return params_.get(Key
 //---------------------------------------------------------------------------
 void OptimizeParameters::set_shared_boundary_weight(double value) { params_.set(Keys::shared_boundary_weight, value); }
 
+//---------------------------------------------------------------------------
+bool OptimizeParameters::get_sampling_scale() { return params_.get(Keys::sampling_scale, false); }
+
+//---------------------------------------------------------------------------
+void OptimizeParameters::set_sampling_scale(bool value) { params_.set(Keys::sampling_scale, value); }
+
+//---------------------------------------------------------------------------
+bool OptimizeParameters::get_sampling_auto_scale() { return params_.get(Keys::sampling_auto_scale, true); }
+
+//---------------------------------------------------------------------------
+void OptimizeParameters::set_sampling_auto_scale(bool value) { params_.set(Keys::sampling_auto_scale, value); }
+
+//---------------------------------------------------------------------------
+double OptimizeParameters::get_sampling_scale_value() { return params_.get(Keys::sampling_scale_value, 1.0); }
+
+//---------------------------------------------------------------------------
+void OptimizeParameters::set_sampling_scale_value(double value) { params_.set(Keys::sampling_scale_value, value); }
+
 //---------------------------------------------------------------------------
 EarlyStoppingConfig OptimizeParameters::get_early_stopping_config() {
   EarlyStoppingConfig config;

Libs/Optimize/OptimizeParameters.h

@@ -138,6 +138,15 @@ class OptimizeParameters {
   double get_shared_boundary_weight();
   void set_shared_boundary_weight(double value);
 
+  bool get_sampling_scale();
+  void set_sampling_scale(bool value);
+
+  bool get_sampling_auto_scale();
+  void set_sampling_auto_scale(bool value);
+
+  double get_sampling_scale_value();
+  void set_sampling_scale_value(double value);
+
   EarlyStoppingConfig get_early_stopping_config();
 
   Parameters get_parameters() const;

Libs/Optimize/Sampler.cpp

@@ -161,6 +161,9 @@ void Sampler::InitializeOptimizationFunctions() {
     m_SamplingFunction->SetSharedBoundaryEnabled(true);
     m_SamplingFunction->SetSharedBoundaryWeight(this->m_SharedBoundaryWeight);
   }
+  m_SamplingFunction->SetSamplingScale(m_SamplingScale);
+  m_SamplingFunction->SetSamplingAutoScale(m_SamplingAutoScale);
+  m_SamplingFunction->SetSamplingScaleValue(m_SamplingScaleValue);
 
   m_LinearRegressionShapeMatrix->Initialize();
   m_MixedEffectsShapeMatrix->Initialize();