Untitled

mail@pastecode.io avatar
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);
}