diff --git a/README.md b/README.md index 06d6223e..4fb02707 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,8 @@ zerobot [-h] [-t token] [-u url] [-n nickname] [-p prefix] [-d|w] [-g 监听地 - [x] 评价图片(发送一张图片让bot评分) - **DeepDanbooru二次元图标签识别** `import _ github.com/FloatTech/ZeroBot-Plugin/plugin_danbooru` - [x] 鉴赏图片[图片] +- **叔叔的AI二次元图片放大** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_scale"` + - [x] 放大图片[图片] - **每日运势** `import _ github.com/FloatTech/ZeroBot-Plugin/plugin_fortune` - [x] 运势 | 抽签 - [x] 设置底图[车万 DC4 爱因斯坦 星空列车 樱云之恋 富婆妹 李清歌 公主连结 原神 明日方舟 碧蓝航线 碧蓝幻想 战双 阴阳师 赛马娘] diff --git a/data b/data index f84539a9..66341033 160000 --- a/data +++ b/data @@ -1 +1 @@ -Subproject commit f84539a9f20934d99d94d62cfb9ad6e193df3617 +Subproject commit 663410335c3a6a340bc9c0df0922c08aa5037edd diff --git a/main.go b/main.go index 9f4d6819..45492d9c 100644 --- a/main.go +++ b/main.go @@ -69,6 +69,7 @@ import ( _ "github.com/FloatTech/ZeroBot-Plugin/plugin_nativewife" // 本地老婆 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_nsfw" // nsfw图片识别 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_saucenao" // 以图搜图 + _ "github.com/FloatTech/ZeroBot-Plugin/plugin_scale" // 叔叔的AI二次元图片放大 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_setutime" // 来份涩图 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_tracemoe" // 搜番 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation" // vtb语录 diff --git a/order/prio.go b/order/prio.go index 740144fb..9ba1be17 100644 --- a/order/prio.go +++ b/order/prio.go @@ -44,6 +44,7 @@ const ( PrioReborn PrioRuncode PrioSauceNao + PrioScale PrioScore PrioSetuTime PrioShaDiao diff --git a/plugin_acgimage/classify.go b/plugin_acgimage/classify.go index fc96f394..b94526f8 100644 --- a/plugin_acgimage/classify.go +++ b/plugin_acgimage/classify.go @@ -89,7 +89,7 @@ func init() { // 插件主体 } }) // 上传一张图进行评价 - engine.OnKeywordGroup([]string{"评价图片"}, zero.OnlyPublic, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). + engine.OnKeywordGroup([]string{"评价图片"}, zero.OnlyGroup, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). Handle(func(ctx *zero.Ctx) { ctx.SendChain(message.Text("少女祈祷中...")) for _, url := range ctx.State["image_url"].([]string) { diff --git a/plugin_danbooru/main.go b/plugin_danbooru/main.go index cab58a6d..3adf0ebb 100644 --- a/plugin_danbooru/main.go +++ b/plugin_danbooru/main.go @@ -32,7 +32,7 @@ func init() { // 插件主体 "- 鉴赏图片[图片]", }) // 上传一张图进行评价 - engine.OnKeywordGroup([]string{"鉴赏图片"}, zero.OnlyPublic, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). + engine.OnKeywordGroup([]string{"鉴赏图片"}, zero.OnlyGroup, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). Handle(func(ctx *zero.Ctx) { ctx.SendChain(message.Text("少女祈祷中...")) for _, url := range ctx.State["image_url"].([]string) { diff --git a/plugin_nsfw/main.go b/plugin_nsfw/main.go index 784c5cad..648a8b4b 100644 --- a/plugin_nsfw/main.go +++ b/plugin_nsfw/main.go @@ -16,7 +16,7 @@ func init() { Help: "nsfw图片识别\n- nsfw打分[图片]", }) // 上传一张图进行评价 - engine.OnKeywordGroup([]string{"nsfw打分"}, zero.OnlyPublic, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). + engine.OnKeywordGroup([]string{"nsfw打分"}, zero.OnlyGroup, ctxext.CmdMatch, ctxext.MustGiven).SetBlock(true). Handle(func(ctx *zero.Ctx) { url := ctx.State["image_url"].([]string) if len(url) > 0 { @@ -51,7 +51,7 @@ func judge(p nsfw.Picture) string { return "普通哦" } c := "" - if p.Drawings > 0.3 { + if p.Drawings > 0.3 || p.Neutral < 0.3 { c = "二次元" } else { c = "三次元" diff --git a/plugin_scale/main.go b/plugin_scale/main.go new file mode 100644 index 00000000..e207576d --- /dev/null +++ b/plugin_scale/main.go @@ -0,0 +1,104 @@ +// Package scale 叔叔的AI二次元图片放大 +package scale + +import ( + "os" + "strconv" + "time" + + "github.com/FloatTech/AnimeAPI/nsfw" + "github.com/FloatTech/AnimeAPI/scale" + "github.com/FloatTech/zbputils/control" + "github.com/FloatTech/zbputils/ctxext" + "github.com/FloatTech/zbputils/file" + zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" + + "github.com/FloatTech/ZeroBot-Plugin/order" +) + +const cachedir = "data/scale/" + +func init() { + _ = os.RemoveAll(cachedir) + err := os.MkdirAll(cachedir, 0755) + if err != nil { + panic(err) + } + engine := control.Register("scale", order.PrioScale, &control.Options{ + DisableOnDefault: false, + Help: "叔叔的AI二次元图片放大\n- 放大图片[图片]", + }) + // 上传一张图进行评价 + engine.OnKeywordGroup([]string{"放大图片"}, zero.OnlyGroup, ctxext.CmdMatch, ctxext.MustGiven, getPara).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + url := ctx.State["image_url"].([]string) + if len(url) > 0 { + ctx.SendChain(message.Text("少女祈祷中...")) + p, err := nsfw.Classify(url[0]) + if err != nil { + ctx.SendChain(message.Text("ERROR:", err)) + return + } + if p[0].Drawings < 0.1 || p[0].Neutral > 0.8 { + ctx.SendChain(message.Text("请发送二次元图片!")) + return + } + paras := ctx.State["scale_paras"].([2]int) + data, err := scale.Get(url[0], paras[0], paras[1], 2) + if err != nil { + ctx.SendChain(message.Text("ERROR:", err)) + return + } + n := cachedir + strconv.Itoa(int(ctx.Event.UserID)) + f, err := os.Create(n) + if err != nil { + ctx.SendChain(message.Text("ERROR:", err)) + return + } + f.Write(data) + f.Close() + ctx.SendChain(message.Image("file:///" + file.BOTPATH + "/" + n)) + } + }) +} + +func getPara(ctx *zero.Ctx) bool { + next := zero.NewFutureEvent("message", 999, false, zero.CheckUser(ctx.Event.UserID)) + recv, cancel := next.Repeat() + i := 0 + paras := [2]int{} + ctx.SendChain(message.Text("请输入模型序号\n0.", scale.Models[0], "\n1.", scale.Models[1], "\n2.", scale.Models[2], "\n3.", scale.Models[3], "\n4.", scale.Models[4])) + for { + select { + case <-time.After(time.Second * 120): + return false + case e := <-recv: + msg := e.Message.ExtractPlainText() + num, err := strconv.Atoi(msg) + if err != nil { + ctx.SendChain(message.Text("请输入数字!")) + continue + } + switch i { + case 0: + if num < 0 || num > 4 { + ctx.SendChain(message.Text("模型序号非法!")) + continue + } + paras[0] = num + ctx.SendChain(message.Text("请输入放大倍数(2-4)")) + case 1: + if num < 2 || num > 4 { + ctx.SendChain(message.Text("放大倍数非法!")) + continue + } + cancel() + paras[1] = num + ctx.State["scale_paras"] = paras + return true + } + i++ + } + } +}