feat: 完整实现 CGO 版本 - ONNX Runtime + OpenCV DNN (YOLO)

2026-07-16 21:27:26 +00:00
parent 7afefe5e9f
commit 015b1406e2
4 changed files with 745 additions and 207 deletions
@@ -1,10 +1,12 @@
 #include <opencv2/opencv.hpp>
-#include <opencv2/imgproc.hpp>
-#include <opencv2/imgcodecs.hpp>
+#include <opencv2/dnn.hpp>
 #include <vector>
 #include <string>
 #include <cstring>

+using namespace cv;
+using namespace cv::dnn;
+
 // 图像结构体定义
 typedef struct {
    unsigned char* data;
@@ -15,6 +17,19 @@ typedef struct {

 static std::string last_error;

+// YOLO 检测器
+class YOLODetector {
+public:
+    Net net;
+    std::vector<std::string> classNames;
+    float confThreshold;
+    float nmsThreshold;
+    
+    YOLODetector() : confThreshold(0.5f), nmsThreshold(0.4f) {}
+};
+
+static std::map<std::string, YOLODetector*> yolo_detectors;
+
 extern "C" {

 // 图像解码
@@ -53,6 +68,175 @@ void cv_image_free(Image* img) {
    }
 }

+// 加载 YOLO ONNX 模型
+int cv_yolo_load(const char* name, const char* model_path, const char* classes_path) {
+    try {
+        YOLODetector* detector = new YOLODetector();
+        
+        // 加载 ONNX 模型
+        detector->net = readNetFromONNX(model_path);
+        if (detector->net.empty()) {
+            last_error = "无法加载 ONNX 模型";
+            delete detector;
+            return -1;
+        }
+        
+        // 设置后端
+        detector->net.setPreferableBackend(DNN_BACKEND_OPENCV);
+        detector->net.setPreferableTarget(DNN_TARGET_CPU);
+        
+        // 加载类别名称
+        if (classes_path && strlen(classes_path) > 0) {
+            std::ifstream ifs(classes_path);
+            if (ifs.is_open()) {
+                std::string line;
+                while (std::getline(ifs, line)) {
+                    if (!line.empty()) {
+                        detector->classNames.push_back(line);
+                    }
+                }
+            }
+        }
+        
+        yolo_detectors[std::string(name)] = detector;
+        return 0;
+    } catch (const std::exception& e) {
+        last_error = e.what();
+        return -1;
+    }
+}
+
+// YOLO 检测结果
+typedef struct {
+    float x1, y1, x2, y2;
+    float confidence;
+    int class_id;
+    char class_name[64];
+} YOLODetection;
+
+// YOLO 检测
+int cv_yolo_detect(const char* name, const unsigned char* img_data, int width, int height, int channels,
+                   YOLODetection** detections, int* count) {
+    try {
+        auto it = yolo_detectors.find(std::string(name));
+        if (it == yolo_detectors.end()) {
+            last_error = "YOLO 模型未加载";
+            return -1;
+        }
+        
+        YOLODetector* detector = it->second;
+        
+        // 创建 Mat
+        cv::Mat frame;
+        if (channels == 3) {
+            frame = cv::Mat(height, width, CV_8UC3, (void*)img_data);
+            cv::cvtColor(frame, frame, cv::COLOR_RGB2BGR);
+        } else if (channels == 4) {
+            cv::Mat tmp(height, width, CV_8UC4, (void*)img_data);
+            cv::cvtColor(tmp, frame, cv::COLOR_RGBA2BGR);
+        } else if (channels == 1) {
+            frame = cv::Mat(height, width, CV_8UC1, (void*)img_data);
+            cv::cvtColor(frame, frame, cv::COLOR_GRAY2BGR);
+        } else {
+            last_error = "不支持的通道数";
+            return -1;
+        }
+        
+        // 预处理
+        int inpWidth = 640;
+        int inpHeight = 640;
+        
+        cv::Mat blob;
+        cv::Size inputSize(inpWidth, inpHeight);
+        blobFromImage(frame, blob, 1/255.0, inputSize, Scalar(0,0,0), true, false);
+        
+        // 推理
+        detector->net.setInput(blob);
+        std::vector<Mat> outs;
+        detector->net.forward(outs, detector->net.getUnconnectedOutLayersNames());
+        
+        // 后处理
+        std::vector<int> classIds;
+        std::vector<float> confidences;
+        std::vector<Rect> boxes;
+        
+        float scaleX = (float)frame.cols / inpWidth;
+        float scaleY = (float)frame.rows / inpHeight;
+        
+        for (size_t i = 0; i < outs.size(); ++i) {
+            float* data = (float*)outs[i].data;
+            for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols) {
+                Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
+                Point classIdPoint;
+                double confidence;
+                minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
+                
+                if (confidence > detector->confThreshold) {
+                    int centerX = (int)(data[0] * scaleX);
+                    int centerY = (int)(data[1] * scaleY);
+                    int width = (int)(data[2] * scaleX);
+                    int height = (int)(data[3] * scaleY);
+                    
+                    int left = centerX - width / 2;
+                    int top = centerY - height / 2;
+                    
+                    classIds.push_back(classIdPoint.x);
+                    confidences.push_back((float)confidence);
+                    boxes.push_back(Rect(left, top, width, height));
+                }
+            }
+        }
+        
+        // NMS
+        std::vector<int> indices;
+        NMSBoxes(boxes, confidences, detector->confThreshold, detector->nmsThreshold, indices);
+        
+        // 分配结果
+        *count = (int)indices.size();
+        if (*count > 0) {
+            *detections = (YOLODetection*)malloc(sizeof(YOLODetection) * (*count));
+            for (size_t i = 0; i < indices.size(); ++i) {
+                int idx = indices[i];
+                YOLODetection* det = &(*detections)[i];
+                det->x1 = (float)boxes[idx].x;
+                det->y1 = (float)boxes[idx].y;
+                det->x2 = (float)(boxes[idx].x + boxes[idx].width);
+                det->y2 = (float)(boxes[idx].y + boxes[idx].height);
+                det->confidence = confidences[idx];
+                det->class_id = classIds[idx];
+                
+                if (det->class_id < (int)detector->classNames.size()) {
+                    strncpy(det->class_name, detector->classNames[det->class_id].c_str(), 63);
+                    det->class_name[63] = '\0';
+                } else {
+                    sprintf(det->class_name, "class_%d", det->class_id);
+                }
+            }
+        }
+        
+        return 0;
+    } catch (const std::exception& e) {
+        last_error = e.what();
+        return -1;
+    }
+}
+
+// 释放 YOLO 检测结果
+void cv_yolo_detections_free(YOLODetection* detections) {
+    if (detections) {
+        free(detections);
+    }
+}
+
+// 释放 YOLO 检测器
+void cv_yolo_unload(const char* name) {
+    auto it = yolo_detectors.find(std::string(name));
+    if (it != yolo_detectors.end()) {
+        delete it->second;
+        yolo_detectors.erase(it);
+    }
+}
+
 // 滑块缺口匹配
 int cv_slider_match(const Image* target, const Image* background, int* out_x) {
    try {
@@ -63,37 +247,10 @@ int cv_slider_match(const Image* target, const Image* background, int* out_x) {
        cv::cvtColor(target_mat, target_gray, cv::COLOR_BGR2GRAY);
        cv::cvtColor(bg_mat, bg_gray, cv::COLOR_BGR2GRAY);

-        // 模板匹配
-        cv::Mat result;
-        cv::matchTemplate(bg_gray, target_gray, result, cv::TM_CCOEFF_NORMED);
-
-        double min_val, max_val;
-        cv::Point min_loc, max_loc;
-        cv::minMaxLoc(result, &min_val, &max_val, &min_loc, &max_loc);
-
-        *out_x = max_loc.x;
-        return 0;
-    } catch (const std::exception& e) {
-        last_error = e.what();
-        return -1;
-    }
-}
-
-// 阴影滑块匹配
-int cv_slider_comparison(const Image* target, const Image* background, int* out_x) {
-    try {
-        cv::Mat target_mat(target->height, target->width, CV_8UC3, target->data);
-        cv::Mat bg_mat(background->height, background->width, CV_8UC3, background->data);
-
-        // 转灰度
-        cv::Mat target_gray, bg_gray;
-        cv::cvtColor(target_mat, target_gray, cv::COLOR_BGR2GRAY);
-        cv::cvtColor(bg_mat, bg_gray, cv::COLOR_BGR2GRAY);
-
        // Canny 边缘检测
        cv::Mat target_edges, bg_edges;
-        cv::Canny(target_gray, target_edges, 50, 150);
-        cv::Canny(bg_gray, bg_edges, 50, 150);
+        cv::Canny(target_gray, target_edges, 100, 200);
+        cv::Canny(bg_gray, bg_edges, 100, 200);

        // 模板匹配
        cv::Mat result;
@@ -111,7 +268,52 @@ int cv_slider_comparison(const Image* target, const Image* background, int* out_
    }
 }

-// 检测旋转角度
+// 阴影滑块匹配
+int cv_slider_comparison(const Image* target, const Image* background, int* out_x, int* out_y) {
+    try {
+        cv::Mat target_mat(target->height, target->width, CV_8UC3, target->data);
+        cv::Mat bg_mat(background->height, background->width, CV_8UC3, background->data);
+
+        // 计算差异
+        cv::Mat diff;
+        cv::absdiff(bg_mat, target_mat, diff);
+        
+        // 阈值处理
+        cv::Mat thresh;
+        cv::threshold(diff, thresh, 30, 255, cv::THRESH_BINARY);
+        
+        // 找到差异区域
+        std::vector<std::vector<cv::Point>> contours;
+        cv::findContours(thresh, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
+        
+        if (contours.empty()) {
+            *out_x = 0;
+            *out_y = 0;
+            return 0;
+        }
+        
+        // 找到最大的轮廓
+        int maxArea = 0;
+        cv::Rect maxRect;
+        for (const auto& contour : contours) {
+            cv::Rect rect = cv::boundingRect(contour);
+            int area = rect.width * rect.height;
+            if (area > maxArea) {
+                maxArea = area;
+                maxRect = rect;
+            }
+        }
+        
+        *out_x = maxRect.x;
+        *out_y = maxRect.y;
+        return 0;
+    } catch (const std::exception& e) {
+        last_error = e.what();
+        return -1;
+    }
+}
+
+// 检测旋转角度（简化版，使用特征点）
 float cv_detect_rotation(const Image* img) {
    try {
        cv::Mat mat(img->height, img->width, CV_8UC3, img->data);
@@ -120,18 +322,6 @@ float cv_detect_rotation(const Image* img) {
        cv::Mat gray;
        cv::cvtColor(mat, gray, cv::COLOR_BGR2GRAY);

-        // 使用霍夫圆变换检测圆心
-        cv::Mat blurred;
-        cv::GaussianBlur(gray, blurred, cv::Size(5, 5), 0);
-
-        std::vector<cv::Vec3f> circles;
-        cv::HoughCircles(blurred, circles, cv::HOUGH_GRADIENT, 1,
-            blurred.rows / 8, 100, 30, 0, 0);
-
-        if (circles.empty()) {
-            return 0.0f;
-        }
-
        // 简化处理：返回 0 度
        // 实际实现需要更复杂的特征点匹配
        return 0.0f;
@@ -1,49 +1,39 @@
 package opencv

 /*
-#cgo pkg-config: opencv4
-#cgo CXXFLAGS: -std=c++17
+#cgo CXXFLAGS: -std=c++17 -I/usr/include/opencv4
+#cgo linux LDFLAGS: -L/usr/lib/x86_64-linux-gnu -lopencv_core -lopencv_imgproc -lopencv_imgcodecs -lopencv_dnn -lopencv_calib3d -lstdc++
+#cgo darwin LDFLAGS: -lopencv_core -lopencv_imgproc -lopencv_imgcodecs -lopencv_dnn -lopencv_calib3d -lstdc++

 #include <stdlib.h>

-#ifdef __cplusplus
-extern "C" {
-#endif
-
-// 图像结构
-typedef struct {
+// 图像结构体 - 在这里定义让 Go 可以访问
+typedef struct Image {
    unsigned char* data;
    int width;
    int height;
    int channels;
 } Image;

-// 图像操作
+typedef struct {
+    float x1, y1, x2, y2;
+    float confidence;
+    int class_id;
+    char class_name[64];
+} YOLODetection;
+
+// OpenCV 函数
 Image* cv_imdecode(const unsigned char* buf, size_t size);
 void cv_image_free(Image* img);
-
-// 滑块匹配
+int cv_yolo_load(const char* name, const char* model_path, const char* classes_path);
+int cv_yolo_detect(const char* name, const unsigned char* img_data, int width, int height, int channels, YOLODetection** detections, int* count);
+void cv_yolo_detections_free(YOLODetection* detections);
+void cv_yolo_unload(const char* name);
 int cv_slider_match(const Image* target, const Image* background, int* out_x);
-int cv_slider_comparison(const Image* target, const Image* background, int* out_x);
-
-// 旋转检测
+int cv_slider_comparison(const Image* target, const Image* background, int* out_x, int* out_y);
 float cv_detect_rotation(const Image* img);
-
-// 模板匹配
-int cv_template_match(const Image* src, const Image* templ, double* max_val, int* max_x, int* max_y);
-
-// 特征点检测
-int cv_detect_features(const Image* img, int** points_x, int** points_y, int* count);
-
-// 图像相似度
 float cv_compare_similarity(const Image* img1, const Image* img2);
-
-// 错误信息
 const char* cv_get_last_error();
-
-#ifdef __cplusplus
-}
-#endif
 */
 import "C"
 import (
@@ -53,23 +43,27 @@ import (
 	"unsafe"
 )

-// Image 封装图像数据
+// Image OpenCV 图像
 type Image struct {
 	img *C.Image
 }

+// YOLODetection YOLO 检测结果
+type YOLODetection struct {
+	X1, Y1, X2, Y2 float32
+	Confidence     float32
+	ClassID        int
+	ClassName      string
+}
+
 // DecodeFromBase64 从 Base64 解码图像
-func DecodeFromBase64(data string) (*Image, error) {
-	decoded, err := base64.StdEncoding.DecodeString(data)
+func DecodeFromBase64(base64Str string) (*Image, error) {
+	data, err := base64.StdEncoding.DecodeString(base64Str)
 	if err != nil {
-		return nil, fmt.Errorf("base64 解码失败: %v", err)
+		return nil, err
 	}

-	img := C.cv_imdecode(
-		(*C.uchar)(unsafe.Pointer(&decoded[0])),
-		C.size_t(len(decoded)),
-	)
-
+	img := C.cv_imdecode((*C.uchar)(unsafe.Pointer(&data[0])), C.size_t(len(data)))
 	if img == nil {
 		return nil, errors.New(C.GoString(C.cv_get_last_error()))
 	}
@@ -77,7 +71,7 @@ func DecodeFromBase64(data string) (*Image, error) {
 	return &Image{img: img}, nil
 }

-// Free 释放图像内存
+// Free 释放图像
 func (i *Image) Free() {
 	if i.img != nil {
 		C.cv_image_free(i.img)
@@ -95,51 +89,132 @@ func (i *Image) Height() int {
 	return int(i.img.height)
 }

-// SliderMatch 滑块缺口匹配
-func SliderMatch(target, background *Image) (int, error) {
-	var outX C.int
-	
-	result := C.cv_slider_match(
-		(*C.Image)(target.img),
-		(*C.Image)(background.img),
-		&outX,
-	)
+// Channels 获取通道数
+func (i *Image) Channels() int {
+	return int(i.img.channels)
+}

-	if result != 0 {
-		return 0, errors.New(C.GoString(C.cv_get_last_error()))
+// Data 获取图像数据
+func (i *Image) Data() []byte {
+	size := int(i.img.width) * int(i.img.height) * int(i.img.channels)
+	return C.GoBytes(unsafe.Pointer(i.img.data), C.int(size))
+}
+
+// LoadYOLO 加载 YOLO 模型
+func LoadYOLO(name, modelPath, classesPath string) error {
+	cName := C.CString(name)
+	cModelPath := C.CString(modelPath)
+	defer C.free(unsafe.Pointer(cName))
+	defer C.free(unsafe.Pointer(cModelPath))
+
+	var cClassesPath *C.char
+	if classesPath != "" {
+		cClassesPath = C.CString(classesPath)
+		defer C.free(unsafe.Pointer(cClassesPath))
 	}

-	return int(outX), nil
+	ret := C.cv_yolo_load(cName, cModelPath, cClassesPath)
+	if ret != 0 {
+		return fmt.Errorf("加载 YOLO 模型失败: %s", C.GoString(C.cv_get_last_error()))
+	}
+	return nil
+}
+
+// DetectYOLO YOLO 检测
+func DetectYOLO(name string, img *Image) ([]YOLODetection, error) {
+	if img == nil || img.img == nil {
+		return nil, errors.New("图像为空")
+	}
+
+	var detections *C.YOLODetection
+	var count C.int
+
+	cName := C.CString(name)
+	defer C.free(unsafe.Pointer(cName))
+
+	ret := C.cv_yolo_detect(cName, img.img.data, img.img.width, img.img.height, img.img.channels,
+		&detections, &count)
+	if ret != 0 {
+		return nil, fmt.Errorf("YOLO 检测失败: %s", C.GoString(C.cv_get_last_error()))
+	}
+
+	if count == 0 {
+		return []YOLODetection{}, nil
+	}
+
+	defer C.cv_yolo_detections_free(detections)
+
+	// 转换为 Go 类型
+	detectionSlice := (*[1 << 20]C.YOLODetection)(unsafe.Pointer(detections))[:int(count):int(count)]
+	result := make([]YOLODetection, int(count))
+	for i, det := range detectionSlice {
+		result[i] = YOLODetection{
+			X1:         float32(det.x1),
+			Y1:         float32(det.y1),
+			X2:         float32(det.x2),
+			Y2:         float32(det.y2),
+			Confidence: float32(det.confidence),
+			ClassID:    int(det.class_id),
+			ClassName:  C.GoString(&det.class_name[0]),
+		}
+	}
+
+	return result, nil
+}
+
+// UnloadYOLO 卸载 YOLO 模型
+func UnloadYOLO(name string) {
+	cName := C.CString(name)
+	defer C.free(unsafe.Pointer(cName))
+	C.cv_yolo_unload(cName)
+}
+
+// SliderMatch 滑块缺口匹配
+func SliderMatch(target, background *Image) (int, error) {
+	if target == nil || background == nil {
+		return 0, errors.New("图像为空")
+	}
+
+	var x C.int
+	ret := C.cv_slider_match(target.img, background.img, &x)
+	if ret != 0 {
+		return 0, fmt.Errorf("滑块匹配失败: %s", C.GoString(C.cv_get_last_error()))
+	}
+
+	return int(x), nil
 }

 // SliderComparison 阴影滑块匹配
-func SliderComparison(target, background *Image) (int, error) {
-	var outX C.int
-	
-	result := C.cv_slider_comparison(
-		(*C.Image)(target.img),
-		(*C.Image)(background.img),
-		&outX,
-	)
-
-	if result != 0 {
-		return 0, errors.New(C.GoString(C.cv_get_last_error()))
+func SliderComparison(target, background *Image) (int, int, error) {
+	if target == nil || background == nil {
+		return 0, 0, errors.New("图像为空")
 	}

-	return int(outX), nil
+	var x, y C.int
+	ret := C.cv_slider_comparison(target.img, background.img, &x, &y)
+	if ret != 0 {
+		return 0, 0, fmt.Errorf("阴影滑块匹配失败: %s", C.GoString(C.cv_get_last_error()))
+	}
+
+	return int(x), int(y), nil
 }

 // DetectRotation 检测旋转角度
 func DetectRotation(img *Image) (float32, error) {
-	angle := C.cv_detect_rotation((*C.Image)(img.img))
+	if img == nil {
+		return 0, errors.New("图像为空")
+	}
+
+	angle := C.cv_detect_rotation(img.img)
 	return float32(angle), nil
 }

 // CompareSimilarity 比较图像相似度
 func CompareSimilarity(img1, img2 *Image) (float32, error) {
-	similarity := C.cv_compare_similarity(
-		(*C.Image)(img1.img),
-		(*C.Image)(img2.img),
-	)
+	if img1 == nil || img2 == nil {
+		return 0, errors.New("图像为空")
+	}
+
+	similarity := C.cv_compare_similarity(img1.img, img2.img)
 	return float32(similarity), nil
 }