|
| 1 | +// Package animetrace AnimeTrace 动画/Galgame识别 |
| 2 | +package animetrace |
| 3 | + |
| 4 | +import ( |
| 5 | + "bytes" |
| 6 | + "encoding/base64" |
| 7 | + "errors" |
| 8 | + "fmt" |
| 9 | + "image" |
| 10 | + "image/jpeg" |
| 11 | + "mime/multipart" |
| 12 | + "strings" |
| 13 | + |
| 14 | + "github.com/FloatTech/floatbox/web" |
| 15 | + "github.com/FloatTech/imgfactory" |
| 16 | + ctrl "github.com/FloatTech/zbpctrl" |
| 17 | + "github.com/FloatTech/zbputils/control" |
| 18 | + "github.com/FloatTech/zbputils/ctxext" |
| 19 | + "github.com/disintegration/imaging" |
| 20 | + "github.com/tidwall/gjson" |
| 21 | + zero "github.com/wdvxdr1123/ZeroBot" |
| 22 | + "github.com/wdvxdr1123/ZeroBot/message" |
| 23 | +) |
| 24 | + |
| 25 | +func init() { |
| 26 | + engine := control.Register("animetrace", &ctrl.Options[*zero.Ctx]{ |
| 27 | + DisableOnDefault: false, |
| 28 | + Brief: "AnimeTrace 动画/Galgame识别插件", |
| 29 | + Help: "- Gal识图\n- 动漫识图\n- 动漫识图 2\n- 动漫识图 [模型名]\n- Gal识图 [模型名]", |
| 30 | + }) |
| 31 | + |
| 32 | + engine.OnPrefix("gal识图", zero.OnlyGroup, zero.MustProvidePicture).SetBlock(true).Handle(func(ctx *zero.Ctx) { |
| 33 | + args := ctx.State["args"].(string) |
| 34 | + var model string |
| 35 | + switch strings.TrimSpace(args) { |
| 36 | + case "": |
| 37 | + model = "full_game_model_kira" // 默认使用的模型 |
| 38 | + default: |
| 39 | + model = args // 自定义设置模型 |
| 40 | + } |
| 41 | + processImageRecognition(ctx, model) |
| 42 | + }) |
| 43 | + |
| 44 | + engine.OnPrefix("动漫识图", zero.OnlyGroup, zero.MustProvidePicture).SetBlock(true).Handle(func(ctx *zero.Ctx) { |
| 45 | + args := ctx.State["args"].(string) |
| 46 | + var model string |
| 47 | + switch strings.TrimSpace(args) { |
| 48 | + case "": |
| 49 | + model = "anime_model_lovelive" |
| 50 | + case "2": |
| 51 | + model = "pre_stable" |
| 52 | + default: |
| 53 | + model = args |
| 54 | + } |
| 55 | + processImageRecognition(ctx, model) |
| 56 | + }) |
| 57 | +} |
| 58 | + |
| 59 | +// 处理图片识别 |
| 60 | +func processImageRecognition(ctx *zero.Ctx, model string) { |
| 61 | + urls := ctx.State["image_url"].([]string) |
| 62 | + if len(urls) == 0 { |
| 63 | + return |
| 64 | + } |
| 65 | + imageData, err := imgfactory.Load(urls[0]) |
| 66 | + if err != nil { |
| 67 | + ctx.Send(message.Text("下载图片失败: ", err)) |
| 68 | + return |
| 69 | + } |
| 70 | + //ctx.Send(message.Text(model)) |
| 71 | + respBody, err := createAndSendMultipartRequest("https://api.animetrace.com/v1/search", imageData, map[string]string{ |
| 72 | + "is_multi": "0", |
| 73 | + "model": model, |
| 74 | + "ai_detect": "0", |
| 75 | + }) |
| 76 | + if err != nil { |
| 77 | + ctx.Send(message.Text("识别请求失败: ", err)) |
| 78 | + return |
| 79 | + } |
| 80 | + code := gjson.Get(string(respBody), "code").Int() |
| 81 | + if code != 0 { |
| 82 | + ctx.Send(message.Text("错误: ", gjson.Get(string(respBody), "zh_message").String())) |
| 83 | + return |
| 84 | + } |
| 85 | + dataArray := gjson.Get(string(respBody), "data").Array() |
| 86 | + if len(dataArray) == 0 { |
| 87 | + ctx.Send(message.Text("未识别到任何角色")) |
| 88 | + return |
| 89 | + } |
| 90 | + var sk message.Message |
| 91 | + sk = append(sk, ctxext.FakeSenderForwardNode(ctx, message.Text("共识别到 ", len(dataArray), " 个角色,可能是以下来源"))) |
| 92 | + for _, value := range dataArray { |
| 93 | + boxArray := value.Get("box").Array() |
| 94 | + imgWidth, imgHeight := imageData.Bounds().Dx(), imageData.Bounds().Dy() // 你可以从 `imageData.Bounds()` 获取 |
| 95 | + box := []int{ |
| 96 | + int(boxArray[0].Float() * float64(imgWidth)), |
| 97 | + int(boxArray[1].Float() * float64(imgHeight)), |
| 98 | + int(boxArray[2].Float() * float64(imgWidth)), |
| 99 | + int(boxArray[3].Float() * float64(imgHeight)), |
| 100 | + } |
| 101 | + croppedImg := imaging.Crop(imageData, image.Rect(box[0], box[1], box[2], box[3])) |
| 102 | + var buf bytes.Buffer |
| 103 | + if err := imaging.Encode(&buf, croppedImg, imaging.JPEG, imaging.JPEGQuality(80)); err != nil { |
| 104 | + ctx.Send(message.Text("图片编码失败: ", err)) |
| 105 | + continue |
| 106 | + } |
| 107 | + |
| 108 | + base64Str := base64.StdEncoding.EncodeToString(buf.Bytes()) |
| 109 | + var sb strings.Builder |
| 110 | + value.Get("character").ForEach(func(_, character gjson.Result) bool { |
| 111 | + sb.WriteString(fmt.Sprintf("《%s》的角色 %s\n", character.Get("work").String(), character.Get("character").String())) |
| 112 | + return true |
| 113 | + }) |
| 114 | + sk = append(sk, ctxext.FakeSenderForwardNode(ctx, message.Image("base64://"+base64Str), message.Text(sb.String()))) |
| 115 | + } |
| 116 | + ctx.SendGroupForwardMessage(ctx.Event.GroupID, sk) |
| 117 | +} |
| 118 | + |
| 119 | +// 发送图片识别请求 |
| 120 | +func createAndSendMultipartRequest(url string, img image.Image, formFields map[string]string) ([]byte, error) { |
| 121 | + body := &bytes.Buffer{} |
| 122 | + writer := multipart.NewWriter(body) |
| 123 | + |
| 124 | + // 直接编码图片 |
| 125 | + part, err := writer.CreateFormFile("file", "image.jpg") |
| 126 | + if err != nil { |
| 127 | + return nil, errors.New("创建文件字段失败: " + err.Error()) |
| 128 | + } |
| 129 | + if err := jpeg.Encode(part, img, &jpeg.Options{Quality: 80}); err != nil { |
| 130 | + return nil, errors.New("图片编码失败: " + err.Error()) |
| 131 | + } |
| 132 | + |
| 133 | + // 写入其他字段 |
| 134 | + for key, value := range formFields { |
| 135 | + if err := writer.WriteField(key, value); err != nil { |
| 136 | + return nil, errors.New("写入表单字段失败 (" + key + "): " + err.Error()) |
| 137 | + } |
| 138 | + } |
| 139 | + |
| 140 | + if err := writer.Close(); err != nil { |
| 141 | + return nil, errors.New("关闭 multipart writer 失败: " + err.Error()) |
| 142 | + } |
| 143 | + |
| 144 | + return web.PostData(url, writer.FormDataContentType(), body) |
| 145 | +} |
0 commit comments