#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);
}