Untitled
unknown
plain_text
a year ago
2.2 kB
2
Indexable
Never
void ObjectDetector::detect(cv::Mat& image, bool showConf) { cv::Mat copy = image; // Convert the image to grayscale cv::cvtColor(copy, copy, cv::COLOR_BGR2GRAY); // Enhance the image contrast cv::equalizeHist(copy, copy); // Resize the image to increase speed cv::resize(copy, copy, cv::Size(224, 224), 6); // Convert resized image to BGR for compatibility cv::cvtColor(copy, copy, cv::COLOR_BGRA2BGR); cv::Mat blob; cv::dnn::blobFromImage(copy, blob, 1. / 255., cv::Size(224, 224)); model.setInput(blob); std::vector<cv::Mat> outputs; model.forward(outputs); float* data = (float*)outputs[0].data; const int rows = outputs[0].size[1]; float scaleX = image.cols / (float)copy.cols; float scaleY = image.rows / (float)copy.rows; std::vector<cv::Rect> boxes; std::vector<float> confidences; std::vector<int> classes; for (int i = 0; i < rows; i++) { float confidence = data[4]; if (confidence > minConfidence) { int classId = 0; float maxConfidence = 0; for (int j = 0; j < classNames.size(); j++) if (data[5 + j] > maxConfidence) { classId = j; maxConfidence = data[5 + j]; } if (classNames[classId].second) { int box_width = int(data[2] * scaleX); int box_height = int(data[3] * scaleY); int box_x = int(data[0] * scaleX - 0.5 * box_width); int box_y = int(data[1] * scaleY - 0.5 * box_height); lastRect = cv::Rect(box_x, box_y, box_width, box_height); boxes.push_back(lastRect); confidences.push_back(confidence); classes.push_back(classId); } } data += 5 + classNames.size(); } float nmsThreshold = 0.1; std::vector<int> indices; cv::dnn::NMSBoxes(boxes, confidences, minConfidence, nmsThreshold, indices); for (const auto &index : indices) { cv::Rect box = boxes[index]; float confidence = confidences[index]; int classId = classes[index]; currentClassName = classNames[classId].first.c_str(); std::stringstream ss; ss << currentClassName; if (showConf) ss << ": confidence = " + std::to_string((int)(confidence * 100)) + "%"; drawLabel(image, ss.str(), lastRect); } cv::cvtColor(image, image, cv::COLOR_BGRA2BGR); }