videoStreamInfer/src/nodes/draw/drawNode.cpp

#include "common/utils.hpp"
#include "nodes/base/base.hpp"
#include "nodes/draw/drawNode.hpp"
#include "nodes/draw/position.hpp"
#include <opencv2/opencv.hpp>
#include <chrono>
#include <tuple>

namespace GNode
{

const std::vector<std::pair<int, int>> coco_pairs = {
    {0, 1}, {0, 2}, {0, 11}, {0, 12}, {1, 3}, {2, 4},
    {5, 6}, {5, 7}, {7, 9}, {6, 8}, {8, 10},
    {11, 12}, {5, 11}, {6, 12},
    {11, 13}, {13, 15}, {12, 14}, {14, 16}
};

static std::tuple<int, int, int> getFontSize(const std::string& text)
{
    int baseline = 0;
    cv::Size textSize = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 1.0, 2, &baseline);
    // std::cout << textSize << std::endl;
    return std::make_tuple(textSize.width, textSize.height, baseline);
}

static void overlay_mask(
    cv::Mat& image, const cv::Mat& smallMask,
    int roiX, int roiY,
    const cv::Scalar& color, double alpha) 
{
    if (image.empty() || smallMask.empty() ||
        image.type() != CV_8UC3 || smallMask.type() != CV_8UC1) {
        return;
    }
    alpha = std::max(0.0, std::min(1.0, alpha));
    
    cv::Rect roiRect(roiX, roiY, smallMask.cols, smallMask.rows);
    
    cv::Rect imageRect(0, 0, image.cols, image.rows);
    cv::Rect intersectionRect = roiRect & imageRect; // 使用 & 操作符计算交集
    
    if (intersectionRect.width <= 0 || intersectionRect.height <= 0) {
        return;
    }
    
    cv::Mat originalROI = image(intersectionRect); // ROI 指向 image 的数据
    
    int maskStartX = intersectionRect.x - roiX;
    int maskStartY = intersectionRect.y - roiY;
    cv::Rect maskIntersectionRect(maskStartX, maskStartY, intersectionRect.width, intersectionRect.height);
    cv::Mat smallMaskROI = smallMask(maskIntersectionRect);
    
    cv::Mat colorPatchROI(intersectionRect.size(), image.type(), color);
    
    cv::Mat tempColoredROI = originalROI.clone(); // 需要一个临时区域进行覆盖
    colorPatchROI.copyTo(tempColoredROI, smallMaskROI);
    
    cv::addWeighted(originalROI, 1.0 - alpha, tempColoredROI, alpha, 0.0, originalROI);
}


static std::tuple<uint8_t, uint8_t, uint8_t> hsv2bgr(float h, float s, float v) 
{
    const int h_i = static_cast<int>(h * 6);
    const float f = h * 6 - h_i;
    const float p = v * (1 - s);
    const float q = v * (1 - f * s);
    const float t = v * (1 - (1 - f) * s);
    float r, g, b;
    switch (h_i) {
      case 0:
        r = v, g = t, b = p;
        break;
      case 1:
        r = q, g = v, b = p;
        break;
      case 2:
        r = p, g = v, b = t;
        break;
      case 3:
        r = p, g = q, b = v;
        break;
      case 4:
        r = t, g = p, b = v;
        break;
      case 5:
        r = v, g = p, b = q;
        break;
      default:
        r = 1, g = 1, b = 1;
        break;
    }
    return std::make_tuple(static_cast<uint8_t>(b * 255), static_cast<uint8_t>(g * 255),
                      static_cast<uint8_t>(r * 255));
  }

static std::tuple<uint8_t, uint8_t, uint8_t> random_color(int id) 
{
    float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f;
    float s_plane = ((((unsigned int)id << 3) ^ 0x315793) % 100) / 100.0f;
    return hsv2bgr(h_plane, s_plane, 1);
}


static float box_iou(const data::Box& box1, const data::Box& box2) 
{
    float inter_x1 = std::max(box1.left, box2.left);
    float inter_y1 = std::max(box1.top, box2.top);
    float inter_x2 = std::min(box1.right, box2.right);
    float inter_y2 = std::min(box1.bottom, box2.bottom);

    float inter_area = std::max(0.0f, inter_x2 - inter_x1) * std::max(0.0f, inter_y2 - inter_y1);
    float box1_area = (box1.right - box1.left) * (box1.bottom - box1.top);
    float box2_area = (box2.right - box2.left) * (box2.bottom - box2.top);

    return inter_area / (box1_area + box2_area - inter_area + 1e-6f); // Avoid division by zero
}

void DrawNode::handle_data(std::shared_ptr<meta::BaseData>& meta_data) 
{
    auto frame_data = std::dynamic_pointer_cast<meta::FrameData>(meta_data);
    bool show_track = true;
    cv::Mat image = frame_data->image.clone();
    int image_width = image.cols;
    int image_height = image.rows;
    PositionManager<float> pm(getFontSize);
    for (auto& box : frame_data->boxes)
    {
        if (show_track)
        {
          float max_iou = 0.0f;
          for (const auto& track_box : frame_data->track_boxes)
          {
            float iou = box_iou(box, track_box);
            if (iou > max_iou && iou > 0.7f)
            {
              max_iou = iou;
              box.class_id = track_box.class_id;
            }   
          }
        }
        uint8_t b, g, r;
        std::tie(b, g, r) = random_color(box.class_id);
        cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(b, g, r), 2);

        std::tuple<int, int, int, int> pbox = std::make_tuple(box.left, box.top, box.right, box.bottom);
        int x, y;
        std::string text;
        if (box.class_id >= 0) 
        {
            text = str_format("%s %.2f id=%d", box.label.c_str(), box.score, box.class_id);
        }
        else
        {
            text = str_format("%s %.2f", box.label.c_str(), box.score);
        }
        std::tie(x, y) = pm.selectOptimalPosition(pbox, image_width, image_height, text);
        
        cv::putText(image, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2);

        if (!box.seg_mask.empty())
        {
            overlay_mask(image, box.seg_mask, box.left, box.top, cv::Scalar(b, g, r), 0.6);
        }
        if (!box.keypoints.empty())
        {
            for (const auto& point : box.keypoints)
            {
                cv::circle(image, cv::Point(point.x, point.y), 5, cv::Scalar(b, g, r), -1);
            }
            for (const auto& pair : coco_pairs) 
            {
                int startIdx = pair.first;
                int endIdx = pair.second;

                if (startIdx < box.keypoints.size() && endIdx < box.keypoints.size()) 
                {
                    int x1 = (int)box.keypoints[startIdx].x;
                    int y1 = (int)box.keypoints[startIdx].y;
                    int x2 = (int)box.keypoints[endIdx].x;
                    int y2 = (int)box.keypoints[endIdx].y;
                    cv::line(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(b, g, r), 2);
                }
            }
        }
    }
    frame_data->draw_image = image;
}

}