{Feature} Core - Multithread undistortion (#29)

Summary: Pull Request resolved: #29 add dispenso parallel_for for cpu speed up on distortion function. Undistortion speedup 10x for slam image is 30ms to 3ms from Jupyter notebook A new dependency is introduced (Dispenso) and gathered with CMake Fetch Reviewed By: SeaOtocinclus Differential Revision: D48628380 fbshipit-source-id: ca443c22aa8a64a79ef0b4c2b2551057c26d6c0d
facebookresearch · Oct 27, 2023 · 01dff69 · 01dff69
1 parent d025981
commit 01dff69
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 18 deletions.
diff --git a/cmake/Setup3rdParty.cmake b/cmake/Setup3rdParty.cmake
@@ -14,6 +14,12 @@
 
 include(FetchContent) # once in the project to include the module
 
+FetchContent_Declare(
+  Dispenso
+  GIT_REPOSITORY      https://github.com/facebookincubator/dispenso.git
+  GIT_TAG             9360e1e214a4a295e44174d7703676a94d9ebfce # v1.1.0
+)
+
 FetchContent_Declare(
   vrs
   GIT_REPOSITORY  https://github.com/facebookresearch/vrs.git
@@ -49,7 +55,7 @@ FetchContent_Declare(
   GIT_TAG        v2.3.2
 )
 
-set(dependencies cli11 eigen Sophus)
+set(dependencies cli11 eigen Sophus Dispenso)
 foreach(X IN LISTS dependencies)
   message("Pulling deps: {${X}}")
   FetchContent_MakeAvailable(${X})

diff --git a/core/image/utility/CMakeLists.txt b/core/image/utility/CMakeLists.txt
@@ -14,7 +14,7 @@
 
 add_library(image_distort Distort.cpp Distort.h)
 target_include_directories(image_distort PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/../..)
-target_link_libraries(image_distort PUBLIC image)
+target_link_libraries(image_distort PUBLIC image dispenso)
 
 add_library(image_debayer Debayer.cpp Debayer.h)
 target_include_directories(image_debayer PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/../..)

diff --git a/core/image/utility/Distort.cpp b/core/image/utility/Distort.cpp
@@ -15,7 +15,9 @@
  */
 
 #include "Distort.h"
+#include <dispenso/parallel_for.h>
 #include <iostream>
+#include <vector>
 
 namespace projectaria::tools::image {
 
@@ -27,27 +29,36 @@ ManagedImage<T, DefaultImageAllocator<T>, MaxVal> distortImage(
     const InterpolationMethod method) {
   ManagedImage<T, DefaultImageAllocator<T>, MaxVal> dst(imageSize(0), imageSize(1));
 
+  std::vector<Eigen::Vector2i> dstPixels;
+  dstPixels.reserve(dst.height() * dst.width());
   for (int y = 0; y < dst.height(); y++) {
     for (int x = 0; x < dst.width(); x++) {
-      Eigen::Vector2f dstPixel(static_cast<float>(x), static_cast<float>(y));
-      std::optional<Eigen::Vector2f> maybeSrcPixel = inverseWarp(dstPixel);
-      if (!maybeSrcPixel || !src.inBounds((*maybeSrcPixel)(0), (*maybeSrcPixel)(1), 0.5f)) {
-        dst(x, y) = Zero<T>::val();
-      } else {
-        switch (method) {
-          case InterpolationMethod::Bilinear:
-            dst(x, y) = src((*maybeSrcPixel)(0), (*maybeSrcPixel)(1));
-            break;
-          case InterpolationMethod::NearestNeighbor:
-            Eigen::Vector2i nearestPixel =
-                (*maybeSrcPixel + Eigen::Vector2f(0.5, 0.5)).template cast<int>();
-            dst(x, y) = src(nearestPixel(0), nearestPixel(1));
-            break;
-        }
-      }
+      dstPixels.push_back(Eigen::Vector2i{x, y});
     }
   }
 
+  dispenso::parallel_for(
+      0, dstPixels.size(), [&src, &dst, &inverseWarp, &dstPixels, &method](size_t index) {
+        Eigen::Vector2f pixel(
+            static_cast<float>(dstPixels[index](0)), static_cast<float>(dstPixels[index](1)));
+        std::optional<Eigen::Vector2f> maybeSrcPixel = inverseWarp(pixel);
+        dst(dstPixels[index].x(), dstPixels[index].y()) = Zero<T>::val();
+        if (maybeSrcPixel && src.inBounds((*maybeSrcPixel)(0), (*maybeSrcPixel)(1), 0.5f)) {
+          switch (method) {
+            case InterpolationMethod::Bilinear:
+              dst(dstPixels[index].x(), dstPixels[index].y()) =
+                  src((*maybeSrcPixel)(0), (*maybeSrcPixel)(1));
+              break;
+            case InterpolationMethod::NearestNeighbor:
+              Eigen::Vector2i nearestPixel =
+                  (*maybeSrcPixel + Eigen::Vector2f(0.5, 0.5)).template cast<int>();
+              dst(dstPixels[index].x(), dstPixels[index].y()) =
+                  src(nearestPixel(0), nearestPixel(1));
+              break;
+          }
+        }
+      });
+
   return dst;
 }