sladro
e13cb3659c
- 创建基本项目结构和目录 - 添加CMake构建系统 - 实现基础的配置解析功能 - 添加YOLO推理框架支持 - 集成RTSP和视频流处理功能 - 添加性能监控和日志系统
131 lines
4.2 KiB
Plaintext
131 lines
4.2 KiB
Plaintext
#include <cuda_runtime.h>
|
|
#include "preprocess.hpp"
|
|
#include "../common/logger.hpp"
|
|
|
|
namespace pipeline {
|
|
namespace detail {
|
|
|
|
// CUDA kernel配置
|
|
constexpr int BLOCK_SIZE = 256;
|
|
|
|
// 计算网格大小
|
|
inline int divUp(int total, int block) {
|
|
return (total + block - 1) / block;
|
|
}
|
|
|
|
// BGR转RGB并归一化的kernel
|
|
__global__ void bgr2rgbAndNormalizeKernel(
|
|
const unsigned char* input,
|
|
float* output,
|
|
int height,
|
|
int width,
|
|
float scale
|
|
) {
|
|
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
const int total = height * width;
|
|
|
|
if (idx < total) {
|
|
const int h = idx / width;
|
|
const int w = idx % width;
|
|
const int input_idx = (h * width + w) * 3;
|
|
const int output_idx = (h * width + w) * 3;
|
|
|
|
// BGR转RGB并归一化
|
|
output[output_idx + 0] = input[input_idx + 2] * scale; // R = B
|
|
output[output_idx + 1] = input[input_idx + 1] * scale; // G = G
|
|
output[output_idx + 2] = input[input_idx + 0] * scale; // B = R
|
|
}
|
|
}
|
|
|
|
// 批处理kernel
|
|
__global__ void batchPreprocessKernel(
|
|
const unsigned char* const* inputs,
|
|
float* output,
|
|
int batch_size,
|
|
int height,
|
|
int width,
|
|
float scale
|
|
) {
|
|
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
const int total_pixels = height * width;
|
|
|
|
if (idx < total_pixels * batch_size) {
|
|
const int batch_idx = idx / total_pixels;
|
|
const int pixel_idx = idx % total_pixels;
|
|
const int h = pixel_idx / width;
|
|
const int w = pixel_idx % width;
|
|
|
|
const unsigned char* input = inputs[batch_idx];
|
|
const int input_idx = (h * width + w) * 3;
|
|
const int output_idx = (batch_idx * total_pixels + h * width + w) * 3;
|
|
|
|
// BGR转RGB并归一化
|
|
output[output_idx + 0] = input[input_idx + 2] * scale; // R = B
|
|
output[output_idx + 1] = input[input_idx + 1] * scale; // G = G
|
|
output[output_idx + 2] = input[input_idx + 0] * scale; // B = R
|
|
}
|
|
}
|
|
|
|
// 预处理CUDA实现
|
|
bool preprocessGPU(
|
|
const std::vector<cv::Mat>& images,
|
|
detail::CudaBuffer& output_buffer,
|
|
const std::vector<std::unique_ptr<detail::CudaBuffer>>& temp_buffers,
|
|
cudaStream_t stream,
|
|
float scale
|
|
) {
|
|
const int batch_size = static_cast<int>(images.size());
|
|
const int height = images[0].rows;
|
|
const int width = images[0].cols;
|
|
const int total_pixels = height * width;
|
|
|
|
try {
|
|
// 1. 分配和拷贝输入数据
|
|
std::vector<unsigned char*> d_inputs(batch_size);
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
// 使用预分配的临时缓冲区
|
|
auto& buffer = temp_buffers[i];
|
|
const size_t image_size = images[i].total() * 3;
|
|
|
|
// 拷贝数据到设备
|
|
cudaMemcpyAsync(buffer->devicePtr(), images[i].data, image_size,
|
|
cudaMemcpyHostToDevice, stream);
|
|
d_inputs[i] = static_cast<unsigned char*>(buffer->devicePtr());
|
|
}
|
|
|
|
// 2. 分配和拷贝设备指针数组
|
|
unsigned char** d_input_ptrs;
|
|
cudaMalloc(&d_input_ptrs, batch_size * sizeof(unsigned char*));
|
|
cudaMemcpyAsync(d_input_ptrs, d_inputs.data(), batch_size * sizeof(unsigned char*),
|
|
cudaMemcpyHostToDevice, stream);
|
|
|
|
// 3. 启动kernel
|
|
const int grid_size = divUp(total_pixels * batch_size, BLOCK_SIZE);
|
|
batchPreprocessKernel<<<grid_size, BLOCK_SIZE, 0, stream>>>(
|
|
d_input_ptrs,
|
|
static_cast<float*>(output_buffer.devicePtr()),
|
|
batch_size,
|
|
height,
|
|
width,
|
|
scale
|
|
);
|
|
|
|
// 4. 清理
|
|
cudaFree(d_input_ptrs);
|
|
|
|
// 5. 检查错误
|
|
cudaError_t err = cudaGetLastError();
|
|
if (err != cudaSuccess) {
|
|
throw std::runtime_error(cudaGetErrorString(err));
|
|
}
|
|
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
pipeline::Logger::error("Error in GPU preprocessing: " + std::string(e.what()));
|
|
return false;
|
|
}
|
|
}
|
|
|
|
} // namespace detail
|
|
} // namespace pipeline |