学习c#的第十五天

3D障碍物目标框（中心点坐标XYZ、长宽高lwh、朝向角theta）的非极大值抑制

#include <iostream>
#include <vector>
#include <algorithm>
#include <opencv2/opencv.hpp>

// 定义3D目标框的结构体
struct BoundingBox3D
{
    double centerX, centerY, centerZ; // 中心点坐标
    double length, width, height;     // 长宽高
    double theta;                     // 朝向角
    double score;                     // 目标框得分

    BoundingBox3D(double x, double y, double z, double l, double w, double h, double t, double s)
        : centerX(x), centerY(y), centerZ(z), length(l), width(w), height(h), theta(t), score(s) {}
};

class NMS3D
{
public:
    // 构造函数，传入IoU阈值
    NMS3D(double iouThreshold) : iouThreshold_(iouThreshold) {}

    // 执行NMS
    std::vector<BoundingBox3D> executeNMS(const std::vector<BoundingBox3D> &boxes)
    {
        std::vector<BoundingBox3D> resultBoxes;

        // 按得分降序排序
        std::vector<BoundingBox3D> sortedBoxes = sortBoxesByScore(boxes);

        // 遍历排序后的框
        while (!sortedBoxes.empty())
        {
            // 保留得分最高的框
            BoundingBox3D topBox = sortedBoxes[0];
            resultBoxes.push_back(topBox);

            // 移除与当前框IoU大于阈值的框
            sortedBoxes.erase(sortedBoxes.begin());
            sortedBoxes = removeOverlappingBoxes(topBox, sortedBoxes);
        }

        return resultBoxes;
    }

private:
    // 按得分降序排序
    std::vector<BoundingBox3D> sortBoxesByScore(const std::vector<BoundingBox3D> &boxes)
    {
        std::vector<BoundingBox3D> sortedBoxes = boxes;
        std::sort(sortedBoxes.begin(), sortedBoxes.end(),
                  [](const BoundingBox3D &a, const BoundingBox3D &b)
                  {
                      return a.score > b.score;
                  });
        return sortedBoxes;
    }

    // 移除与指定框IoU大于阈值的框
    std::vector<BoundingBox3D> removeOverlappingBoxes(const BoundingBox3D &box,
                                                      const std::vector<BoundingBox3D> &boxes)
    {
        std::vector<BoundingBox3D> filteredBoxes;
        for (const auto &b : boxes)
        {
            if (calculateIoU(box, b) < iouThreshold_)
            {
                filteredBoxes.push_back(b);
            }
        }
        return filteredBoxes;
    }

    // 计算两个框的IoU（Intersection over Union）
    double calculateIoU(const BoundingBox3D &box1, const BoundingBox3D &box2)
    {
        // 计算两个框的相交部分的体积
        double intersectionVolume = calculateIntersectionVolume(box1, box2);

        // 计算两个框的并集部分的体积
        double unionVolume = box1.length * box1.width * box1.height +
                             box2.length * box2.width * box2.height -
                             intersectionVolume;

        // 计算IoU
        return intersectionVolume / unionVolume;
    }

    // 计算两个框的相交部分的体积
    double calculateIntersectionVolume(const BoundingBox3D &box1, const BoundingBox3D &box2)
    {
        // 计算平面重叠面积
        double intersectArea = calIntersectionArea(box1, box2);
        
        double intersectHeight = calculateOverlap(box1.centerZ, box1.height, box2.centerZ, box2.height);

        // 计算相交部分的体积
        return intersectArea * intersectHeight;
    }

    cv::Point rotatePoint(const cv::Point &point, double angle)
    {
        double rotatedX = point.x * cos(angle) - point.y * sin(angle);
        double rotatedY = point.x * sin(angle) + point.y * cos(angle);

        return cv::Point(rotatedX, rotatedY);
    }

    double calIntersectionArea(const BoundingBox3D &box1, const BoundingBox3D &box2)
    {

        std::vector<cv::Point> triangle1 = {rotatePoint(cv::Point(box1.centerX - box1.width / 2, box1.centerY - box1.height / 2), box1.theta),
                                            rotatePoint(cv::Point(box1.centerX + box1.width / 2, box1.centerY - box1.height / 2), box1.theta),
                                            rotatePoint(cv::Point(box1.centerX + box1.width / 2, box1.centerY + box1.height / 2), box1.theta),
                                            rotatePoint(cv::Point(box1.centerX - box1.width / 2, box1.centerY + box1.height / 2), box1.theta)};

        std::vector<cv::Point> triangle2 = {rotatePoint(cv::Point(box2.centerX - box2.width / 2, box2.centerY - box2.height / 2), box2.theta),
                                            rotatePoint(cv::Point(box2.centerX + box2.width / 2, box2.centerY - box2.height / 2), box2.theta),
                                            rotatePoint(cv::Point(box2.centerX + box2.width / 2, box2.centerY + box2.height / 2), box2.theta),
                                            rotatePoint(cv::Point(box2.centerX - box2.width / 2, box2.centerY + box2.height / 2), box2.theta)};

        int height = -1;
        int width = -1;
        for (auto point : triangle1)
        {
            if (width < point.x)
                width = point.x;
            if (height < point.y)
                height = point.y;
        }

        for (auto point : triangle2)
        {
            if (width < point.x)
                width = point.x;
            if (height < point.y)
                height = point.y;
        }

        // 创建一个空白图像作为画布
        cv::Mat img = cv::Mat::zeros(height, width, CV_8UC3);

        // 在画布上填充两个多边形
        cv::fillConvexPoly(img, polygon1, cv::Scalar(0, 0, 1));
        cv::fillConvexPoly(img, polygon2, cv::Scalar(0, 1, 0));

        // 创建一个用于存储交集的图像
        cv::Mat intersection = cv::Mat::zeros(img.size(), img.type());

        // 计算交集
        cv::bitwise_and(img, img, intersection);

        double union_area = cv::sum(intersection)[0];
    }

    // 计算两个轴上的重叠部分长度
    double calculateOverlap(double center1, double size1, double center2, double size2)
    {
        double halfSize1 = size1 / 2;
        double halfSize2 = size2 / 2;
        double min1 = center1 - halfSize1;
        double max1 = center1 + halfSize1;
        double min2 = center2 - halfSize2;
        double max2 = center2 + halfSize2;

        // 计算重叠部分长度
        return std::max(0.0, std::min(max1, max2) - std::max(min1, min2));
    }

    double iouThreshold_; // IoU阈值
};

int main()
{
    std::vector<BoundingBox3D> inputBoxes;
    inputBoxes.push_back(BoundingBox3D(0.0, 0.0, 0.0, 2.0, 1.0, 1.0, 0, 0.9));
    inputBoxes.push_back(BoundingBox3D(0.1, 0.1, 0.1, 2.0, 1.0, 1.0, 0, 0.8));
    inputBoxes.push_back(BoundingBox3D(2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 0, 0.7));

    double iouThreshold = 0.5; // 可根据实际情况调整IoU阈值
    NMS3D nms(iouThreshold);
    std::vector<BoundingBox3D> resultBoxes = nms.executeNMS(inputBoxes);

    // 输出结果框
    for (const auto &box : resultBoxes)
    {
        std::cout << "Center: (" << box.centerX << ", " << box.centerY << ", " << box.centerZ << "), "
                  << "Dimensions: (" << box.length << ", " << box.width << ", " << box.height << "), "
                  << "Theta: " << box.theta << ", "
                  << "Score: " << box.score << std::endl;
    }

    return 0;
}