zhipuzi_pos_windows/zhipuzi_pos_windows/ai/YoloFeatureManager.cpp

#include "../pch/pch.h"

#include "YoloFeatureManager.h"
#include <fstream>
#include <algorithm>
#include <iostream>
#include <functional>
#include <numeric>
#include <sstream>

#include "../tool/CVideoCapture.h"

#include "YoloClassName.h"

YoloFeatureManager::YoloFeatureManager()
{
	inputWidth = 448;
	inputHeight = 448	;

	CONF_THRESHOLD = 0.5f; // 可以根据需要调整置信度阈值
	NMS_THRESHOLD = 0.4f;  // 可以根据需要调整NMS阈值

	FRUIT_VEGETABLE_COUNT = sizeof(FRUIT_VEGETABLE_NAMES) / sizeof(FRUIT_VEGETABLE_NAMES[0]);
}

YoloFeatureManager::~YoloFeatureManager()
{
	
}

void YoloFeatureManager::loadModel()
{
	try
	{
		std::wstring wsProgramDir = CSystem::GetProgramDir();
		std::filesystem::path mainDir = wsProgramDir;
		std::string sMainDir = mainDir.string();

		// YOLO2026模型路径
		std::string modelPath = sMainDir + "/ai/yolo26s-cls-zhipuzi-448.onnx";

		net = cv::dnn::readNetFromONNX(modelPath);
	}
	catch (const std::exception& e)
	{
		std::string aa = std::string(e.what());
		CLewaimaiLog::OutputDebugMessage("加载模型失败: " + std::string(e.what()));
		return;
	}
}

void YoloFeatureManager::loadModelForOpenVINO()
{
	try
	{
		std::wstring wsProgramDir = CSystem::GetProgramDir();
		std::filesystem::path mainDir = wsProgramDir;
		std::string sMainDir = mainDir.string();

		// YOLO2026模型路径
		std::string modelPath = sMainDir + "/ai/best.xml";
		std::string configPath = sMainDir + "/ai/best.bin";

		net = cv::dnn::readNetFromModelOptimizer(modelPath, configPath);

		// 设置目标设备 (可选: CPU, GPU, MYRIAD等)
		net.setPreferableBackend(cv::dnn::DNN_BACKEND_INFERENCE_ENGINE);
		net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);  // 或DNN_TARGET_MYRIAD等
	}
	catch (const std::exception& e)
	{
		std::string aa = std::string(e.what());
		CLewaimaiLog::OutputDebugMessage(("加载模型失败: " + std::string(e.what())).c_str());
		return;
	}
}

// 寻找置信度最高的类别
int YoloFeatureManager::getTopClass(const cv::Mat& output)
{
	// 将输出展平为一维数组
	cv::Mat flatOutput = output.reshape(1, 1);

	double maxVal;

	cv::Point maxLoc;

	// 找到最大值的位置（即最高置信度类别索引）
	cv::minMaxLoc(flatOutput, nullptr, &maxVal, nullptr, &maxLoc);

	return maxLoc.x;
}

//	获取类别名称
std::string YoloFeatureManager::getClassName(std::size_t classId) const
{
	if (classId >= 0 && classId < FRUIT_VEGETABLE_COUNT)
	{
		std::string cnName = FRUIT_VEGETABLE_NAMES[classId];

		return cnName;
	}

	return "Unknown";
}

std::vector<float> YoloFeatureManager::extractFeatures(const std::string & imagePath)
{
	cv::Mat image = cv::imread(imagePath);

	return extractFeatures(image);
}

std::vector<float> YoloFeatureManager::extractFeatures(cv::Mat& image)
{
	std::lock_guard<std::mutex> lock(m_mutex);

	try
	{
		auto time_1 = std::chrono::high_resolution_clock::now();

		if (image.empty())
		{
			throw std::runtime_error("Could not load image");
		}

		// 转换为blob（归一化+通道转换）
		cv::Mat blob;
		cv::dnn::blobFromImage(image, blob, 1.0 / 255, cv::Size(inputWidth, inputHeight), cv::Scalar(0, 0, 0), true, true);
		net.setInput(blob);

		auto time_2 = std::chrono::high_resolution_clock::now();

		//获取模型的所有层名称（调试用）
		std::vector<cv::String> layerNames = net.getLayerNames();

		// 获取Flatten层输出（yolo26s-cls的Flatten层名称为 "onnx_node!/model.10/Flatten"，这是GAP后分类头前的一层）'
		// GAP层是onnx_node!/model.10/pool/GlobalAveragePool
		cv::Mat featureMat = net.forward("onnx_node!/model.10/Flatten");

		// 检查输出是否有效
		if (featureMat.empty())
		{
			throw std::runtime_error("模型前向传播未产生有效输出");
		}

		if (featureMat.type() != CV_32F)
		{
			throw std::runtime_error("Mat类型错误");
		}

		float norm_before = cv::norm(featureMat, cv::NORM_L2);
		CLewaimaiLog::OutputDebugMessageFormat("归一化前 norm：%.6f\n", norm_before);

		cv::normalize(featureMat, featureMat, 1.0, 0.0, cv::NORM_L2); //L2归一化

		float norm_after = cv::norm(featureMat, cv::NORM_L2);
		CLewaimaiLog::OutputDebugMessageFormat("归一化后 norm：%.6f\n", norm_after);

		// 将Mat格式的特征转换为vector<float>（方便后续计算/存储）
		std::vector<float> feature_vector;
		feature_vector.assign((float*)featureMat.data, (float*)featureMat.data + featureMat.total());

		//进行时间统计
		auto time_3 = std::chrono::high_resolution_clock::now();

		auto duration_1 = std::chrono::duration_cast<std::chrono::milliseconds>(time_2 - time_1);
		std::wstring msg = L"图片处理耗时: " + std::to_wstring(duration_1.count()) + L" 毫秒";
		CLewaimaiLog::OutputDebugMessage(msg.c_str());

		auto duration_2 = std::chrono::duration_cast<std::chrono::milliseconds>(time_3 - time_2);
		std::wstring msg2 = L"模型推理耗时: " + std::to_wstring(duration_2.count()) + L" 毫秒";
		CLewaimaiLog::OutputDebugMessage(msg2.c_str());

		auto totalDuration = std::chrono::duration_cast<std::chrono::milliseconds>(time_3 - time_1);
		std::wstring msg4 = L"总耗时: " + std::to_wstring(totalDuration.count()) + L" 毫秒";
		CLewaimaiLog::OutputDebugMessage(msg4.c_str());

		return feature_vector;
	}
	catch (const std::exception& e)
	{
		std::string aa = std::string(e.what());
		CLewaimaiLog::OutputDebugMessage("提取特征失败: " + std::string(e.what()));
		return {};
	}
}

void YoloFeatureManager::DebugTopResults(const cv::Mat& output, int topK)
{
	cv::Mat scores = output.reshape(1, 1);
	const int count = scores.cols;
	topK = std::min(topK, count);

	std::vector<int> indices(count);
	std::iota(indices.begin(), indices.end(), 0);

	std::partial_sort(indices.begin(), indices.begin() + topK, indices.end(),
		[&scores](int left, int right)
		{
			return scores.at<float>(0, left) > scores.at<float>(0, right);
		});

	std::string message = "YOLO分类 Top " + std::to_string(topK) + " 结果:";
	for (int i = 0; i < topK; ++i)
	{
		const int classId = indices[i];
		const float confidence = scores.at<float>(0, classId);

		message += "\nTop " + std::to_string(i + 1) +
			": id=" + std::to_string(classId) +
			": name=" + this->getClassName(classId) +
			", confidence=" + std::to_string(confidence);
	}

	m_topResultMessage = message;
}

std::string YoloFeatureManager::Class(cv::Mat & image)
{
	std::lock_guard<std::mutex> lock(m_mutex);

	try
	{ 
		std::string className = "";

		// ====================== 图像预处理 ======================
		// 转换为blob格式：归一化(0-1)、通道转换(BGR->RGB)、调整尺寸
		cv::Mat blob;
		cv::dnn::blobFromImage(image, blob, 1.0 / 255.0, cv::Size(inputWidth, inputHeight), cv::Scalar(0, 0, 0), true, true);
		net.setInput(blob);

		auto time_1 = std::chrono::high_resolution_clock::now();

		// ====================== 模型推理 ======================
		cv::Mat output = net.forward(); // 输出形状：1x1000（对应ImageNet 1000类）

		this->DebugTopResults(output);

		auto time_2 = std::chrono::high_resolution_clock::now();
		auto totalDuration = std::chrono::duration_cast<std::chrono::milliseconds>(time_2 - time_1);

		std::wstring msg = L"Class 总耗时: " + std::to_wstring(totalDuration.count()) + L" 毫秒\r\n";
		msg += CLewaimaiString::ANSIToUnicode(m_topResultMessage);

		CLewaimaiLog::OutputDebugMessage(msg.c_str());
		CLewaimaiLog::OutputMessageBox(msg.c_str());

		// ====================== 解析结果 ======================
		int topClassIdx = this->getTopClass(output);
		float topConfidence = output.at<float>(topClassIdx);

		if (topConfidence > 0.8)
		{
			className = this->getClassName(topClassIdx);
		}
		else
		{
			className = "Unknown";
		}

		return className;
	}
	catch (const std::exception& e)
	{
		CLewaimaiLog::OutputDebugMessage(("YOLO分类失败: " + std::string(e.what())).c_str());
		return {};
	}
	
}

void YoloFeatureManager::drawChineseText(cv::Mat & img, const wchar_t * text, cv::Point pos, cv::Scalar color, int fontSize)
{
	// 1. 检查输入有效性
	if (img.empty() || text == nullptr || wcslen(text) == 0)
	{
		return;
	}
	if (img.type() != CV_8UC3)
	{ 
		// 仅支持 3 通道彩色图像
		cvtColor(img, img, cv::COLOR_GRAY2BGR);
	}

	// 2. 创建内存 DC 并关联临时位图（关键：基于图像的 DC 创建，而非屏幕 DC）
	HDC hScreenDC = GetDC(NULL);
	HDC hMemDC = CreateCompatibleDC(hScreenDC);
	// 创建与原图像尺寸、格式匹配的位图
	HBITMAP hMemBmp = CreateCompatibleBitmap(hScreenDC, img.cols, img.rows);
	// 保存原始位图句柄，用于后续恢复
	HBITMAP hOldBmp = (HBITMAP)SelectObject(hMemDC, hMemBmp);

	// 3. 将 OpenCV 图像数据复制到内存位图（保留原图像内容，而非黑色）
	BITMAPINFO bmi = { 0 };
	bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
	bmi.bmiHeader.biWidth = img.cols;
	bmi.bmiHeader.biHeight = -img.rows; // 翻转 Y 轴（OpenCV 与 GDI 坐标方向相反）
	bmi.bmiHeader.biPlanes = 1;
	bmi.bmiHeader.biBitCount = 24;
	bmi.bmiHeader.biCompression = BI_RGB;
	// 将 OpenCV 图像写入内存位图
	SetDIBits(hScreenDC, hMemBmp, 0, img.rows, img.data, &bmi, DIB_RGB_COLORS);

	// 4. 设置中文字体（修复字体创建参数，增加容错）
	HFONT hFont = CreateFont(
		fontSize, 0, 0, 0, FW_NORMAL, 0, 0, 0,
		GB2312_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS,
		DEFAULT_QUALITY, DEFAULT_PITCH | FF_DONTCARE, L"黑体"
	);
	HFONT hOldFont = (HFONT)SelectObject(hMemDC, hFont);

	// 5. 设置文字绘制属性（背景透明、颜色正确）
	SetBkMode(hMemDC, TRANSPARENT);
	// OpenCV 是 BGR，GDI 是 RGB，需转换
	SetTextColor(hMemDC, RGB((int)color[2], (int)color[1], (int)color[0]));

	// 6. 绘制中文字符（确保坐标在图像范围内）
	int textLen = wcslen(text);
	if (pos.x >= 0 && pos.y >= 0 && pos.x < img.cols && pos.y < img.rows)
	{
		TextOutW(hMemDC, pos.x, pos.y, text, textLen);
	}

	// 7. 将绘制后的位图数据复制回 OpenCV 图像
	GetDIBits(hScreenDC, hMemBmp, 0, img.rows, img.data, &bmi, DIB_RGB_COLORS);

	// 8. 释放资源（关键：恢复原始句柄后再删除，避免内存泄漏）
	SelectObject(hMemDC, hOldFont); // 恢复原始字体
	DeleteObject(hFont);            // 删除自定义字体
	SelectObject(hMemDC, hOldBmp);  // 恢复原始位图
	DeleteObject(hMemBmp);          // 删除内存位图
	DeleteDC(hMemDC);               // 删除内存 DC
	ReleaseDC(NULL, hScreenDC);     // 释放屏幕 DC
}