From 81475815fba9793c350f89535697238a08eab1ab Mon Sep 17 00:00:00 2001
From: uvos <devnull@uvos.xyz>
Date: Fri, 5 Apr 2024 12:23:11 +0200
Subject: [PATCH] more effectively multithread

---
 main.cpp | 50 +++++++++++++++++++++++++++++++++++++-------------
 yolo.cpp |  2 --
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/main.cpp b/main.cpp
index 834bfff..a368aef 100644
--- a/main.cpp
+++ b/main.cpp
@@ -253,7 +253,7 @@ static void reduceSize(cv::Mat& image, const cv::Size& targetSize)
 	}
 }
 
-void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yolo, std::mutex& yoloMutex, FaceRecognizer* recognizer,
+void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yolo, FaceRecognizer* recognizer,
 	std::mutex& reconizerMutex, const std::filesystem::path& debugOutputPath)
 {
 	InteligentRoi intRoi(yolo);
@@ -266,9 +266,7 @@ void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yol
 
 	reduceSize(image, config.targetSize);
 
-	yoloMutex.lock();
 	std::vector<Yolo::Detection> detections = yolo.runInference(image);
-	yoloMutex.unlock();
 
 	Log(Log::DEBUG)<<"Got "<<detections.size()<<" detections for "<<path;
 	for(Yolo::Detection& detection : detections)
@@ -297,9 +295,7 @@ void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yol
 		bool ret = seamCarveResize(image, detections, config.targetSize.aspectRatio());
 		if(ret && image.size().aspectRatio() != config.targetSize.aspectRatio())
 		{
-			yoloMutex.lock();
 			detections = yolo.runInference(image);
-			yoloMutex.unlock();
 		}
 	}
 
@@ -336,6 +332,35 @@ void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yol
 		Log(Log::WARN)<<"could not save image to "<<config.outputDir/path.filename()<<" skipping";
 }
 
+void threadFn(const std::vector<std::filesystem::path>& images, const Config& config, FaceRecognizer* recognizer,
+		std::mutex& reconizerMutex, const std::filesystem::path& debugOutputPath)
+{
+	Yolo yolo(config.modelPath, {640, 480}, config.classesPath, false);
+	for(std::filesystem::path path : images)
+		pipeline(path, config, yolo, recognizer, reconizerMutex, debugOutputPath);
+}
+
+template<typename T>
+std::vector<std::vector<T>> splitVector(const std::vector<T>& vec, size_t parts)
+{
+	std::vector<std::vector<T>> out;
+
+	size_t length = vec.size()/parts;
+	size_t remain = vec.size() % parts;
+
+	size_t begin = 0;
+	size_t end = 0;
+
+	for (size_t i = 0; i < std::min(parts, vec.size()); ++i)
+	{
+		end += (remain > 0) ? (length + !!(remain--)) : length;
+		out.push_back(std::vector<T>(vec.begin() + begin, vec.begin() + end));
+		begin = end;
+	}
+
+	return out;
+}
+
 int main(int argc, char* argv[])
 {
 	Log::level = Log::INFO;
@@ -370,8 +395,6 @@ int main(int argc, char* argv[])
 		return 1;
 	}
 
-	Yolo yolo(config.modelPath, {640, 480}, config.classesPath, true);
-
 	if(!std::filesystem::exists(config.outputDir))
 	{
 		if(!std::filesystem::create_directory(config.outputDir))
@@ -403,13 +426,14 @@ int main(int argc, char* argv[])
 		recognizer->setThreshold(config.threshold);
 	}
 
-	std::mutex yoloMutex;
+	std::vector<std::thread> threads;
+	std::vector<std::vector<std::filesystem::path>> imagePathParts = splitVector(imagePaths, std::thread::hardware_concurrency());
 
-	auto pipelineLambda = [&yolo, &debugOutputPath, &config, &yoloMutex, &recognizer, &recognizerMutex](const std::filesystem::path& path)
-	{
-		pipeline(path, config, yolo, yoloMutex, recognizer, recognizerMutex, debugOutputPath);
-	};
-	std::for_each(std::execution::par_unseq, imagePaths.begin(), imagePaths.end(), pipelineLambda);
+	for(size_t i = 0; i < std::thread::hardware_concurrency(); ++i)
+		threads.push_back(std::thread(threadFn, imagePathParts[i], std::ref(config),  recognizer, std::ref(recognizerMutex), std::ref(debugOutputPath)));
+
+	for(std::thread& thread : threads)
+		thread.join();
 
 	return 0;
 }
diff --git a/yolo.cpp b/yolo.cpp
index 74d9f2a..9379ddd 100644
--- a/yolo.cpp
+++ b/yolo.cpp
@@ -42,13 +42,11 @@ Yolo::Yolo(const std::filesystem::path &onnxModelPath, const cv::Size &modelInpu
 	}
 	if(runWithOCl)
 	{
-		std::cout << "\nRunning on OCV" << std::endl;
 		net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
 		net.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL);
 	}
 	else
 	{
-		std::cout << "\nRunning on CPU" << std::endl;
 		net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
 		net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
 	}