Merge pull request #1476 from alicevision/dev/onnxSegmentation

image segmentation module
alicevision · Jun 27, 2023 · 3472073 · 3472073
2 parents a0b3408 + 935f7b2
commit 3472073
Show file tree

Hide file tree

Showing 12 changed files with 849 additions and 6 deletions.
diff --git a/src/aliceVision/CMakeLists.txt b/src/aliceVision/CMakeLists.txt
@@ -56,9 +56,14 @@ if(ALICEVISION_BUILD_MVS)
   if(ALICEVISION_HAVE_CUDA)
     add_subdirectory(depthMap)
   endif()
+
+  if(ALICEVISION_HAVE_ONNX)
+    add_subdirectory(segmentation)
+  endif()
 endif()
 
 
+
 if(ALICEVISION_BUILD_SFM AND ALICEVISION_BUILD_MVS)
   add_subdirectory(sfmMvsUtils)
 endif()

diff --git a/src/aliceVision/feature/FeatureExtractor.cpp b/src/aliceVision/feature/FeatureExtractor.cpp
@@ -182,9 +182,9 @@ void FeatureExtractor::computeViewJob(const FeatureExtractorViewJob& job, bool u
     {
         const auto masksFolder = fs::path(_masksFolder);
         const auto idMaskPath = masksFolder /
-                fs::path(std::to_string(job.view().getViewId())).replace_extension("png");
+                fs::path(std::to_string(job.view().getViewId())).replace_extension(_maskExtension);
         const auto nameMaskPath = masksFolder /
-                fs::path(job.view().getImagePath()).filename().replace_extension("png");
+                fs::path(job.view().getImagePath()).filename().replace_extension(_maskExtension);
 
         if (fs::exists(idMaskPath))
         {
@@ -233,7 +233,7 @@ void FeatureExtractor::computeViewJob(const FeatureExtractorViewJob& job, bool u
                 bool masked = false;
                 if (x < mask.Width() && y < mask.Height())
                 {
-                    if (mask(y, x) == 0)
+                    if ((mask(y, x) == 0 && !_maskInvert) || (mask(y, x) != 0 && _maskInvert))
                     {
                         masked = true;
                     }

diff --git a/src/aliceVision/feature/FeatureExtractor.hpp b/src/aliceVision/feature/FeatureExtractor.hpp
@@ -79,9 +79,11 @@ class FeatureExtractor
       _rangeSize = rangeSize;
     }
 
-    void setMasksFolder(const std::string& folder)
+    void setMasksFolder(const std::string& folder, const std::string& ext, bool invert)
     {
       _masksFolder = folder;
+      _maskExtension = ext;
+      _maskInvert = invert;
     }
 
     void setOutputFolder(const std::string& folder)
@@ -103,6 +105,8 @@ class FeatureExtractor
     const sfmData::SfMData& _sfmData;
     std::vector<std::shared_ptr<feature::ImageDescriber>> _imageDescribers;
     std::string _masksFolder;
+    std::string _maskExtension;
+    bool _maskInvert;
     std::string _outputFolder;
     int _rangeStart = -1;
     int _rangeSize = -1;

diff --git a/src/aliceVision/image/imageAlgo.cpp b/src/aliceVision/image/imageAlgo.cpp
@@ -1,3 +1,9 @@
+// This file is part of the AliceVision project.
+// Copyright (c) 2019 AliceVision contributors.
+// This Source Code Form is subject to the terms of the Mozilla Public License,
+// v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
 #include "imageAlgo.hpp"
 
 #include <aliceVision/image/Image.hpp>
@@ -423,6 +429,50 @@ void resizeImage(int downscale, image::Image<image::RGBAfColor>& inoutImage,
     inoutImage.swap(rescaled);
 }
 
+void resizeImage(const int outWidth, const int outHeight, const image::Image<image::RGBfColor> &inImage,
+                 image::Image<image::RGBfColor> &outImage, const std::string &filter,
+                 float filterSize)
+{
+    outImage.resize(outWidth, outHeight);
+    resizeImage(oiio::TypeDesc::FLOAT, inImage.Width(), inImage.Height(), outWidth, outHeight, 3,
+                inImage.data(), outImage.data(), filter, filterSize);
+}
+
+void resizeImage(const int outWidth, const int outHeight, const image::Image<IndexT> &inImage,
+                 image::Image<IndexT> &outImage, const std::string &filter,
+                 float filterSize)
+{
+    outImage.resize(outWidth, outHeight);
+    resizeImage(oiio::TypeDesc::UINT32, inImage.Width(), inImage.Height(), outWidth, outHeight, 1,
+                inImage.data(), outImage.data(), filter, filterSize);
+}
+
+template<typename T>
+void resampleImage(oiio::TypeDesc typeDesc,
+                 int inWidth,
+                 int inHeight,
+                 int outWidth,
+                 int outHeight,
+                 int nchannels,
+                 const T* inBuffer,
+                 T* outBuffer,
+                 bool interpolate)
+{
+    const oiio::ImageBuf inBuf(oiio::ImageSpec(inWidth, inHeight, nchannels, typeDesc),
+                               const_cast<T*>(inBuffer));
+    oiio::ImageBuf outBuf(oiio::ImageSpec(outWidth, outHeight, nchannels, typeDesc), outBuffer);
+
+    oiio::ImageBufAlgo::resample(outBuf, inBuf, interpolate);
+}
+
+void resampleImage(int outWidth, int outHeight, const image::Image<IndexT>& inImage,
+                 image::Image<IndexT>& outImage, bool interpolate)
+{
+    outImage.resize(outWidth, outHeight);
+    resampleImage(oiio::TypeDesc::UINT32, inImage.Width(), inImage.Height(), outWidth, outHeight, 1,
+        inImage.data(), outImage.data(), interpolate);
+}
+
 template<typename T>
 void convolveImage(oiio::TypeDesc typeDesc,
                    int inWidth,

diff --git a/src/aliceVision/image/imageAlgo.hpp b/src/aliceVision/image/imageAlgo.hpp
@@ -75,6 +75,24 @@ void resizeImage(int downscale, const image::Image<image::RGBAfColor>& inImage,
                  image::Image<image::RGBAfColor>& outImage,
                  const std::string& filter = "", float filterSize = 0);
 
+/**
+ * @brief Resize a given image buffer.
+ * @param[in] newWidth The destination width
+ * @param[in] newHeight The destination height
+ * @param[in] inImage The input image buffer
+ * @param[out] outImage The output image buffer
+ * @param[in] filter The name of a high-quality filter to use when resampling
+ *            Default is bilinear resampling
+ *            See openImageIO documentation "ImageBufAlgo filtername"
+ * @param[in] filterSize The resize filter size
+ */
+void resizeImage(int newWidth, int newHeight, const image::Image<IndexT>& inImage,
+                 image::Image<IndexT>& outImage,
+                 const std::string& filter = "", float filterSize = 0);
+void resizeImage(int newWidth, int newHeight, const image::Image<image::RGBfColor>& inImage,
+                 image::Image<image::RGBfColor>& outImage,
+                 const std::string& filter = "", float filterSize = 0);
+
 /**
  * @brief Resize a given image buffer in place.
  * @param[in] downscale The resize downscale
@@ -97,6 +115,18 @@ void resizeImage(int downscale, image::Image<image::RGBAColor>& inoutImage,
 void resizeImage(int downscale, image::Image<image::RGBAfColor>& inoutImage,
                  const std::string& filter = "", float filterSize = 0);
 
+
+/**
+ * @brief resample a given image buffer.
+ * @param[in] newWidth The destination width
+ * @param[in] newHeight The destination height
+ * @param[in] inImage The input image buffer
+ * @param[out] outImage The output image buffer
+ * @param[in] interpolate use interpolation (bilinear) ?
+ */
+void resampleImage(int newWidth, int newHeight, const image::Image<IndexT>& inImage,
+                 image::Image<IndexT>& outImage, bool interpolate);
+
 /**
  * @brief convolve a given image buffer
  * @param[in] inBuffer The input image buffer

diff --git a/src/aliceVision/image/pixelTypes.hpp b/src/aliceVision/image/pixelTypes.hpp
@@ -169,6 +169,17 @@ namespace aliceVision
         return Rgb( ((*this)(0) + other(0)), ((*this)(1) + other(1)), ((*this)(2) + other(2)));
       }
 
+      /**
+      * @brief Elementwise multiplication
+      * @param other the other element to multiply
+      * @return Rgb color after multiply
+      * @note This does not modify the Rgb value (ie: only return a modified copy)
+      */
+      inline Rgb operator *( const Rgb& other ) const
+      {
+        return Rgb( ((*this)(0) * other(0)), ((*this)(1) * other(1)), ((*this)(2) * other(2)));
+      }
+
 
       /**
       * @brief scalar division
@@ -183,6 +194,7 @@ namespace aliceVision
                     T( ( Z )( ( *this )( 1 ) ) / val ),
                     T( ( Z )( ( *this )( 2 ) ) / val ) );
       }
+
 
       /**
       * @brief scalar multiplication

diff --git a/src/aliceVision/segmentation/CMakeLists.txt b/src/aliceVision/segmentation/CMakeLists.txt
@@ -0,0 +1,34 @@
+# Headers
+set(segmentation_files_headers
+    segmentation.hpp
+)
+
+# Sources
+set(segmentation_files_sources
+    segmentation.cpp
+)
+
+set(SEGMENTATION_PRIVATE_LINKS "")
+set(SEGMENTATION_PRIVATE_INCLUDE_DIRS "")
+if(ALICEVISION_HAVE_CUDA)
+  set(SEGMENTATION_PRIVATE_LINKS ${CUDA_LIBRARIES})
+  set(SEGMENTATION_PRIVATE_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
+endif()
+
+alicevision_add_library(aliceVision_segmentation
+  SOURCES ${segmentation_files_headers} ${segmentation_files_sources}
+  PUBLIC_LINKS
+    aliceVision_system
+    aliceVision_numeric
+    aliceVision_image
+    ONNXRuntime::ONNXRuntime
+  PRIVATE_LINKS
+    ${SEGMENTATION_PRIVATE_LINKS}
+  PRIVATE_INCLUDE_DIRS
+    ${SEGMENTATION_PRIVATE_INCLUDE_DIRS}
+)
+
+
+
+
+