libtorch使用

#include <iostream>
#include <iomanip>
#include <sstream>
#include <chrono>
#include <random>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <torch/torch.h>
#include <torch/script.h>

// 按照指定的宽度width和高度high进行缩放和填充,以生成一个新图像
cv::Mat GetProcessedImage(cv::Mat &image, const int targetWidth, const int targetHeight) {
    cv::Mat outputImage;

    // 计算缩放比例
    float ratioWidth = static_cast<float>(targetWidth) / static_cast<float>(image.cols);
    float ratioHeight = static_cast<float>(targetHeight) / static_cast<float>(image.rows);
    float resizeScale = std::min(ratioHeight, ratioWidth);

    // 缩放后宽高
    int newShapeWidth = std::round(image.cols * resizeScale);
    int newShapeHigh = std::round(image.rows * resizeScale);
    // 需填充宽高
    float padWidth = (targetWidth - newShapeWidth) / 2.;
    float padHeight = (targetHeight - newShapeHigh) / 2.;

    int top = std::round(padHeight - 0.1);
    int bottom = std::round(padHeight + 0.1);
    int left = std::round(padWidth - 0.1);
    int right = std::round(padWidth + 0.1);

    // 缩放
    cv::resize(image, outputImage, cv::Size(newShapeWidth, newShapeHigh), 0, 0, cv::INTER_AREA);
    // 添加边框, 常量边框(BORDER_CONSTANT)
    cv::copyMakeBorder(outputImage, outputImage, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(114.));
    return outputImage;
}


// 将格式为 [x_center, y_center, width, height](即XYWH格式)的边界框坐标转换为格式为 [x_min, y_min, x_max, y_max](即XYXY格式)的边界框坐标
torch::Tensor XYWHToXYXY(const torch::Tensor& x) {
    torch::Tensor y = torch::empty_like(x);
    torch::Tensor dw = x.index({"...", 2}).div(2);
    torch::Tensor dh = x.index({"...", 3}).div(2);

    // 矩形框左上角坐标
    y.index_put_({"...", 0}, x.index({"...", 0}) - dw);
    y.index_put_({"...", 1}, x.index({"...", 1}) - dh);
    // 矩形框右下角坐标
    y.index_put_({"...", 2}, x.index({"...", 0}) + dw);
    y.index_put_({"...", 3}, x.index({"...", 1}) + dh);
    return y;
}


// Reference: https://github.com/pytorch/vision/blob/main/torchvision/csrc/ops/cpu/nms_kernel.cpp
// NMS(非极大值抑制)函数是目标检测任务中关键步骤,旨在从一组重叠的边界框中选出最佳边界框
// 接受边界框坐标(bboxes)、对应的得分(scores)以及阈值作为输入,并返回一个包含保留边界框索引的Tenso
torch::Tensor NMS(const torch::Tensor& bboxes, const torch::Tensor& scores, float IOUThreshold) {
    if (bboxes.numel() == 0) {
        return torch::empty({0}, bboxes.options().dtype(torch::kLong));
    }
   
    // 提取边界框左上角(x,y)坐标与右下角(x,y)坐标
    torch::Tensor x1_t = bboxes.select(1, 0).contiguous();
    torch::Tensor y1_t = bboxes.select(1, 1).contiguous();
    torch::Tensor x2_t = bboxes.select(1, 2).contiguous();
    torch::Tensor y2_t = bboxes.select(1, 3).contiguous();

    // 边界框面积
    torch::Tensor areas_t = (x2_t - x1_t) * (y2_t - y1_t);

    // 按得分降序对边界框进行排序
    torch::Tensor order_t = std::get<1>(scores.sort(/*stable=*/true, /*dim=*/0, /* descending=*/true));

    // 边界框数量
    int64_t ndets = bboxes.size(0);
    // 创建一个用于标记哪些边界框被抑制的Tensor, 初始化为0(未被抑制)
    torch::Tensor suppressed_t = torch::zeros({ndets}, bboxes.options().dtype(torch::kByte));
    // 创建一个用于存储保留边界框索引的Tensor,初始化为0(这里0是占位符,后续会被替换)
    torch::Tensor keep_t = torch::zeros({ndets}, bboxes.options().dtype(torch::kLong));

    // 获取指向底层数据的指针,以便高效访问
    uint8_t* suppressed = suppressed_t.data_ptr<uint8_t>();
    int64_t* keep = keep_t.data_ptr<int64_t>();
    int64_t* order = order_t.data_ptr<int64_t>();
    float* x1 = x1_t.data_ptr<float>();
    float* y1 = y1_t.data_ptr<float>();
    float* x2 = x2_t.data_ptr<float>();
    float* y2 = y2_t.data_ptr<float>();
    float* areas = areas_t.data_ptr<float>();

    // 初始化保留边界框的数量
    int64_t numToKeep = 0;
    for (int64_t i = 0; i < ndets; i++) {
        // 获取当前边界框的索引
        int64_t index = order[i];
        // 如果当前边界框已被抑制,则跳过
        if (suppressed[index] == 1) {
            continue;
        }

        // 将当前边界框的索引保存到保留列表中
        keep[numToKeep++] = index;
        // 提取当前边界框的坐标和面积
        float ix1 = x1[index];
        float iy1 = y1[index];
        float ix2 = x2[index];
        float iy2 = y2[index];
        float iarea = areas[index];
        for (int64_t j = i + 1; j < ndets; j++) {
            int64_t k = order[j];
            if (suppressed[k] == 1) {
                continue;
            }
            float xx1 = std::max(ix1, x1[k]);
            float yy1 = std::max(iy1, y1[k]);
            float xx2 = std::min(ix2, x2[k]);
            float yy2 = std::min(iy2, y2[k]);

            float w = std::max(static_cast<float>(0), xx2 - xx1);
            float h = std::max(static_cast<float>(0), yy2 - yy1);
            float inter = w * h;
            float ovr = inter / (iarea + areas[k] - inter);
            if (ovr > IOUThreshold) {
                suppressed[k] = 1;
            }
        }
    }
    // 返回包含保留边界框索引的Tensor(只包含前numToKeep个元素)
    return keep_t.narrow(0, 0, numToKeep);
}

// prediction: 一个四维张量, 通常形状为[batch_size, num_anchors * (num_CLASSES + 5), height, width], 其中5代表边界框的坐标(通常是中心点坐标加上宽高)和置信度, num_CLASSES是类别数
// conf_thres: 置信度阈值, 低于此阈值的检测框将被忽略
// iou_thres: IoU阈值, 用于判断两个检测框是否重叠过多, 如果重叠超过此阈值, 则较低置信度的检测框将被抑制
// max_det: 每个图像最多保留的检测框数量
torch::Tensor NonMaxSuppression(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
    // 批次大小
    int64_t bs = prediction.size(0);
    // 类别数
    int64_t nc = prediction.size(1) - 4;
    // 每个锚点的额外信息数量
    int64_t nm = prediction.size(1) - nc - 4;
    // 边界框坐标和类别置信度的总维度
    int64_t mi = 4 + nc;

    // 计算超过置信度阈值的掩码, 使用amax(1)找到每行(每个锚点)中的最大置信度,并与置信度阈值进行比较
    torch::Tensor xc = prediction.index({torch::indexing::Slice(), torch::indexing::Slice(4, mi)}).amax(1) > conf_thres;

    // 调整预测张量的形状,将最后一个维度(通常是坐标和置信度)移动到倒数第二个维度, 方便后续对边界框和置信度的处理
    prediction = prediction.transpose(-1, -2);
    // 将边界框的格式从xywh(中心点坐标+宽高)转换为xyxy(左上角和右下角坐标)
    prediction.index_put_({"...", torch::indexing::Slice({torch::indexing::None, 4})}, XYWHToXYXY(prediction.index({"...", torch::indexing::Slice(torch::indexing::None, 4)})));

    // 初始化输出张量列表,用于存储每个批次的非极大值抑制结果
    std::vector<torch::Tensor> output;
    for (int i = 0; i < bs; i++) {
        // 为每个批次创建一个空的输出张量,其形状为(0, 6+nm),其中6代表边界框坐标、置信度和类别索引,nm代表额外信息数量
        output.push_back(torch::zeros({0, 6 + nm}, prediction.device()));
    }

    // 遍历每个批次进行处理
    for (int xi = 0; xi < prediction.size(0); xi++) {
        // 获取当前批次的预测张量
        torch::Tensor x = prediction[xi];
        // 使用之前计算的掩码过滤掉低于置信度阈值的边界框
        x = x.index({xc[xi]});
        // 分离边界框、类别置信度和额外信息
        // 边界框坐标
        torch::Tensor box = x.narrow(1, 0, 4);
        // 类别置信度和额外信息
        torch::Tensor cls_conf_mask = x.narrow(1, 4, nc + nm);
        // 类别置信度
        torch::Tensor cls = x.narrow(1, 4, nc + nm).narrow(1, 0, nc);
        // 额外信息
        torch::Tensor mask = x.narrow(1, 4, nc + nm).narrow(1, nc, nm);

        // 找到每个边界框的最高置信度类别及其索引
        auto [conf, j] = cls.max(1, true);
        // 将边界框、最高置信度、类别索引和额外信息合并为一个新的张量
        x = torch::cat({box, conf, j.toType(torch::kFloat), mask}, 1);
        // 再次使用置信度阈值过滤边界框
        x = x.index({conf.view(-1) > conf_thres});
        // 如果当前批次有剩余的边界框,则进行非极大值抑制
        if (x.size(0)) {
            // NMS
            torch::Tensor c = x.index({torch::indexing::Slice(), torch::indexing::Slice{5, 6}}) * 7680;
            torch::Tensor boxes = x.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, 4)}) + c;
            // 置信度作为分数
            torch::Tensor scores = x.index({torch::indexing::Slice(), 4});
            // 调用NMS函数进行非极大值抑制,返回保留的边界框索引
            torch::Tensor indices  = NMS(boxes, scores, iou_thres);
            // 限制保留的边界框数量不超过max_det
            indices = indices .index({torch::indexing::Slice(torch::indexing::None, max_det)});
            // 根据保留的索引从原始张量x中选择对应的边界框
            output[xi] = x.index({indices});
        }
    }

    // 将所有批次的输出张量堆叠成一个新的张量并返回
    return torch::stack(output);
}


// 根据图像尺寸的变化来调整边界框(bounding boxes)的尺寸和位置
torch::Tensor ScaleBoxes(const std::vector<int>& image1Shape, torch::Tensor& boxes, const std::vector<int>& image0Shape) {
    // 计算缩放比例
    float gain = (std::min)((float)image1Shape[0] / image0Shape[0], (float)image1Shape[1] / image0Shape[1]);
    // 计算填充值
    float pad0 = std::round((float)(image1Shape[1] - image0Shape[1] * gain) / 2. - 0.1);
    float pad1 = std::round((float)(image1Shape[0] - image0Shape[0] * gain) / 2. - 0.1);
    // 调整边界框位置
    boxes.index_put_({"...", 0}, boxes.index({"...", 0}) - pad0);
    boxes.index_put_({"...", 2}, boxes.index({"...", 2}) - pad0);
    boxes.index_put_({"...", 1}, boxes.index({"...", 1}) - pad1);
    boxes.index_put_({"...", 3}, boxes.index({"...", 3}) - pad1);
    // 调整边界框尺寸
    boxes.index_put_({"...", torch::indexing::Slice(torch::indexing::None, 4)}, boxes.index({"...", torch::indexing::Slice(torch::indexing::None, 4)}).div(gain));
    return boxes;
}


// 生产随机色
cv::Scalar GetRandomColor() {
    // 创建一个随机设备用于生成非确定性种子
    std::random_device rd;
    // 使用随机设备生成的种子初始化 Mersenne Twister 生成器
    std::mt19937 gen(rd());
    // 创建一个在 0 到 255 之间均匀分布的整数随机数生成器
    std::uniform_int_distribution<int> dis(0, 255);
    int r = dis(gen);
    int g = dis(gen);
    int b = dis(gen);
    // 用于opencv
    return cv::Scalar(b, g, r);
}


int main(int argc, char* argv[]) {
    const int BOX_FONT_THICKNESS     = 2;
    const int NAME_FONT_THICKNESS    = 2;
    const double NAME_FONT_SCALE     = 0.8;
    const int PROCESSED_IMAGE_WIDTH  = 640;
    const int PROCESSED_IMAGE_HEIGHT = 640;

    const std::string MODEL_PATH  = "yolov8m.torchscript";

    const std::vector<std::string> CLASSES {
        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
        "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
        "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
        "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife",
        "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
        "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
    };

    std::chrono::high_resolution_clock::time_point start = std::chrono::high_resolution_clock::now();

    if (argc < 2) {
        std::cout << "please input image path" << std::endl;
        return 0;
    }

    try {       
        torch::jit::script::Module yolov8Model = torch::jit::load(MODEL_PATH);
        // 将模型设置为评估模式
        yolov8Model.eval();
        // 将模型和数据移动到CPU上,并设置数据类型为Float32
        yolov8Model.to(torch::kCPU, torch::kFloat32);

        cv::Mat baseImage = cv::imread(std::string(argv[1]));
        cv::Mat processedImage = GetProcessedImage(baseImage, PROCESSED_IMAGE_WIDTH, PROCESSED_IMAGE_HEIGHT);

        // 将处理后的图像转换为PyTorch张量
        torch::Tensor imageTensor = torch::from_blob(
            processedImage.data,
            // 张量的形状,即图像的高度、宽度和通道数
            { processedImage.rows, processedImage.cols, 3 },
            torch::kByte
        ).to(torch::kCPU);

        // 将张量的数据类型转换为Float32, 然后将张量中的每个值除以255, 以将其从[0, 255]范围归一化到[0, 1]范围
        imageTensor = imageTensor.toType(torch::kFloat32).div(255);
        // 对张量改变其维度顺序, 转换为pytorch中CHW格式
        imageTensor = imageTensor.permute({2, 0, 1});
        // 在张量前面增加一个批次维度, 在0维上增加一个大小为1的维度
        imageTensor = imageTensor.unsqueeze(0);

        // 将处理好的图像张量封装在一个std::vector<torch::jit::IValue>中,作为模型的输入
        std::vector<torch::jit::IValue> inputs {imageTensor};
        // 对模型进行前向传播,获取输出然后将输出转换为一个张量,并移动到CPU上
        torch::Tensor output = yolov8Model.forward(inputs).toTensor().cpu();

        // NMS
        torch::Tensor keep = NonMaxSuppression(output)[0];
        torch::Tensor boxes = keep.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, 4)});
        keep.index_put_({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, 4)}, ScaleBoxes({processedImage.rows, processedImage.cols}, boxes, {baseImage.rows, baseImage.cols}));

        // 设置输出格式为固定小数点表示法,并精确到小数点后两位
        std::ostringstream out;
        for (int i = 0; i < keep.size(0); i++) {
            int x1 = keep[i][0].item().toFloat();
            int y1 = keep[i][1].item().toFloat();
            int x2 = keep[i][2].item().toFloat();
            int y2 = keep[i][3].item().toFloat();
            float conf = keep[i][4].item().toFloat();
            int cls = keep[i][5].item().toInt();
            if (conf >= 0.5) {
                std::cout << "rect: [" << x1 << "," << y1 << "," << x2 << "," << y2 << "]  conf: " << conf << "  class: " << CLASSES[cls] << std::endl;

                cv::Scalar color =  GetRandomColor();
                cv::rectangle(baseImage, cv::Rect(cv::Point(x1, y1), cv::Point(x2, y2)), color, BOX_FONT_THICKNESS);

                out << std::fixed << std::setprecision(2) << conf;
                cv::putText(baseImage, CLASSES[cls] + ":" + out.str(), cv::Point(x1, y1 - 5), cv::FONT_HERSHEY_SIMPLEX, NAME_FONT_SCALE, color, NAME_FONT_THICKNESS, cv::LINE_AA);
                out.str("");
                out.clear();
            }
        }

        cv::imwrite("detected_image.jpg", baseImage);
    } catch (const c10::Error& e) {
        std::cout << e.msg() << std::endl;
    }

    std::chrono::high_resolution_clock::time_point end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed_seconds = end - start;
    std::cout << "used time: " << elapsed_seconds.count() << " s" << std::endl;
    return 0;
}