56 lines
2.0 KiB
C++
56 lines
2.0 KiB
C++
#include "ImagePreprocessor.h"
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
|
|
ImagePreprocessor::ImagePreprocessor() {
|
|
// Constructor implementation
|
|
}
|
|
|
|
ImagePreprocessor::~ImagePreprocessor() {
|
|
// Destructor implementation
|
|
}
|
|
|
|
cv::Mat ImagePreprocessor::get_mask(const cv::Mat& frame, const std::string& target_color) {
|
|
cv::Mat hsv_frame;
|
|
cv::cvtColor(frame, hsv_frame, cv::COLOR_BGR2HSV);
|
|
|
|
cv::Mat mask;
|
|
if (target_color == "red") {
|
|
// Red color range in HSV
|
|
cv::Mat mask1, mask2;
|
|
cv::inRange(hsv_frame, cv::Scalar(0, 43, 49), cv::Scalar(25, 255, 255), mask1);
|
|
cv::inRange(hsv_frame, cv::Scalar(156, 46, 49), cv::Scalar(180, 255, 255), mask2);
|
|
mask = mask1 | mask2;
|
|
} else if (target_color == "blue") {
|
|
// Blue color range in HSV
|
|
cv::inRange(hsv_frame, cv::Scalar(83, 200, 44), cv::Scalar(130, 255, 255), mask);
|
|
} else {
|
|
// Default to empty mask if color not recognized
|
|
mask = cv::Mat::zeros(frame.size(), CV_8UC1);
|
|
}
|
|
|
|
// Apply morphological operations to reduce noise
|
|
cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
|
|
cv::morphologyEx(mask, mask, cv::MORPH_OPEN, kernel);
|
|
cv::morphologyEx(mask, mask, cv::MORPH_CLOSE, kernel);
|
|
|
|
return mask;
|
|
}
|
|
|
|
cv::Mat ImagePreprocessor::get_color_only_frame(const cv::Mat& frame, const std::string& target_color) {
|
|
cv::Mat mask = get_mask(frame, target_color);
|
|
|
|
// 不再反转掩膜，直接使用原始掩膜
|
|
// 目标颜色区域为白色(255)，非目标颜色区域为黑色(0)
|
|
|
|
cv::Mat color_only_frame;
|
|
frame.copyTo(color_only_frame, mask); // 使用掩膜复制原始帧，只保留目标颜色部分
|
|
|
|
return color_only_frame;
|
|
}
|
|
|
|
void ImagePreprocessor::process_frame(const cv::Mat& frame, const std::string& target_color,
|
|
cv::Mat& mask, cv::Mat& color_only_frame) {
|
|
mask = get_mask(frame, target_color);
|
|
color_only_frame = get_color_only_frame(frame, target_color);
|
|
} |