#include"belt.hpp"
|
|
#define checkRuntime(op) __check_cuda_runtimeBelt((op), #op, __FILE__, __LINE__)
|
|
bool __check_cuda_runtimeBelt(cudaError_t code, const char* op, const char* file, int line) {
|
if (code != cudaSuccess) {
|
const char* err_name = cudaGetErrorName(code);
|
const char* err_message = cudaGetErrorString(code);
|
printf("runtime error %s:%d %s failed. \n code = %s, message = %s\n", file, line, op, err_name, err_message);
|
return false;
|
}
|
return true;
|
}
|
|
|
inline const char* severity_string(nvinfer1::ILogger::Severity t) {
|
switch (t) {
|
case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
|
case nvinfer1::ILogger::Severity::kERROR: return "error";
|
case nvinfer1::ILogger::Severity::kWARNING: return "warning";
|
case nvinfer1::ILogger::Severity::kINFO: return "info";
|
case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
|
default: return "unknow";
|
}
|
}
|
|
|
class Logger : public ILogger
|
{
|
void log(Severity severity, const char* msg) noexcept
|
{
|
// suppress info-level messages
|
if (severity != Severity::kINFO)
|
printf("%s: %s\n", severity_string(severity), msg);
|
}
|
};
|
|
|
void warp_affine_bilinear(uint8_t* src, int src_line_size, int src_width, int src_height,
|
float* dst, int dst_width, int dst_height, uint8_t fill_value, AffineMatrix matrix, cudaStream_t stream);
|
|
|
BELT::~BELT() {
|
|
checkRuntime(cudaStreamSynchronize(stream));
|
if (stream != nullptr)
|
{
|
checkRuntime(cudaStreamDestroy(stream));
|
}
|
|
if (!context) {
|
context->destroy();
|
}
|
if (!engine) {
|
engine->destroy();
|
}
|
if (!runtime) {
|
runtime->destroy();
|
}
|
|
if (!output_data_host) {
|
delete[] output_data_host;
|
}
|
|
if (!input_data_host) {
|
delete[] input_data_host;
|
}
|
|
if (!input_data_device) {
|
checkRuntime(cudaFree(input_data_device));
|
}
|
|
|
}
|
|
bool BELT::initConfig(const char* enginefile, int* leftarea, int* rightarea, int* beltarea, float bigcoal_threshold, float beltdeviation_threshold, float confThreshold) {
|
|
std::ifstream file(enginefile, std::ios::binary);
|
char* trtModelStream = NULL;
|
int size = 0;
|
if (file.good()) {
|
file.seekg(0, file.end);
|
size = file.tellg();
|
file.seekg(0, file.beg);
|
trtModelStream = new char[size];
|
assert(trtModelStream);
|
file.read(trtModelStream, size);
|
file.close();
|
}
|
else {
|
return false;
|
}
|
|
Logger logger;
|
this->runtime = createInferRuntime(logger);
|
assert(this->runtime != nullptr);
|
this->engine = runtime->deserializeCudaEngine(trtModelStream, size);
|
assert(this->engine != nullptr);
|
this->context = engine->createExecutionContext();
|
assert(this->context != nullptr);
|
delete[] trtModelStream;
|
|
auto input_dims = engine->getBindingDimensions(0);
|
input_batch = input_dims.d[0];
|
input_channel = input_dims.d[1];
|
input_height = input_dims.d[2];
|
input_width = input_dims.d[3];
|
input_numel = input_batch * input_channel * input_height * input_width;
|
|
input_data_host = new float[input_numel * sizeof(float)] { 0 };
|
|
auto output_dims = engine->getBindingDimensions(1);
|
output_numbox = output_dims.d[1];
|
output_numprob = output_dims.d[2];
|
num_classes = output_numprob - 5;
|
output_numel = input_batch * output_numbox * output_numprob;
|
|
output_data_host = new float[sizeof(float) * output_numel] { 0 };
|
|
checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));
|
checkRuntime(cudaMalloc(&output_data_device, output_numel * sizeof(float)));
|
|
checkRuntime(cudaStreamCreate(&stream));
|
|
|
point left_p, right_p, belt_p;
|
for (int i = 0; i < 8; i += 2) {
|
|
left_p.x = leftarea[i];
|
left_p.y = leftarea[i + 1];
|
leftarea_.push_back(left_p);
|
}
|
|
for (int j = 0; j < 8; j += 2) {
|
|
right_p.x = rightarea[j];
|
right_p.y = rightarea[j + 1];
|
rightarea_.push_back(right_p);
|
}
|
|
for (int q = 0; q < 8; q += 2) {
|
|
belt_p.x = beltarea[q];
|
belt_p.y = beltarea[q + 1];
|
beltarea_.push_back(belt_p);
|
}
|
|
//Ƥ������
|
belt_top_cy = (beltarea[1] + beltarea[3]) / 2;
|
belt_bottom_cy = (beltarea[5] + beltarea[7]) / 2;
|
|
belt_top_length = abs(beltarea[2] - beltarea[0]);
|
belt_bottom_length = abs(beltarea[4] - beltarea[6]);
|
|
k = (1 / belt_top_length - 1 / belt_bottom_length) / (belt_top_cy - belt_bottom_cy);
|
if (abs(k) < 0.000001) {
|
k += (abs(k) / k * 0.000001);
|
}
|
|
confThreshold_ = confThreshold;
|
threshold = bigcoal_threshold;
|
threshold_ = beltdeviation_threshold;
|
|
return true;
|
|
}
|
|
int BELT::detect(cv::Mat& image, std::vector<Box>& boxes) {
|
|
int width = image.cols;
|
int height = image.rows;
|
int channels = image.channels();
|
int src_size = width * height * channels;
|
uint8_t* psrc_device = nullptr;
|
checkRuntime(cudaMalloc(&psrc_device, src_size));
|
checkRuntime(cudaMemcpyAsync(psrc_device, image.data, src_size, cudaMemcpyHostToDevice, stream));
|
|
AffineMatrix affine;
|
affine.compute(width, height, input_width, input_height);
|
warp_affine_bilinear(psrc_device, width * 3, width, height, input_data_device, input_width, input_height, 114, affine, stream);
|
|
float* bindings[] = { input_data_device,output_data_device };
|
context->enqueueV2((void**)bindings, stream, nullptr);
|
|
checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream));
|
checkRuntime(cudaFree(psrc_device));
|
|
std::vector<cv::Rect> bboxes;
|
std::vector<int> classIds;
|
std::vector<float> scores;
|
for (int i = 0; i < output_numbox; ++i) {
|
float* ptr = output_data_host + i * output_numprob;
|
float objness = ptr[4];
|
if (objness < confThreshold_) {
|
continue;
|
}
|
float* pclass = ptr + 5;
|
int label = std::max_element(pclass, pclass + num_classes) - pclass;
|
float prob = pclass[label] * objness;
|
if (prob < confThreshold_)
|
continue;
|
|
float cx = ptr[0];
|
float cy = ptr[1];
|
float w1 = ptr[2];
|
float h1 = ptr[3];
|
|
cv::Rect box;
|
box.x = cx;
|
box.y = cy;
|
box.width = w1;
|
box.height = h1;
|
|
bboxes.emplace_back(box);
|
classIds.emplace_back(label);
|
scores.emplace_back(prob);
|
|
}
|
|
|
if (bboxes.size() == 0) {
|
|
return 1;
|
}
|
|
|
point box_p;
|
float left_roller_num = 0, right_roller_num = 0;//���������
|
float left_roller_length = 0, right_roller_length = 0;//�����й���������ռƤ��������ȵı���
|
struct Box box;
|
|
std::vector<int> indexes;
|
cv::dnn::NMSBoxes(bboxes, scores, confThreshold_, nmsThreshold_, indexes);
|
for (size_t i = 0; i < indexes.size(); i++) {
|
int idx = indexes[i];
|
auto& ibox = bboxes[idx];
|
float x = ibox.x;
|
float y = ibox.y;
|
float w = ibox.width;
|
float h = ibox.height;
|
int id = classIds[idx];
|
|
int image_base_cx = static_cast<int>(affine.d2i[0] * x + affine.d2i[2]);
|
int image_base_cy = static_cast<int>(affine.d2i[0] * y + affine.d2i[5]);
|
int image_base_width = static_cast<int>(affine.d2i[0] * w);
|
int image_base_height = static_cast<int>(affine.d2i[0] * h);
|
if ((image_base_width <= 0) || (image_base_height <= 0))
|
continue;
|
|
box_p.x = image_base_cx;
|
box_p.y = image_base_cy;
|
|
float mmpx = k * (image_base_cy - belt_top_cy) + 1 / belt_top_length;
|
|
//�������й��Ҳ���Ƥ������
|
if ((id != 1) && (!panduan(beltarea_, box_p))) {
|
continue;
|
}
|
|
//�ж�λ��Ƥ�������ú���Ƿ��ú��
|
float width_ratio = mmpx * image_base_width;
|
float height_ratio = mmpx * image_base_height;
|
float diameter = std::max(width_ratio, height_ratio);
|
if (id == 0) {
|
if (diameter < threshold) {
|
continue;
|
}
|
|
}
|
|
if (id == 1) {
|
//λ����������й�
|
if (panduan(leftarea_, box_p)) {
|
left_roller_num += 1;
|
left_roller_length += (mmpx * image_base_width);
|
|
}
|
//λ����������й�
|
else if (panduan(rightarea_, box_p)) {
|
right_roller_num += 1;
|
right_roller_length += (mmpx * image_base_width);
|
}
|
|
else {
|
continue;
|
}
|
}
|
|
box.class_id = id;
|
box.x = image_base_cx;
|
box.y = image_base_cy;
|
box.width = image_base_width;
|
box.height = image_base_height;
|
box.score = scores[idx];
|
boxes.push_back(box);
|
}
|
|
if (left_roller_num == 0) {
|
left_roller_num = 1;
|
}
|
|
if (right_roller_num == 0) {
|
right_roller_num = 1;
|
}
|
|
//printf("�������=%f,��=%f;�������=%f,��=%f\n", left_roller_num,left_roller_length / left_roller_num, right_roller_num,right_roller_length / right_roller_num);
|
|
if (abs(left_roller_length / left_roller_num - right_roller_length / right_roller_num) > threshold_) {
|
|
return 0;//��ƫ
|
}
|
|
else {
|
return 1;//����ƫ
|
}
|
}
|
