feat: 完整实现 CGO 版本 - ONNX Runtime + OpenCV DNN (YOLO)
This commit is contained in:
+142
-67
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user