Merge pull request #22 from originlake/fast_undistort

added lru caching for faster undistortion

Author: pierotofy

opensfm/commands/undistort.py

@@ -16,6 +16,7 @@ def run_impl(self, dataset: DataSet, args: argparse.Namespace) -> None:
             args.reconstruction_index,
             args.tracks,
             args.output,
+            None,
             args.skip_images
         )
 

opensfm/context.py

@@ -11,7 +11,9 @@
 
 import cv2
 from joblib import Parallel, delayed, parallel_backend
-
+from multiprocessing import Lock as ProcessLock, Manager as ProcessManager
+from multiprocessing.dummy import Lock as ThreadLock
+import functools
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -61,7 +63,38 @@ def parallel_map(func, args, num_proc: int, max_batch_size: int = 1, backend="th
     cv2.setNumThreads(threads_used)
     return res
 
-
+def lru_cache(maxsize=128, use_multiprocessing=False):
+    """Least-recently-used cache decorator."""
+    def decorator(func):
+        if use_multiprocessing:
+            manager = ProcessManager()
+            cache = manager.dict()
+            keys = manager.list()
+            lock = ProcessLock()
+        else:
+            cache = {}
+            keys = []
+            lock = ThreadLock()
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            key = (args, tuple(kwargs.items()))
+            with lock:
+                if key in cache:
+                    keys.remove(key)
+                    keys.append(key)
+                    return cache[key]
+                result = func(*args, **kwargs)
+                if len(keys) >= maxsize:
+                    old_key = keys.pop(0)
+                    del cache[old_key]
+                keys.append(key)
+                cache[key] = result
+                return result
+
+        return wrapper
+
+    return decorator
 # Memory usage
 
 if sys.platform == "win32":

opensfm/src/robust/CMakeLists.txt

@@ -14,6 +14,7 @@ set(ROBUST_FILES
     src/relative_pose_model.cc
     src/line_model.cc
     src/instanciations.cc
+    src/similarity_model.cc
 )
 add_library(robust ${ROBUST_FILES})
 target_link_libraries(robust

opensfm/src/robust/src/similarity_model.cc

@@ -0,0 +1,4 @@
+#include "robust/similarity_model.h"
+
+
+const int Similarity::MINIMAL_SAMPLES;

opensfm/undistort.py

@@ -13,7 +13,7 @@
     types,
     features_processing,
 )
-from opensfm.context import parallel_map
+from opensfm.context import parallel_map, lru_cache
 from opensfm.dataset import UndistortedDataSet
 from opensfm.dataset_base import DataSetBase
 
@@ -153,6 +153,7 @@ def undistort_image_and_masks(arguments) -> None:
         for k, v in undistorted.items():
             udata.save_undistorted_segmentation(k, v)
 
+compute_camera_mapping_lru = lru_cache(maxsize=100)(pygeometry.compute_camera_mapping)
 
 def undistort_image(
     shot: pymap.Shot,
@@ -180,7 +181,7 @@ def undistort_image(
         [undistorted_shot] = undistorted_shots
         new_camera = undistorted_shot.camera
         height, width = original.shape[:2]
-        map1, map2 = pygeometry.compute_camera_mapping(
+        map1, map2 = compute_camera_mapping_lru(
             shot.camera, new_camera, width, height
         )
         undistorted = cv2.remap(original, map1, map2, interpolation)