From b54433be82595ed0e723908b349e859f40debae9 Mon Sep 17 00:00:00 2001
From: Borong Yuan <yuanborong@hotmail.com>
Date: Sun, 21 Jan 2024 15:53:43 +0800
Subject: [PATCH] support color histogram equalization

---
 corelib/src/CameraThread.cpp | 90 ++++++++++++++++++++++++++----------
 1 file changed, 66 insertions(+), 24 deletions(-)

diff --git a/corelib/src/CameraThread.cpp b/corelib/src/CameraThread.cpp
index a8c099a26e..a18fc2c18f 100644
--- a/corelib/src/CameraThread.cpp
+++ b/corelib/src/CameraThread.cpp
@@ -38,6 +38,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "rtabmap/core/IMUFilter.h"
 #include "rtabmap/core/Features2d.h"
 #include "rtabmap/core/clams/discrete_depth_distortion_model.h"
+#include <opencv2/imgproc/types_c.h>
 #include <opencv2/stitching/detail/exposure_compensate.hpp>
 #include <rtabmap/utilite/UTimer.h>
 #include <rtabmap/utilite/ULogger.h>
@@ -543,46 +544,87 @@ void CameraThread::postUpdate(SensorData * dataPtr, CameraInfo * info) const
 
 	if(_histogramMethod && !data.imageRaw().empty())
 	{
-		if(data.imageRaw().type() == CV_8UC1)
+		UDEBUG("");
+		UTimer timer;
+		cv::Mat image;
+		if(_histogramMethod == 1)
 		{
-			UDEBUG("");
-			UTimer timer;
-			cv::Mat image;
-			if(_histogramMethod == 1)
+			if(data.imageRaw().type() == CV_8UC1)
 			{
 				cv::equalizeHist(data.imageRaw(), image);
-				if(!data.depthRaw().empty())
+			}
+			else if(data.imageRaw().type() == CV_8UC3)
+			{
+				cv::Mat channels[3];
+				cv::cvtColor(data.imageRaw(), image, CV_BGR2YCrCb);
+				cv::split(image, channels);
+				cv::equalizeHist(channels[0], channels[0]);
+				cv::merge(channels, 3, image);
+				cv::cvtColor(image, image, CV_YCrCb2BGR);
+			}
+			if(!data.depthRaw().empty())
+			{
+				data.setRGBDImage(image, data.depthRaw(), data.cameraModels());
+			}
+			else if(!data.rightRaw().empty())
+			{
+				cv::Mat right;
+				if(data.rightRaw().type() == CV_8UC1)
 				{
-					data.setRGBDImage(image, data.depthRaw(), data.cameraModels());
+					cv::equalizeHist(data.rightRaw(), right);
 				}
-				else if(!data.rightRaw().empty())
+				else if(data.rightRaw().type() == CV_8UC3)
 				{
-					cv::Mat right;
-					cv::equalizeHist(data.rightRaw(), right);
-					data.setStereoImage(image, right, data.stereoCameraModels()[0]);
+					cv::Mat channels[3];
+					cv::cvtColor(data.rightRaw(), right, CV_BGR2YCrCb);
+					cv::split(right, channels);
+					cv::equalizeHist(channels[0], channels[0]);
+					cv::merge(channels, 3, right);
+					cv::cvtColor(right, right, CV_YCrCb2BGR);
 				}
+				data.setStereoImage(image, right, data.stereoCameraModels()[0]);
 			}
-			else if(_histogramMethod == 2)
+		}
+		else if(_histogramMethod == 2)
+		{
+			cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0);
+			if(data.imageRaw().type() == CV_8UC1)
 			{
-				cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0);
 				clahe->apply(data.imageRaw(), image);
-				if(!data.depthRaw().empty())
+			}
+			else if(data.imageRaw().type() == CV_8UC3)
+			{
+				cv::Mat channels[3];
+				cv::cvtColor(data.imageRaw(), image, CV_BGR2YCrCb);
+				cv::split(image, channels);
+				clahe->apply(channels[0], channels[0]);
+				cv::merge(channels, 3, image);
+				cv::cvtColor(image, image, CV_YCrCb2BGR);
+			}
+			if(!data.depthRaw().empty())
+			{
+				data.setRGBDImage(image, data.depthRaw(), data.cameraModels());
+			}
+			else if(!data.rightRaw().empty())
+			{
+				cv::Mat right;
+				if(data.rightRaw().type() == CV_8UC1)
 				{
-					data.setRGBDImage(image, data.depthRaw(), data.cameraModels());
+					clahe->apply(data.rightRaw(), right);
 				}
-				else if(!data.rightRaw().empty())
+				else if(data.rightRaw().type() == CV_8UC3)
 				{
-					cv::Mat right;
-					clahe->apply(data.rightRaw(), right);
-					data.setStereoImage(image, right, data.stereoCameraModels()[0]);
+					cv::Mat channels[3];
+					cv::cvtColor(data.rightRaw(), right, CV_BGR2YCrCb);
+					cv::split(right, channels);
+					clahe->apply(channels[0], channels[0]);
+					cv::merge(channels, 3, right);
+					cv::cvtColor(right, right, CV_YCrCb2BGR);
 				}
+				data.setStereoImage(image, right, data.stereoCameraModels()[0]);
 			}
-			if(info) info->timeHistogramEqualization = timer.ticks();
-		}
-		else
-		{
-			UWARN("Histogram equalization only supports grayscale images...");
 		}
+		if(info) info->timeHistogramEqualization = timer.ticks();
 	}
 
 	if(_stereoExposureCompensation && !data.imageRaw().empty() && !data.rightRaw().empty())