uvos/SDImagePreprocess
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add face recognition support to the system

Browse source
This commit is contained in:
uvos 2024-04-05 11:24:04 +02:00
parent b2ffbfa530
commit a279001151
5 changed files with 283 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

60
main.cpp
View file

@ -6,6 +6,7 @@
#include <opencv2/highgui.hpp>
#include <algorithm>
#include <execution>
#include <string>
#include <vector>
#include <numeric>
@ -15,6 +16,7 @@
#include "utils.h"
#include "intelligentroi.h"
#include "seamcarving.h"
#include "facerecognizer.h"
const Yolo::Detection* pointInDetectionHoriz(int x, const std::vector<Yolo::Detection>& detections, const Yolo::Detection* ignore = nullptr)
{
@ -223,7 +225,7 @@ void drawDebugInfo(cv::Mat &image, const cv::Rect& rect, const std::vector<Yolo:
for(const Yolo::Detection& detection : detections)
{
cv::rectangle(image, detection.box, detection.color, 3);
std::string label = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
std::string label = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4) + ' ' + std::to_string(detection.priority);
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX, 1, 1, 0);
cv::Rect textBox(detection.box.x, detection.box.y - 40, labelSize.width + 10, labelSize.height + 20);
cv::rectangle(image, textBox, detection.color, cv::FILLED);
@ -251,7 +253,8 @@ static void reduceSize(cv::Mat& image, const cv::Size& targetSize)
}
}
void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yolo, const std::filesystem::path& debugOutputPath)
void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yolo, std::mutex& yoloMutex, FaceRecognizer* recognizer,
std::mutex& reconizerMutex, const std::filesystem::path& debugOutputPath)
{
InteligentRoi intRoi(yolo);
cv::Mat image = cv::imread(path);
@ -263,11 +266,28 @@ 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(const Yolo::Detection& detection : detections)
Log(Log::DEBUG)<<detection.class_id<<": "<<detection.className<<" at "<<detection.box<<" with prio "<<detection.priority;
for(Yolo::Detection& detection : detections)
{
bool hasmatch = false;
if(recognizer && detection.className == "person")
{
cv::Mat person = image(detection.box);
reconizerMutex.lock();
std::pair<int, double> match = recognizer->isMatch(person);
reconizerMutex.unlock();
if(match.first >= 0)
{
detection.priority += 10;
hasmatch = true;
}
}
Log(Log::DEBUG)<<detection.class_id<<": "<<detection.className<<" at "<<detection.box<<" with prio "<<detection.priority<<(hasmatch ? " has match" : "");
}
cv::Rect crop;
bool incompleate = intRoi.getCropRectangle(crop, detections, image.size());
@ -276,7 +296,11 @@ 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();
}
}
cv::Mat croppedImage;
@ -306,7 +330,7 @@ void pipeline(const std::filesystem::path& path, const Config& config, Yolo& yol
}
cv::Mat resizedImage;
cv::resize(croppedImage, resizedImage, {512, 512}, 0, 0, cv::INTER_CUBIC);
cv::resize(croppedImage, resizedImage, config.targetSize, 0, 0, cv::INTER_CUBIC);
bool ret = cv::imwrite(config.outputDir/path.filename(), resizedImage);
if(!ret)
Log(Log::WARN)<<"could not save image to "<<config.outputDir/path.filename()<<" skipping";
@ -346,7 +370,7 @@ int main(int argc, char* argv[])
return 1;
}
Yolo yolo(config.modelPath, {640, 480}, config.classesPath, false);
Yolo yolo(config.modelPath, {640, 480}, config.classesPath, true);
if(!std::filesystem::exists(config.outputDir))
{
@ -364,8 +388,28 @@ int main(int argc, char* argv[])
std::filesystem::create_directory(debugOutputPath);
}
std::for_each(std::execution::parallel_unsequenced_policy(),
imagePaths.begin(), imagePaths.end(), [&yolo, &debugOutputPath, &config](const std::filesystem::path& path){pipeline(path, config, yolo, debugOutputPath);});
FaceRecognizer* recognizer = nullptr;
std::mutex recognizerMutex;
if(!config.focusPersonImage.empty())
{
cv::Mat personImage = cv::imread(config.focusPersonImage);
if(personImage.empty())
{
Log(Log::ERROR)<<"Could not load image from "<<config.focusPersonImage;
return 1;
}
recognizer = new FaceRecognizer();
recognizer->addReferances({personImage});
recognizer->setThreshold(config.threshold);
}
std::mutex yoloMutex;
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);
return 0;
}