mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2026-03-01 02:20:25 +00:00
分离搜图与随机图(随机图需要cgo而搜图不用)
This commit is contained in:
fumiama
136
picsearcher/pic_searcher.go
Normal file
136
picsearcher/pic_searcher.go
Normal file
@@ -0,0 +1,136 @@
|
||||
package picsearcher
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
zero "github.com/wdvxdr1123/ZeroBot"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
|
||||
"github.com/Yiwen-Chan/ZeroBot-Plugin/api/pixiv"
|
||||
utils "github.com/Yiwen-Chan/ZeroBot-Plugin/picsearcher/utils"
|
||||
)
|
||||
|
||||
var (
|
||||
BOTPATH = pixiv.PathExecute() // 当前bot运行目录
|
||||
DATAPATH = BOTPATH + "data/SetuTime/" // 数据目录
|
||||
CACHEPATH = DATAPATH + "cache/" // 缓冲图片路径
|
||||
)
|
||||
|
||||
func init() { // 插件主体
|
||||
// 根据PID搜图
|
||||
zero.OnRegex(`^搜图(\d+)$`).SetBlock(true).SetPriority(30).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
id := pixiv.Str2Int(ctx.State["regex_matched"].([]string)[1])
|
||||
ctx.Send("少女祈祷中......")
|
||||
// 获取P站插图信息
|
||||
illust := &pixiv.Illust{}
|
||||
if err := illust.IllustInfo(id); err != nil {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
return
|
||||
}
|
||||
// 下载P站插图
|
||||
if _, err := illust.PixivPicDown(CACHEPATH); err != nil {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
return
|
||||
}
|
||||
// 发送搜索结果
|
||||
ctx.Send(illust.DetailPic)
|
||||
return
|
||||
})
|
||||
// 以图搜图
|
||||
zero.OnMessage(FullMatchText("以图搜图", "搜索图片", "以图识图"), MustHasPicture()).SetBlock(true).SetPriority(999).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
// 开始搜索图片
|
||||
ctx.Send("少女祈祷中......")
|
||||
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.SendChain(message.Text("ERROR: ", err))
|
||||
}
|
||||
if m, err := utils.Ascii2dSearch(pic); err == nil {
|
||||
ctx.SendChain(m...) // 返回Ascii2d的结果
|
||||
continue
|
||||
} else {
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
98
picsearcher/utils/ascii2d.go
Normal file
98
picsearcher/utils/ascii2d.go
Normal file
@@ -0,0 +1,98 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/Yiwen-Chan/ZeroBot-Plugin/api/pixiv"
|
||||
xpath "github.com/antchfx/htmlquery"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
)
|
||||
|
||||
// Ascii2dSearch Ascii2d 以图搜图
|
||||
// 第一个参数 返回错误
|
||||
// 第二个参数 返回的信息
|
||||
func Ascii2dSearch(pic string) (message.Message, error) {
|
||||
var (
|
||||
api = "https://ascii2d.net/search/uri"
|
||||
)
|
||||
transport := http.Transport{
|
||||
DisableKeepAlives: true,
|
||||
}
|
||||
client := &http.Client{
|
||||
Transport: &transport,
|
||||
}
|
||||
|
||||
// 包装请求参数
|
||||
data := url.Values{}
|
||||
data.Set("uri", pic) // 图片链接
|
||||
fromData := strings.NewReader(data.Encode())
|
||||
|
||||
// 网络请求
|
||||
req, _ := http.NewRequest("POST", api, fromData)
|
||||
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 nil, err
|
||||
}
|
||||
// 色合检索改变到特征检索
|
||||
var bovwUrl = strings.ReplaceAll(resp.Request.URL.String(), "color", "bovw")
|
||||
bovwReq, _ := http.NewRequest("POST", bovwUrl, nil)
|
||||
bovwReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
bovwReq.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.104 Safari/537.36")
|
||||
bovwResp, err := client.Do(bovwReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer bovwResp.Body.Close()
|
||||
// 解析XPATH
|
||||
doc, err := xpath.Parse(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 取出每个返回的结果
|
||||
list := xpath.Find(doc, `//div[@class="row item-box"]`)
|
||||
var link string
|
||||
// 遍历取出第一个返回的PIXIV结果
|
||||
for _, n := range list {
|
||||
linkPath := xpath.Find(n, `//div[2]/div[3]/h6/a[1]`)
|
||||
picPath := xpath.Find(n, `//div[1]/img`)
|
||||
if len(linkPath) != 0 && len(picPath) != 0 {
|
||||
link = xpath.SelectAttr(linkPath[0], "href")
|
||||
if strings.Contains(link, "www.pixiv.net") {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
// 链接取出PIXIV id
|
||||
var index = strings.LastIndex(link, "/")
|
||||
if link == "" || index == -1 {
|
||||
return nil, fmt.Errorf("Ascii2d not found")
|
||||
}
|
||||
var id = pixiv.Str2Int(link[index+1:])
|
||||
if id == 0 {
|
||||
return nil, fmt.Errorf("convert to pid error")
|
||||
}
|
||||
// 根据PID查询插图信息
|
||||
var illust = &pixiv.Illust{}
|
||||
if err := illust.IllustInfo(id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if illust.AgeLimit != "all-age" {
|
||||
return nil, fmt.Errorf("Ascii2d not found")
|
||||
}
|
||||
// 返回插图信息文本
|
||||
return message.Message{
|
||||
message.Text(
|
||||
"[SetuTime] emmm大概是这个？", "\n",
|
||||
"标题：", illust.Title, "\n",
|
||||
"插画ID：", illust.Pid, "\n",
|
||||
"画师：", illust.UserName, "\n",
|
||||
"画师ID：", illust.UserId, "\n",
|
||||
"直链：", "https://pixivel.moe/detail?id=", illust.Pid,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
77
picsearcher/utils/saucenao.go
Normal file
77
picsearcher/utils/saucenao.go
Normal file
@@ -0,0 +1,77 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
)
|
||||
|
||||
// SauceNaoSearch SauceNao 以图搜图 需要链接 返回错误和信息
|
||||
func SauceNaoSearch(pic string) (message.Message, error) {
|
||||
var (
|
||||
api = "https://saucenao.com/search.php"
|
||||
apiKey = "2cc2772ca550dbacb4c35731a79d341d1a143cb5"
|
||||
|
||||
minSimilarity = 70.0 // 返回图片结果的最小相似度
|
||||
)
|
||||
|
||||
// 包装请求参数
|
||||
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()
|
||||
|
||||
// 网络请求
|
||||
client := &http.Client{}
|
||||
req, err := http.NewRequest("GET", link.String(), nil)
|
||||
if err != nil {
|
||||
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 nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := ioutil.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// 如果返回不是200则立刻抛出错误
|
||||
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 nil, fmt.Errorf("SauceNAO not found, status %d", status)
|
||||
}
|
||||
if content.Get("results.0.header.similarity").Float() < minSimilarity {
|
||||
return nil, fmt.Errorf("SauceNAO not found")
|
||||
}
|
||||
result := content.Get("results.0")
|
||||
// 正常发送
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user