Merge pull request #12 from iago-suarez/batched_lsd

Batched LSD

.github/workflows/build-new.yml

@@ -34,44 +34,44 @@ jobs:
           name: pytlsd-${{ matrix.platform }}
           path: wheelhouse/pytlsd-*.whl
 
-  mac-build:
-    name: Wrapper macOS Build
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ macos-11, macos-12 ]
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-      - name: Build the macOS wheels
-        run: |
-          ./package/build-wheels-macos.sh
-      - name: Archive wheels
-        uses: actions/upload-artifact@v3
-        with:
-          name: pytlsd-${{ matrix.os }}
-          path: ./wheelhouse/pytlsd-*.whl
+#  mac-build:
+#    name: Wrapper macOS Build
+#    runs-on: ${{ matrix.os }}
+#    strategy:
+#      matrix:
+#        os: [ macos-11, macos-12 ]
+#    steps:
+#      - name: Checkout
+#        uses: actions/checkout@v3
+#        with:
+#          submodules: recursive
+#      - name: Build the macOS wheels
+#        run: |
+#          ./package/build-wheels-macos.sh
+#      - name: Archive wheels
+#        uses: actions/upload-artifact@v3
+#        with:
+#          name: pytlsd-${{ matrix.os }}
+#          path: ./wheelhouse/pytlsd-*.whl
 
-  pypi-publish:
-    name: Publish wheels to PyPI
-    needs: [ linux-build, mac-build ]
-    runs-on: ubuntu-latest
-    # We publish the wheel to pypi when a new tag is pushed,
-    # either by creating a new GitHub release or explictly with `git tag`
-    if: ${{ github.event_name == 'release' || startsWith(github.ref, 'refs/tags') }}
-    steps:
-      - name: Download wheels
-        uses: actions/download-artifact@v3
-        with:
-          path: ./artifacts/
-      - name: Move wheels
-        run: mkdir ./wheelhouse && mv ./artifacts/**/*.whl ./wheelhouse/
-      - name: Publish package
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          skip_existing: true
-          user: __token__
-          password: ${{ secrets.PYPI_API_TOKEN }}
-          packages_dir: ./wheelhouse/
+#  pypi-publish:
+#    name: Publish wheels to PyPI
+#    needs: [ linux-build, mac-build ]
+#    runs-on: ubuntu-latest
+#    # We publish the wheel to pypi when a new tag is pushed,
+#    # either by creating a new GitHub release or explictly with `git tag`
+#    if: ${{ github.event_name == 'release' || startsWith(github.ref, 'refs/tags') }}
+#    steps:
+#      - name: Download wheels
+#        uses: actions/download-artifact@v3
+#        with:
+#          path: ./artifacts/
+#      - name: Move wheels
+#        run: mkdir ./wheelhouse && mv ./artifacts/**/*.whl ./wheelhouse/
+#      - name: Publish package
+#        uses: pypa/gh-action-pypi-publish@release/v1
+#        with:
+#          skip_existing: true
+#          user: __token__
+#          password: ${{ secrets.PYPI_API_TOKEN }}
+#          packages_dir: ./wheelhouse/

.github/workflows/pip.yml

@@ -26,12 +26,6 @@ jobs:
           pip install -r requirements.txt
       - name: Build python package
         run: pip install .
-      - name: Lint with ruff
-        run: |
-          # stop the build if there are Python syntax errors or undefined names
-          ruff --format=github --select=E9,F63,F7,F82 --target-version=py37 .
-          # default set of ruff rules with GitHub Annotations
-          ruff --format=github --target-version=py37 .
       - name: Test with pytest
         run: |
           pytest -v tests/tests.py --junitxml=test-reports/report.xml

CMakeLists.txt

@@ -14,6 +14,11 @@ message( STATUS "    OpenCV_INCLUDE_DIRS: " ${OpenCV_INCLUDE_DIRS})
 message( STATUS "    OpenCV_LIBS: " ${OpenCV_LIBS})
 include_directories(${OpenCV_INCLUDE_DIRS})
 
+find_package(OpenMP)
+if(OpenMP_CXX_FOUND)
+    set(OpenMP_LIBS OpenMP::OpenMP_CXX)
+endif()
+
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
 # Declare the library containing the Line segment Detection (LSD) code
@@ -24,7 +29,7 @@ add_library(tlsd src/lsd.cpp)
 pybind11_add_module(pytlsd src/PYAPI.cpp)
 
 # Add the dependency between the LSD code and the bindings
-target_link_libraries(pytlsd PRIVATE tlsd ${OpenCV_LIBS})
+target_link_libraries(pytlsd PRIVATE tlsd ${OpenCV_LIBS} ${OpenMP_LIBS})
 
 target_compile_definitions(pytlsd PRIVATE VERSION_INFO=${EXAMPLE_VERSION_INFO})
 

README.md

@@ -21,3 +21,4 @@ pip3 install .
 ```
 python3 tests/test.py
 ```
+

setup.py

@@ -137,7 +137,7 @@ def build_extension(self, ext: CMakeExtension) -> None:
 # logic and declaration, and simpler if you include description/version in a file.
 setup(
     name="pytlsd",
-    version="0.0.4",
+    version="0.0.5",
     author="Iago Suarez",
     author_email="iagoh92@gmail.com",
     description="Transparent bindings of LSD (Line Segment Detector)",

src/PYAPI.cpp

@@ -29,9 +29,10 @@ void check_img_format(const py::buffer_info& correct_info, const py::buffer_info
   }
 }
 
-//  return LineSegmentDetection(n_out, img, X, Y, scale, sigma_scale, quant,
-//                              ang_th, log_eps, density_th, n_bins,
-//                              reg_img, reg_x, reg_y);
+struct LineSegment
+{
+  double x1, y1, x2, y2, /*width, */ p /*, new_log10_NFA*/;
+};
 
 // Passing in a generic array
 // Passing in an array of doubles
@@ -80,7 +81,6 @@ py::array_t<float> run_lsd(const py::array_t<double>& img,
   double *out = LineSegmentDetection(
     &N, imagePtr, info.shape[1], info.shape[0], scale, sigma_scale, quant,
     ang_th, log_eps, density_th, n_bins, grad_nfa, modgrad_ptr, angles_ptr);
-  // std::cout << "Detected " << N << " LSD Segments" << std::endl;
 
   py::array_t<float> segments({N, 5});
   for (int i = 0; i < N; i++) {
@@ -96,6 +96,89 @@ py::array_t<float> run_lsd(const py::array_t<double>& img,
   return segments;
 }
 
+py::list batched_run_lsd(const py::array_t<double>& img,
+                                   double scale=0.8,
+                                   double sigma_scale=0.6,
+                                   double density_th=0.0, /* Minimal density of region points in rectangle. */
+                                   const py::array_t<double>& gradnorm = py::array_t<double>(),
+                                   const py::array_t<double>& gradangle = py::array_t<double>(),
+                                   bool grad_nfa = false) {
+  double quant = 2.0;       /* Bound to the quantization error on the
+                                gradient norm.                                */
+  double ang_th = 22.5;     /* Gradient angle tolerance in degrees.           */
+  // double log_eps = 0.0;     /* Detection threshold: -log10(NFA) > log_eps     */
+  int n_bins = 1024;        /* Number of bins in pseudo-ordering of gradient
+                               modulus.                                       */
+  double log_eps = 0;
+
+  py::buffer_info info = img.request();
+  if (info.format != "d" && info.format != "B" ) {
+    throw py::type_error("Error: The provided numpy array has the wrong type");
+  }
+
+  double *modgrad_ptr{};
+  double *angles_ptr{};
+  if (gradnorm.size() != 0 ) {
+    py::buffer_info gradnorm_info = gradnorm.request();
+    check_img_format(info, gradnorm_info, "Gradnorm");
+    modgrad_ptr = static_cast<double *>(gradnorm_info.ptr);
+  }
+
+  if (gradangle.size() != 0) {
+    py::buffer_info gradangle_info = gradangle.request();
+    check_img_format(info, gradangle_info, "Gradangle");
+    angles_ptr = static_cast<double *>(gradangle_info.ptr);
+  }
+
+  if (info.shape.size() != 3) {
+    throw py::type_error("Error: You should provide a 3 dimensional array (batch, height, width)");
+  }
+
+  double *imagePtr = static_cast<double *>(info.ptr);
+
+  const size_t batch_size = info.shape[0];
+  const size_t img_size = info.shape[2] * info.shape[1];
+
+  std::vector<std::shared_ptr<std::vector<LineSegment>>> tmp(batch_size);
+
+  #pragma omp parallel for
+  for (int b = 0 ; b < batch_size ; b++){
+    // LSD call. Returns [x1,y1,x2,y2,width,p,-log10(NFA)] for each segment
+    int N;
+    double *out = LineSegmentDetection(
+      &N, imagePtr + b * img_size, info.shape[2], info.shape[1], scale, sigma_scale, quant,
+      ang_th, log_eps, density_th, n_bins, grad_nfa, modgrad_ptr, angles_ptr);
+
+    tmp[b] = std::make_shared< std::vector<LineSegment> >(N);
+    LineSegment * p_data = tmp[b]->data();
+    for (int i = 0; i < N; i++) {
+      p_data->x1 = out[7 * i + 0];
+      p_data->y1 = out[7 * i + 1];
+      p_data->x2 = out[7 * i + 2];
+      p_data->y2 = out[7 * i + 3];
+      p_data->p = out[7 * i + 5];
+      p_data++;
+    }
+    free(out);
+  }
+
+  py::list segments;
+  for (int b = 0; b < batch_size; b++){
+    py::array_t<float> tmp2({int(tmp[b]->size()), 5});
+    for (int i = 0; i < tmp[b]->size(); i++){
+      tmp2[py::make_tuple(i, 0)] = tmp[b]->at(i).x1;
+      tmp2[py::make_tuple(i, 1)] = tmp[b]->at(i).y1;
+      tmp2[py::make_tuple(i, 2)] = tmp[b]->at(i).x2;
+      tmp2[py::make_tuple(i, 3)] = tmp[b]->at(i).y2;
+      tmp2[py::make_tuple(i, 4)] = tmp[b]->at(i).p;
+    }
+    segments.append(tmp2);
+
+  }
+
+  return segments;
+}
+
 
 PYBIND11_MODULE(pytlsd, m) {
     m.doc() = R"pbdoc(
@@ -121,6 +204,17 @@ PYBIND11_MODULE(pytlsd, m) {
           py::arg("gradangle") = py::array(),
           py::arg("grad_nfa") = false);
 
+  m.def("batched_lsd", &batched_run_lsd, R"pbdoc(
+        Computes Line Segment Detection (LSD) in the image.
+    )pbdoc",
+      py::arg("img"),
+      py::arg("scale") = 0.8,
+      py::arg("sigma_scale") = 0.6,
+      py::arg("density_th") = 0.0,
+      py::arg("gradnorm") = py::array(),
+      py::arg("gradangle") = py::array(),
+      py::arg("grad_nfa") = false);
+
 #ifndef _MSC_VER
 #ifdef VERSION_INFO
     m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);

tests/tests.py

@@ -1,3 +1,4 @@
+import random
 import unittest
 
 import cv2
@@ -61,12 +62,68 @@ def test_triangle(self) -> None:
         self.assertEqual(result.shape, (3, 5))
         self.assert_segs_close(result[:, :4], expected, tol=2.5)
 
+    def test_batched_triangle(self) -> None:
+        img = np.zeros((200, 200), np.uint8)
+        # Define the triangle
+        top = (100, 20)
+        left = (50, 160)
+        right = (150, 160)
+        expected = np.array([[*top, *right], [*right, *left], [*left, *top]])
+        # Draw triangle
+        cv2.drawContours(img, [expected.reshape(-1, 2)], 0, (255,), thickness=cv2.FILLED)
+
+        batch_img = [img]
+        batch_expected = [expected]
+        for _ in range(3):
+            batch_img.append(cv2.rotate(batch_img[-1], cv2.ROTATE_90_CLOCKWISE))
+            x0, y0, x1, y1 = batch_expected[-1].T
+            w = batch_img[-1].shape[1]
+            rotated_pts = np.stack([w - y0, x0, w - y1, x1], axis=-1)
+            batch_expected.append(rotated_pts)
+
+        batch_img = np.array(batch_img)
+
+        # result = [pytlsd.lsd(b) for b in batch_img]
+        result = pytlsd.batched_lsd(batch_img)
+
+        for e, r in zip(batch_expected, result):
+            self.assertEqual(r.shape, (3, 5))
+            self.assert_segs_close(r[:, :4], e, tol=2.5)
+
     def test_real_img(self) -> None:
         img = cv2.imread('resources/ai_001_001.frame.0000.color.jpg', cv2.IMREAD_GRAYSCALE)
         segments = pytlsd.lsd(img)
         # Check that it detects at least 500 segments
         self.assertGreater(len(segments), 500)
 
+    def test_batched_real_imgs(self) -> None:
+        img = cv2.imread('resources/ai_001_001.frame.0000.color.jpg', cv2.IMREAD_GRAYSCALE)
+        batch_size = 8
+        # Generate synthetic variations of img
+        batch = []
+        for i in range(batch_size):
+            H = np.eye(3)
+            H[0, 2] = 200 * (random.random() - 0.5)
+            H[1, 2] = 200 * (random.random() - 0.5)
+            s = 1 + 0.5 * random.random()
+            rot = (random.random() - 0.5) * 45
+            H = cv2.getRotationMatrix2D((img.shape[1] / 2, img.shape[0] / 2), rot, s) @ H
+            new_img = cv2.warpAffine(img, H.astype(np.float32), img.shape[::-1])
+            batch.append(new_img)
+
+        batch = np.array(batch)
+
+        # Compute the result sequentially
+        expected = [pytlsd.lsd(b) for b in batch]
+
+        # Compute the batched result
+        result = pytlsd.batched_lsd(batch)
+
+        # Check that the results are the same
+        for i in range(batch_size):
+            self.assert_segs_close(result[i][:, :4], expected[i][:, :4], tol=0.5)
+
+
     def test_with_grads(self) -> None:
         # Read one image
         gray = cv2.imread('resources/ai_001_001.frame.0000.color.jpg', cv2.IMREAD_GRAYSCALE)