mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2025-12-19 13:59:39 +08:00
⚡️ 更好的识图触发模式
This commit is contained in:
Yiwen-Chan
parent
67622a7668
commit
66ff92e064
@ -2,9 +2,13 @@ package setutime
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
utils "github.com/Yiwen-Chan/ZeroBot-Plugin/setutime/utils"
|
||||
zero "github.com/wdvxdr1123/ZeroBot"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
)
|
||||
|
||||
func init() { // 插件主体
|
||||
@ -28,111 +32,97 @@ func init() { // 插件主体
|
||||
ctx.Send(illust.DetailPic)
|
||||
return
|
||||
})
|
||||
// 通过回复以图搜图
|
||||
zero.OnRegex(`\[CQ:reply,id=(.*?)\](.*)搜索图片`).SetBlock(true).SetPriority(32).
|
||||
// 以图搜图
|
||||
zero.OnMessage(FullMatchText("以图搜图", "搜索图片", "以图识图"), MustHasPicture()).SetBlock(true).SetPriority(999).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
var pics []string // 图片搜索池子
|
||||
// 获取回复的上文图片链接
|
||||
id := utils.Str2Int(ctx.State["regex_matched"].([]string)[1])
|
||||
for _, elem := range ctx.GetMessage(id).Elements {
|
||||
if elem.Type == "image" {
|
||||
pics = append(pics, elem.Data["url"])
|
||||
}
|
||||
}
|
||||
// 没有收到图片则向用户索取
|
||||
if len(pics) == 0 {
|
||||
ctx.Send("请发送多张图片！")
|
||||
next := ctx.FutureEvent("message", ctx.CheckSession())
|
||||
recv, cancel := next.Repeat()
|
||||
for e := range recv { // 循环获取channel发来的信息
|
||||
if len(e.Message) == 1 && e.Message[0].Type == "text" {
|
||||
cancel() // 如果是纯文本则退出索取
|
||||
break
|
||||
}
|
||||
for _, elem := range e.Message {
|
||||
if elem.Type == "image" { // 将信息中的图片添加到搜索池子
|
||||
pics = append(pics, elem.Data["url"])
|
||||
}
|
||||
}
|
||||
if len(pics) >= 5 {
|
||||
cancel() // 如果是图片数量大于等于5则退出索取
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(pics) == 0 {
|
||||
ctx.Send("没有收到图片，搜图结束......")
|
||||
return
|
||||
}
|
||||
// 开始搜索图片
|
||||
ctx.Send("少女祈祷中......")
|
||||
for _, pic := range pics {
|
||||
if text, err := utils.SauceNaoSearch(pic); err == nil {
|
||||
ctx.Send(text) // 返回SauceNAO的结果
|
||||
for _, pic := range ctx.State["image_url"].([]string) {
|
||||
fmt.Println(pic)
|
||||
if m, err := utils.SauceNaoSearch(pic); err == nil {
|
||||
ctx.SendChain(m...) // 返回SauceNAO的结果
|
||||
continue
|
||||
} else {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
}
|
||||
if text, err := utils.Ascii2dSearch(pic); err == nil {
|
||||
ctx.Send(text) // 返回Ascii2d的结果
|
||||
if m, err := utils.Ascii2dSearch(pic); err == nil {
|
||||
ctx.SendChain(m...) // 返回Ascii2d的结果
|
||||
continue
|
||||
} else {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
}
|
||||
}
|
||||
return
|
||||
})
|
||||
// 通过命令以图搜图
|
||||
zero.OnKeywordGroup([]string{"以图识图", "以图搜图", "搜索图片"}).SetBlock(true).SetPriority(33).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
var pics []string // 图片搜索池子
|
||||
// 获取信息中图片链接
|
||||
for _, elem := range ctx.Event.Message {
|
||||
if elem.Type == "image" {
|
||||
pics = append(pics, elem.Data["url"])
|
||||
}
|
||||
}
|
||||
// 没有收到图片则向用户索取
|
||||
if len(pics) == 0 {
|
||||
ctx.Send("请发送多张图片！")
|
||||
next := ctx.FutureEvent("message", zero.CheckUser(ctx.Event.UserID))
|
||||
recv, cancel := next.Repeat()
|
||||
for e := range recv { // 循环获取channel发来的信息
|
||||
if len(e.Message) == 1 && e.Message[0].Type == "text" {
|
||||
cancel() // 如果是纯文本则退出索取
|
||||
break
|
||||
}
|
||||
for _, elem := range e.Message {
|
||||
if elem.Type == "image" { // 将信息中的图片添加到搜索池子
|
||||
pics = append(pics, elem.Data["url"])
|
||||
}
|
||||
}
|
||||
if len(pics) >= 5 {
|
||||
cancel() // 如果是图片数量大于等于5则退出索取
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(pics) == 0 {
|
||||
ctx.Send("没有收到图片，搜图结束......")
|
||||
return
|
||||
}
|
||||
// 开始搜索图片
|
||||
ctx.Send("少女祈祷中......")
|
||||
for _, pic := range pics {
|
||||
if text, err := utils.SauceNaoSearch(pic); err == nil {
|
||||
ctx.Send(text) // 返回SauceNAO的结果
|
||||
continue
|
||||
} else {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
}
|
||||
if text, err := utils.Ascii2dSearch(pic); err == nil {
|
||||
ctx.Send(text) // 返回Ascii2d的结果
|
||||
continue
|
||||
} else {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
}
|
||||
}
|
||||
return
|
||||
})
|
||||
}
|
||||
|
||||
// FullMatchText 如果信息中文本完全匹配则返回 true
|
||||
func FullMatchText(src ...string) zero.Rule {
|
||||
return func(ctx *zero.Ctx) bool {
|
||||
msg := ctx.Event.Message
|
||||
for _, elem := range msg {
|
||||
if elem.Type == "text" {
|
||||
text := elem.Data["text"]
|
||||
text = strings.ReplaceAll(text, " ", "")
|
||||
text = strings.ReplaceAll(text, "\r", "")
|
||||
text = strings.ReplaceAll(text, "\n", "")
|
||||
for _, s := range src {
|
||||
if text == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// HasPicture 消息含有图片返回 true
|
||||
func HasPicture() zero.Rule {
|
||||
return func(ctx *zero.Ctx) bool {
|
||||
msg := ctx.Event.Message
|
||||
url := []string{}
|
||||
// 如果是回复信息则将信息替换成被回复的那条
|
||||
if msg[0].Type == "reply" {
|
||||
id, _ := strconv.Atoi(msg[0].Data["id"])
|
||||
msg = ctx.GetMessage(int64(id)).Elements
|
||||
}
|
||||
// 遍历信息中所有图片
|
||||
for _, elem := range msg {
|
||||
if elem.Type == "image" {
|
||||
url = append(url, elem.Data["url"])
|
||||
}
|
||||
}
|
||||
// 如果有图片就返回true
|
||||
if len(url) > 0 {
|
||||
ctx.State["image_url"] = url
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// MustHasPicture 消息不存在图片阻塞60秒至有图片，超时返回 false
|
||||
func MustHasPicture() zero.Rule {
|
||||
return func(ctx *zero.Ctx) bool {
|
||||
if HasPicture()(ctx) {
|
||||
return true
|
||||
}
|
||||
// 没有图片就索取
|
||||
ctx.Send("请发送一张图片")
|
||||
next := zero.NewFutureEvent("message", 999, false, zero.CheckUser(ctx.Event.UserID), HasPicture())
|
||||
recv, cancel := next.Repeat()
|
||||
select {
|
||||
case e := <-recv:
|
||||
cancel()
|
||||
newCtx := &zero.Ctx{Event: e, State: zero.State{}}
|
||||
if HasPicture()(newCtx) {
|
||||
ctx.State["image_url"] = newCtx.State["image_url"]
|
||||
return true
|
||||
}
|
||||
return false
|
||||
case <-time.After(time.Second * 60):
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,18 +1,17 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
)
|
||||
|
||||
// SauceNaoSearch SauceNao 以图搜图 需要链接 返回错误和信息
|
||||
func SauceNaoSearch(pic string) (text string, err error) {
|
||||
func SauceNaoSearch(pic string) (message.Message, error) {
|
||||
var (
|
||||
api = "https://saucenao.com/search.php"
|
||||
apiKey = "2cc2772ca550dbacb4c35731a79d341d1a143cb5"
|
||||
@ -20,65 +19,59 @@ func SauceNaoSearch(pic string) (text string, err error) {
|
||||
minSimilarity = 70.0 // 返回图片结果的最小相似度
|
||||
)
|
||||
|
||||
transport := http.Transport{
|
||||
DisableKeepAlives: true,
|
||||
}
|
||||
client := &http.Client{
|
||||
Transport: &transport,
|
||||
}
|
||||
|
||||
// 包装请求参数
|
||||
data := url.Values{}
|
||||
data.Set("url", pic) // 图片链接
|
||||
data.Set("api_key", apiKey) // api_key
|
||||
data.Set("db", "5") // 只搜索Pixiv
|
||||
data.Set("numres", "1") // 返回一个结果
|
||||
data.Set("output_type", "2") // 返回JSON格式数据
|
||||
fromData := strings.NewReader(data.Encode())
|
||||
link, _ := url.Parse(api)
|
||||
link.RawQuery = url.Values{
|
||||
"url": []string{pic},
|
||||
"api_key": []string{apiKey},
|
||||
"db": []string{"5"},
|
||||
"numres": []string{"1"},
|
||||
"output_type": []string{"2"},
|
||||
}.Encode()
|
||||
|
||||
// 网络请求
|
||||
req, err := http.NewRequest("POST", api, fromData)
|
||||
client := &http.Client{}
|
||||
req, err := http.NewRequest("GET", link.String(), nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:6.0) Gecko/20100101 Firefox/6.0")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := ioutil.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if code := resp.StatusCode; code != 200 {
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// 如果返回不是200则立刻抛出错误
|
||||
return "", errors.New(fmt.Sprintf("SauceNAO not found, code %d", code))
|
||||
return nil, fmt.Errorf("SauceNAO not found, code %d", resp.StatusCode)
|
||||
}
|
||||
content := gjson.ParseBytes(body)
|
||||
if status := content.Get("header.status").Int(); status != 0 {
|
||||
// 如果json信息返回status不为0则立刻抛出错误
|
||||
return "", errors.New(fmt.Sprintf("SauceNAO not found, status %d", status))
|
||||
return nil, fmt.Errorf("SauceNAO not found, status %d", status)
|
||||
}
|
||||
if content.Get("results.0.header.similarity").Float() < minSimilarity {
|
||||
return "", errors.New("SauceNAO not found")
|
||||
return nil, fmt.Errorf("SauceNAO not found")
|
||||
}
|
||||
result := content.Get("results.0")
|
||||
// 正常发送
|
||||
return fmt.Sprintf(
|
||||
`[SetuTime] 我有把握是这个！[CQ:image,file=%s]相似度：%s%%
|
||||
标题：%s
|
||||
插画ID：%d
|
||||
画师：%s
|
||||
画师ID：%d
|
||||
直链：https://pixivel.moe/detail?id=%d`,
|
||||
content.Get("results.0.header.thumbnail").Str,
|
||||
content.Get("results.0.header.similarity").Str,
|
||||
content.Get("results.0.data.title").Str,
|
||||
content.Get("results.0.data.pixiv_id").Int(),
|
||||
content.Get("results.0.data.member_name").Str,
|
||||
content.Get("results.0.data.member_id").Int(),
|
||||
content.Get("results.0.data.pixiv_id").Int(),
|
||||
), nil
|
||||
return message.Message{
|
||||
message.Text("[SetuTime] 我有把握是这个！"),
|
||||
message.Image(result.Get("header.thumbnail").Str),
|
||||
message.Text(
|
||||
"\n",
|
||||
"相似度：", result.Get("header.similarity").Str, "\n",
|
||||
"标题：", result.Get("data.title").Str, "\n",
|
||||
"插画ID：", result.Get("data.pixiv_id").Int(), "\n",
|
||||
"画师：", result.Get("data.member_name").Str, "\n",
|
||||
"画师ID：", result.Get("data.member_id").Int(), "\n",
|
||||
"直链：", "https://pixivel.moe/detail?id=", result.Get("data.pixiv_id").Int(),
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user