Browse Source

Enhance progress report and input parameters

Steyer, Nick 1 year ago
parent
commit
e4c6177bbd
1 changed files with 47 additions and 23 deletions
  1. 47 23
      hgcalc.py

+ 47 - 23
hgcalc.py

@@ -1,3 +1,4 @@
+import os
 import sys
 
 import numpy as np
@@ -68,7 +69,6 @@ def ransac(pairs, n_iters, k, threshold):
     We execute the ransac loop n_iters times, so that we have a good chance to have a valid homography.
     """
     for iteration in range(n_iters):
-        print(iteration)
         """
         First, we pick a sample of k corresponding point pairs
         """
@@ -107,6 +107,8 @@ def ransac(pairs, n_iters, k, threshold):
             max_inliers = number_inliers_current
             inliers = inliers_current
 
+        report_progress(round((iteration + 1) / n_iters * 100))
+
     """
     We then return the best homography we found as well as the number of inliers found with it and the inliers 
     themselves.
@@ -433,29 +435,51 @@ def transform_pts(p, h):
     return np.array(points_list)
 
 
+def report_progress(percentage):
+    sys.stdout.write('\r')
+    sys.stdout.write("[%-100s] %d%%" % ('=' * percentage, percentage))
+    sys.stdout.flush()
+
+
 def main():
-    if len(sys.argv) != 2:
-        print('Usage: hgcalc <pairs_file_path>')
-    else:
-        pairs_file_path = sys.argv[1]
-        homography_file_path = '.\\homography.csv'
-        # RANSAC Parameters
-        ransac_threshold = 0.02  # inlier threshold
-        p = 0.35  # probability that any given correspondence is valid
-        k = 4  # number of samples drawn per iteration
-        z = 0.99  # total probability of success after all iterations
-
-        pairs = np.loadtxt(pairs_file_path)
-
-        n_iters = ransac_iters(p, k, z)
-        print('Iterations to be done: ', n_iters)
-        h, num_inliers, inliers = ransac(pairs, n_iters, k, ransac_threshold)
-        print('Number of inliers:', num_inliers)
-
-        # recompute homography matrix based on inliers
-        h = recompute_homography(inliers)
-        print(h)
-        np.savetxt(homography_file_path, h)
+    threshold = 0.02  # inlier threshold
+    p = 0.35  # probability that any given correspondence is valid
+    z = 0.99  # total probability of success after all iterations
+
+    if len(sys.argv) == 3:
+        threshold = sys.argv[2]
+    elif len(sys.argv) == 4:
+        threshold = sys.argv[2]
+        p = sys.argv[3]
+    elif len(sys.argv) == 5:
+        threshold = sys.argv[2]
+        p = sys.argv[3]
+        z = sys.argv[4]
+    elif len(sys.argv) != 2:
+        print('Usage: hgcalc <pairs_file_path> [threshold [valid_correspondence_probability [total_probability]]]')
+        return
+
+    calculate_and_save(sys.argv[1], threshold, p, z)
+
+
+def calculate_and_save(pairs_file_path, threshold, p, z):
+    homography_file_path = '.\\homography.csv'
+    pairs = np.loadtxt(pairs_file_path)
+    k = 4  # number of samples drawn at each iteration
+    n_iters = ransac_iters(p, k, z)
+    print('Input file:', pairs_file_path)
+    output_path = os.path.abspath(homography_file_path)
+    print('Output file:', output_path)
+    print('Threshold: %s, Inlier probability: %s, Success probability: %s' % (threshold, p, z))
+    print('Calculating %s iterations' % n_iters)
+    h, num_inliers, inliers = ransac(pairs, n_iters, k, threshold)
+    print()
+    print('Percentage of inliers:', "{0:.0%}".format(num_inliers / len(pairs)))
+
+    # recompute homography matrix based on inliers
+    h = recompute_homography(inliers)
+    np.savetxt(homography_file_path, h)
+    print('Homography saved at', output_path)
 
 
 if __name__ == "__main__":