插件：AnimeTrace 动画/Galgame识别

Author: yexiaoyu123

README.md

@@ -392,6 +392,18 @@ print("run[CQ:image,file="+j["img"]+"]")
 
   - [x] waifu | 随机waifu(从[100000个AI生成的waifu](https://www.thiswaifudoesnotexist.net/)中随机一位)
 
+</details>
+<details>
+  <summary>AnimeTrace 动画/Galgame识别</summary>
+
+  `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/animetrace"`
+
+  基于[AnimeTrace](https://ai.animedb.cn/)API  的识图搜索插件
+  
+  - [x] Gal识图 | Gal识图 [模型名]
+
+  - [x] 动漫识图 | 动漫识图 2 | 动漫识图 [模型名]
+
 </details>
 <details>
   <summary>支付宝到账语音</summary>

main.go

@@ -67,6 +67,7 @@ import (
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aifalse"          // 服务器监控
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiwife"           // 随机老婆
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/alipayvoice"      // 支付宝到账语音
+	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/animetrace"       // AnimeTrace 动画/Galgame识别
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/autowithdraw"     // 触发者撤回时也自动撤回
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/baiduaudit"       // 百度内容审核
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/base16384"        // base16384加解密

plugin/animetrace/main.go

@@ -0,0 +1,205 @@
+package animetrace
+
+import (
+	"bytes"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"image"
+	"image/jpeg"
+	"image/png"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"strings"
+	"time"
+
+	ctrl "github.com/FloatTech/zbpctrl"
+	"github.com/FloatTech/zbputils/control"
+	"github.com/FloatTech/zbputils/ctxext"
+	"github.com/tidwall/gjson"
+	zero "github.com/wdvxdr1123/ZeroBot"
+	"github.com/wdvxdr1123/ZeroBot/message"
+)
+
+var (
+	model string
+)
+
+func init() {
+	engine := control.Register("animetrace", &ctrl.Options[*zero.Ctx]{
+		DisableOnDefault: false,
+		Brief:            "AnimeTrace 动画/Galgame识别",
+		Help:             "- Gal识图\n- 动漫识图\n- 动漫识图 2\n- 动漫识图 [模型名]\n- Gal识图 [模型名]",
+		PublicDataFolder: "Animetrace",
+		OnEnable:         func(ctx *zero.Ctx) { ctx.Send("插件已启用") },
+		OnDisable:        func(ctx *zero.Ctx) { ctx.Send("插件已禁用") },
+	})
+
+	engine.OnPrefix("gal识图", zero.OnlyGroup, zero.MustProvidePicture).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+		args := ctx.State["args"].(string)
+		if strings.TrimSpace(args) == "" {
+			model = "full_game_model_kira" // 默认使用的模型
+		} else {
+			model = args // 自定义设置模型
+		}
+		processImageRecognition(ctx, model)
+	})
+
+	engine.OnPrefix("动漫识图", zero.OnlyGroup, zero.MustProvidePicture).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+		args := ctx.State["args"].(string)
+		if strings.TrimSpace(args) == "" {
+			model = "anime_model_lovelive"
+		} else if strings.TrimSpace(args) == "2" {
+			model = "pre_stable"
+		} else {
+			model = args
+		}
+		processImageRecognition(ctx, model)
+	})
+}
+
+// 处理图片识别
+func processImageRecognition(ctx *zero.Ctx, model string) {
+	urls := ctx.State["image_url"].([]string)
+	if len(urls) == 0 {
+		return
+	}
+
+	imageData, err := downloadImage(urls[0])
+	if err != nil {
+		ctx.Send(message.Text("下载图片失败: ", err))
+		return
+	}
+	//ctx.Send(message.Text(model))
+	respBody, err := createAndSendMultipartRequest("https://api.animetrace.com/v1/search", imageData, map[string]string{
+		"is_multi":  "0",
+		"model":     model,
+		"ai_detect": "0",
+	})
+	if err != nil {
+		ctx.Send(message.Text("识别请求失败: ", err))
+		return
+	}
+
+	code := gjson.Get(string(respBody), "code").Int()
+	if code != 0 {
+		ctx.Send(message.Text("错误: ", gjson.Get(string(respBody), "zh_message").String()))
+		return
+	}
+
+	dataArray := gjson.Get(string(respBody), "data").Array()
+	countStr := fmt.Sprintf("%d", len(dataArray))
+	if countStr == "0" {
+		ctx.Send(message.Text("未识别到任何角色"))
+		return
+	}
+	var sk message.Message
+	sk = append(sk, ctxext.FakeSenderForwardNode(ctx, message.Text("共识别到 "+countStr+" 个角色，可能是以下来源")))
+
+	for _, value := range dataArray {
+		boxArray := value.Get("box").Array()
+		box := make([]float64, len(boxArray))
+		for i, val := range boxArray {
+			box[i] = val.Float()
+		}
+		base64Str, err := cropImage(imageData, box)
+		if err != nil {
+			ctx.Send(message.Text("图片处理失败: ", err))
+			return
+		}
+		var sb strings.Builder
+		value.Get("character").ForEach(func(_, character gjson.Result) bool {
+			sb.WriteString(fmt.Sprintf("《%s》的角色 %s\n", character.Get("work").String(), character.Get("character").String()))
+			return true
+		})
+
+		sk = append(sk, ctxext.FakeSenderForwardNode(ctx, message.Image("base64://"+base64Str), message.Text(sb.String())))
+	}
+
+	ctx.SendGroupForwardMessage(ctx.Event.GroupID, sk)
+}
+
+// 下载图片，设置超时
+func downloadImage(url string) ([]byte, error) {
+	client := &http.Client{Timeout: 10 * time.Second}
+	resp, err := client.Get(url)
+	if err != nil {
+		return nil, fmt.Errorf("下载失败: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("状态码: %d", resp.StatusCode)
+	}
+
+	return io.ReadAll(resp.Body)
+}
+
+// 发送图片识别请求
+func createAndSendMultipartRequest(url string, imageData []byte, formFields map[string]string) ([]byte, error) {
+	body := &bytes.Buffer{}
+	writer := multipart.NewWriter(body)
+
+	part, err := writer.CreateFormFile("file", "image.jpg")
+	if err != nil {
+		return nil, fmt.Errorf("创建文件字段失败: %v", err)
+	}
+	if _, err := io.Copy(part, bytes.NewReader(imageData)); err != nil {
+		return nil, fmt.Errorf("写入文件数据失败: %v", err)
+	}
+
+	for key, value := range formFields {
+		_ = writer.WriteField(key, value)
+	}
+	writer.Close()
+
+	client := &http.Client{Timeout: 15 * time.Second}
+	resp, err := client.Post(url, writer.FormDataContentType(), body)
+	if err != nil {
+		return nil, fmt.Errorf("请求失败: %v", err)
+	}
+	defer resp.Body.Close()
+	return io.ReadAll(resp.Body)
+}
+
+// 裁剪图片并返回 Base64
+func cropImage(buffer []byte, box []float64) (string, error) {
+	contentType := http.DetectContentType(buffer)
+	if contentType != "image/jpeg" && contentType != "image/png" {
+		return "", fmt.Errorf("不支持的图片格式: %s", contentType)
+	}
+
+	img, format, err := image.Decode(bytes.NewReader(buffer))
+	if err != nil {
+		return "", fmt.Errorf("解码失败: %v", err)
+	}
+
+	bounds := img.Bounds()
+	width, height := bounds.Dx(), bounds.Dy()
+	cropX, cropY := int(float64(width)*box[0]), int(float64(height)*box[1])
+	cropW, cropH := int(float64(width)*(box[2]-box[0])), int(float64(height)*(box[3]-box[1]))
+
+	if cropW <= 0 || cropH <= 0 {
+		return "", errors.New("裁剪区域无效")
+	}
+
+	croppedImg := img.(interface {
+		SubImage(r image.Rectangle) image.Image
+	}).SubImage(image.Rect(cropX, cropY, cropX+cropW, cropY+cropH))
+
+	var buf bytes.Buffer
+	switch format {
+	case "jpeg":
+		err = jpeg.Encode(&buf, croppedImg, nil)
+	case "png":
+		err = png.Encode(&buf, croppedImg)
+	default:
+		return "", fmt.Errorf("不支持的格式: %s", format)
+	}
+	if err != nil {
+		return "", fmt.Errorf("图片编码失败: %v", err)
+	}
+
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), nil
+}