浏览代码

Enhance progress report and input parameters

Steyer, Nick 2 年之前
父节点
当前提交
e4c6177bbd
共有 1 个文件被更改,包括 47 次插入23 次删除
  1. 47 23
      hgcalc.py

+ 47 - 23
hgcalc.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import sys
 
 
 import numpy as np
 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.
     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):
     for iteration in range(n_iters):
-        print(iteration)
         """
         """
         First, we pick a sample of k corresponding point pairs
         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
             max_inliers = number_inliers_current
             inliers = 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 
     We then return the best homography we found as well as the number of inliers found with it and the inliers 
     themselves.
     themselves.
@@ -433,29 +435,51 @@ def transform_pts(p, h):
     return np.array(points_list)
     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():
 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__":
 if __name__ == "__main__":