AntiCaptcha/pkg/opencv/opencv.go

package opencv

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

#include <stdlib.h>

// 图像结构体 - 在这里定义让 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* out_y);
float cv_detect_rotation(const Image* img);
float cv_compare_similarity(const Image* img1, const Image* img2);
const char* cv_get_last_error();
*/
import "C"
import (
	"encoding/base64"
	"errors"
	"fmt"
	"unsafe"
)

// 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(base64Str string) (*Image, error) {
	data, err := base64.StdEncoding.DecodeString(base64Str)
	if err != nil {
		return nil, err
	}

	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()))
	}

	return &Image{img: img}, nil
}

// Free 释放图像
func (i *Image) Free() {
	if i.img != nil {
		C.cv_image_free(i.img)
		i.img = nil
	}
}

// Width 获取宽度
func (i *Image) Width() int {
	return int(i.img.width)
}

// Height 获取高度
func (i *Image) Height() int {
	return int(i.img.height)
}

// Channels 获取通道数
func (i *Image) Channels() int {
	return int(i.img.channels)
}

// 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))
	}

	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, int, error) {
	if target == nil || background == nil {
		return 0, 0, errors.New("图像为空")
	}

	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) {
	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) {
	if img1 == nil || img2 == nil {
		return 0, errors.New("图像为空")
	}

	similarity := C.cv_compare_similarity(img1.img, img2.img)
	return float32(similarity), nil
}