From 07d57c56fa910c0dd10d4d21e7064f9c1a6e2dbc Mon Sep 17 00:00:00 2001
From: Pranjal Sahu <pranjalsahu5@gmail.com>
Date: Wed, 21 Dec 2022 17:10:08 -0500
Subject: [PATCH] Check for RMSE when inlier count is same

---
 include/itkLandmarkRegistrationEstimator.h   | 2 +-
 include/itkLandmarkRegistrationEstimator.hxx | 6 +++---
 include/itkParametersEstimator.h             | 2 +-
 include/itkRANSAC.h                          | 1 +
 include/itkRANSAC.hxx                        | 9 +++++++--
 5 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/include/itkLandmarkRegistrationEstimator.h b/include/itkLandmarkRegistrationEstimator.h
index dbe69f1..32c956d 100644
--- a/include/itkLandmarkRegistrationEstimator.h
+++ b/include/itkLandmarkRegistrationEstimator.h
@@ -69,7 +69,7 @@ class LandmarkRegistrationEstimator : public itk::ParametersEstimator<Point<doub
   virtual bool
   Agree(std::vector<double> & parameters, Point<double, Dimension> & data) override;
 
-  virtual std::vector<bool>
+  virtual std::vector<double>
   AgreeMultiple(std::vector<double> & parameters, std::vector<Point<double, Dimension>> & data, unsigned int currentBest) override;
 
   virtual bool
diff --git a/include/itkLandmarkRegistrationEstimator.hxx b/include/itkLandmarkRegistrationEstimator.hxx
index 8fa5870..023e3d1 100644
--- a/include/itkLandmarkRegistrationEstimator.hxx
+++ b/include/itkLandmarkRegistrationEstimator.hxx
@@ -371,7 +371,7 @@ LandmarkRegistrationEstimator<Dimension, TTransform>::CheckCorresspondenceDistan
 }
 
 template <unsigned int Dimension, typename TTransform>
-std::vector<bool>
+std::vector<double>
 LandmarkRegistrationEstimator<Dimension, TTransform>::AgreeMultiple(std::vector<double> & parameters, 
       std::vector<Point<double, Dimension>> & data, unsigned int currentBest)
 {
@@ -398,7 +398,7 @@ LandmarkRegistrationEstimator<Dimension, TTransform>::AgreeMultiple(std::vector<
   transform->SetParameters(optParameters);
 
   std::vector<double> query_pt(3);
-  std::vector<bool> output(data.size(), false);
+  std::vector<double> output(data.size(), -1.0);
   //output.reserve(data.size());
 
   const size_t num_results = 1;
@@ -434,8 +434,8 @@ LandmarkRegistrationEstimator<Dimension, TTransform>::AgreeMultiple(std::vector<
     if (flag)
     {
       localBest++;
+      output[i] = out_dists_sqr[0];
     }
-    output[i] = flag;
   }
 
   return output;
diff --git a/include/itkParametersEstimator.h b/include/itkParametersEstimator.h
index 73f66c9..596a211 100644
--- a/include/itkParametersEstimator.h
+++ b/include/itkParametersEstimator.h
@@ -86,7 +86,7 @@ class ITK_TEMPLATE_EXPORT ParametersEstimator : public Object
   virtual bool
   Agree(std::vector<SType> & parameters, T & data) = 0;
 
-  virtual std::vector<bool>
+  virtual std::vector<double>
   AgreeMultiple(std::vector<SType> & parameters, std::vector<T> & data, unsigned int currentBest) = 0;
 
   virtual bool
diff --git a/include/itkRANSAC.h b/include/itkRANSAC.h
index 25ae0bf..369537d 100644
--- a/include/itkRANSAC.h
+++ b/include/itkRANSAC.h
@@ -203,6 +203,7 @@ class ITK_TEMPLATE_EXPORT RANSAC : public Object
   // agrees with the best model, otherwise false
   bool *       bestVotes;
   unsigned int numVotesForBest;
+  double       bestRMSE;
 
   std::vector<T> data;
   std::vector<T> agreeData;
diff --git a/include/itkRANSAC.hxx b/include/itkRANSAC.hxx
index 09a5d6e..02c99ca 100644
--- a/include/itkRANSAC.hxx
+++ b/include/itkRANSAC.hxx
@@ -345,19 +345,24 @@ RANSAC<T, SType, TTransform>::RANSACThreadCallback(void * arg)
 
         // Expensive Inlier Test
         auto result = caller->paramEstimator->AgreeMultiple(exactEstimateParameters, caller->agreeData, caller->numVotesForBest);
+        double rmse_value = 0.0;
+
         for (m = 0; m < numAgreeObjects; m++)
         { 
-          if (result[m])
+          if (result[m] > 0)
           {
             curVotes[m] = true;
             numVotesForCur++;
+            rmse_value = rmse_value + result[m];
           }
         } // found a larger consensus set?
         
         caller->resultsMutex.lock();
-        if (numVotesForCur > caller->numVotesForBest)
+        if (numVotesForCur > caller->numVotesForBest || (numVotesForCur == caller->numVotesForBest && rmse_value < caller->bestRMSE))
         {
           caller->numVotesForBest = numVotesForCur;
+          caller->bestRMSE = rmse_value;
+
           std::copy(curVotes, curVotes + numAgreeObjects, caller->bestVotes);
 
           caller->parametersRansac.clear();