nanocode/web-opencv/cpp/ColorReader.cpp

#include "ColorReader.h"
#include "GlobalParams.h"
#include <cmath>
#include <algorithm>

ColorReader::ColorReader() {
}

ColorReader::~ColorReader() {
}

std::vector<ColorResult> ColorReader::readColors(
    const cv::Mat& image,
    const std::vector<cv::Point2f>& centers,
    bool smallLocatorOnLeft,
    bool applyWB) {
    
    std::vector<ColorResult> results;
    
    if (image.empty() || centers.empty()) {
        return results;
    }
    
    // 复制图像用于颜色读取（避免绘图干扰）
    cv::Mat colorImage = image.clone();
    
    // 白平衡处理（可选）
    if (applyWB) {
        applyWhiteBalance(colorImage);
    }
    
    // 确定采样半径
    int radius = GlobalParams::DEFAULT_SAMPLE_RADIUS;
    
    // 采样每个色块的颜色
    for (size_t i = 0; i < centers.size(); ++i) {
        ColorResult result;
        result.id = static_cast<int>(i);
        
        // 采样颜色
        cv::Scalar bgr = sampleColor(colorImage, centers[i], radius);
        
        // 转换颜色空间
        convertColorSpaces(bgr, result.h, result.s, result.v,
                          result.l, result.a, result.b);
        
        result.r = bgr[2];
        result.g = bgr[1];
        result.b_val = bgr[0];
        
        results.push_back(result);
    }
    
    // 调整顺序（确保从靠近小定位块的色块开始）
    results = adjustOrder(results, smallLocatorOnLeft);
    
    return results;
}

void ColorReader::applyWhiteBalance(cv::Mat& image) {
    if (image.empty()) {
        return;
    }
    
    // 确保是BGR格式
    cv::Mat bgrImage;
    if (image.channels() == 4) {
        cv::cvtColor(image, bgrImage, cv::COLOR_BGRA2BGR);
    } else if (image.channels() == 1) {
        cv::cvtColor(image, bgrImage, cv::COLOR_GRAY2BGR);
    } else {
        bgrImage = image;
    }
    
    // 计算各通道均值
    cv::Scalar sumChannels = cv::sum(bgrImage);
    int totalPixels = bgrImage.rows * bgrImage.cols;
    
    if (totalPixels == 0) {
        return;
    }
    
    double avgB = sumChannels[0] / totalPixels;
    double avgG = sumChannels[1] / totalPixels;
    double avgR = sumChannels[2] / totalPixels;
    
    // 避免除零
    if (avgB < 1) avgB = 1;
    if (avgG < 1) avgG = 1;
    if (avgR < 1) avgR = 1;
    
    // 计算目标均值（三通道平均）
    double targetAvg = (avgB + avgG + avgR) / 3.0;
    
    if (targetAvg < 1) {
        return;
    }
    
    // 各通道缩放校正
    std::vector<cv::Mat> channels(3);
    cv::split(bgrImage, channels);
    
    channels[0] = channels[0] * (targetAvg / avgB);
    channels[1] = channels[1] * (targetAvg / avgG);
    channels[2] = channels[2] * (targetAvg / avgR);
    
    cv::merge(channels, bgrImage);
    bgrImage.convertTo(image, CV_8UC3);
}

cv::Scalar ColorReader::sampleColor(const cv::Mat& image, const cv::Point2f& center, int radius) {
    int cx = static_cast<int>(center.x);
    int cy = static_cast<int>(center.y);
    
    // 确保采样区域在图像内
    int x1 = std::max(0, cx - radius);
    int y1 = std::max(0, cy - radius);
    int x2 = std::min(image.cols - 1, cx + radius);
    int y2 = std::min(image.rows - 1, cy + radius);
    
    if (x2 <= x1 || y2 <= y1) {
        return cv::Scalar(0, 0, 0);
    }
    
    // 创建采样掩码（圆形）
    cv::Mat roi = image(cv::Rect(x1, y1, x2 - x1 + 1, y2 - y1 + 1));
    cv::Mat mask = cv::Mat::zeros(roi.size(), CV_8U);
    cv::Point maskCenter(cx - x1, cy - y1);
    cv::circle(mask, maskCenter, radius, cv::Scalar(255), -1);
    
    // 计算掩码区域的平均颜色
    cv::Scalar avgColor = cv::mean(roi, mask);
    
    return avgColor;
}

void ColorReader::convertColorSpaces(const cv::Scalar& bgr, 
                                      double& h, double& s, double& v,
                                      double& l, double& a, double& b) {
    // 创建1x1像素的图像进行颜色转换
    cv::Mat bgrMat(1, 1, CV_8UC3, cv::Scalar(bgr[0], bgr[1], bgr[2]));
    
    // 转换到HSV
    cv::Mat hsvMat;
    cv::cvtColor(bgrMat, hsvMat, cv::COLOR_BGR2HSV);
    cv::Vec3b hsv = hsvMat.at<cv::Vec3b>(0, 0);
    h = hsv[0] * 2.0;  // OpenCV的H范围是0-180，转换为0-360
    s = hsv[1];
    v = hsv[2];
    
    // 转换到LAB
    cv::Mat labMat;
    cv::cvtColor(bgrMat, labMat, cv::COLOR_BGR2Lab);
    cv::Vec3b lab = labMat.at<cv::Vec3b>(0, 0);
    l = lab[0];
    a = static_cast<double>(lab[1]) - 128.0;  // a范围从0-255转为-128到127
    b = static_cast<double>(lab[2]) - 128.0;  // b范围从0-255转为-128到127
}

std::vector<ColorResult> ColorReader::adjustOrder(
    std::vector<ColorResult>& colors,
    bool smallLocatorOnLeft) {
    
    // 如果小定位块在左边，颜色顺序已经是正确的
    // 如果小定位块在右边，需要反转顺序
    
    if (!smallLocatorOnLeft && colors.size() > 1) {
        std::reverse(colors.begin(), colors.end());
    }
    
    // 重新编号
    for (size_t i = 0; i < colors.size(); ++i) {
        colors[i].id = static_cast<int>(i);
    }
    
    return colors;
}