PerfectMemory: Accept non-uint8 images as inputs

BrainsOnBoard · Jun 22, 2022 · 987e51d · 987e51d
1 parent df66677
commit 987e51d
Showing 1 changed file with 42 additions and 22 deletions.
diff --git a/include/navigation/perfect_memory.h b/include/navigation/perfect_memory.h
@@ -46,6 +46,7 @@ class PerfectMemory
     PerfectMemory(const cv::Size &unwrapRes, Ts &&... args)
       : m_UnwrapRes(unwrapRes)
       , m_Store(unwrapRes, std::forward<Ts>(args)...)
+      , m_ScratchImage(unwrapRes, CV_8UC1)
     {}
 
     //------------------------------------------------------------------------
@@ -58,13 +59,10 @@ class PerfectMemory
     //------------------------------------------------------------------------
     void train(const cv::Mat &image, const ImgProc::Mask &mask = ImgProc::Mask{})
     {
-        const auto &unwrapRes = getUnwrapResolution();
-        BOB_ASSERT(image.cols == unwrapRes.width);
-        BOB_ASSERT(image.rows == unwrapRes.height);
-        BOB_ASSERT(image.type() == CV_8UC1);
+        convertImage(image);
 
         // Add snapshot
-        m_Store.addSnapshot(image, mask);
+        m_Store.addSnapshot(m_ScratchImage, mask);
     }
 
     float test(const cv::Mat &image, const ImgProc::Mask &mask, const Window &window) const
@@ -116,16 +114,6 @@ class PerfectMemory
     //! Get the resolution of images
     const cv::Size &getUnwrapResolution() const { return m_UnwrapRes; }
 
-protected:
-    //------------------------------------------------------------------------
-    // Protected API
-    //------------------------------------------------------------------------
-    float calcSnapshotDifference(const cv::Mat &image,
-                                 const ImgProc::Mask &mask, size_t snapshot) const
-    {
-        return m_Store.calcSnapshotDifference(image, mask, snapshot);
-    }
-
 private:
     //------------------------------------------------------------------------
     // Private members
@@ -136,11 +124,8 @@ class PerfectMemory
 
     void testInternal(const cv::Mat &image, const ImgProc::Mask &mask, const Window &window) const
     {
-        const auto &unwrapRes = getUnwrapResolution();
-        BOB_ASSERT(image.cols == unwrapRes.width);
-        BOB_ASSERT(image.rows == unwrapRes.height);
-        BOB_ASSERT(image.type() == CV_8UC1);
-        BOB_ASSERT(mask.isValid(unwrapRes));
+        convertImage(image);
+        BOB_ASSERT(mask.isValid(getUnwrapResolution()));
 
         BOB_ASSERT(window.first < getNumSnapshots());
         BOB_ASSERT(window.second <= getNumSnapshots());
@@ -152,10 +137,36 @@ class PerfectMemory
         tbb::parallel_for(tbb::blocked_range<size_t>(window.first, window.second),
             [&](const auto &r) {
                 for (size_t s = r.begin(); s != r.end(); ++s) {
-                    m_Differences[s - window.first] = calcSnapshotDifference(image, mask, s);
+                    m_Differences[s - window.first] = calcSnapshotDifference(m_ScratchImage, mask, s);
                 }
             });
     }
+
+protected:
+    //------------------------------------------------------------------------
+    // Protected API
+    //------------------------------------------------------------------------
+    mutable cv::Mat m_ScratchImage;
+
+    float calcSnapshotDifference(const cv::Mat &image,
+                                 const ImgProc::Mask &mask, size_t snapshot) const
+    {
+        return m_Store.calcSnapshotDifference(image, mask, snapshot);
+    }
+
+    void convertImage(const cv::Mat &image) const
+    {
+        BOB_ASSERT(image.channels() == 1);
+        const auto &unwrapRes = getUnwrapResolution();
+        BOB_ASSERT(image.cols == unwrapRes.width);
+        BOB_ASSERT(image.rows == unwrapRes.height);
+
+        if (image.type() == CV_8UC1) {
+            m_ScratchImage = image;
+        } else {
+            image.convertTo(m_ScratchImage, CV_8UC1);
+        }
+    }
 };
 
 //------------------------------------------------------------------------
@@ -223,7 +234,16 @@ class PerfectMemoryRotater : public PerfectMemory<Store>
     template<class... Ts>
     auto getHeading(const cv::Mat &image, ImgProc::Mask mask, typename PerfectMemory<Store>::Window window, Ts &&... args) const
     {
-        auto rotater = InSilicoRotater::create(this->getUnwrapResolution(), mask, image, std::forward<Ts>(args)...);
+        /*
+         * We don't *need* to do this conversion here, as
+         * PerfectMemory<>::test() will accept non-uint8 images as inputs, but
+         * by doing so we only have to do the conversion once as opposed to once
+         * per rotation.
+         */
+        this->convertImage(image);
+        auto rotater = InSilicoRotater::create(this->getUnwrapResolution(),
+                                               mask, this->m_ScratchImage,
+                                               std::forward<Ts>(args)...);
         calcImageDifferences(window, rotater);
 
         // Now get the minimum for each snapshot and the column this corresponds to