✨ add plugin nsfw

README.md

@@ -147,6 +147,9 @@ zerobot [-h] [-t token] [-u url] [-n nickname] [-p prefix] [-d|w] [-g 监听地
     - [x] 刷新所有本地setu
     - [x] 所有本地setu分类
     - 注：刷新文件夹较慢，请耐心等待刷新完成，会提示“成功”。
+- **nsfw图片识别** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_nsfw"`
+    - [x] nsfw打分[图片]
+    - [x] 当图片属于非 neutral 类别时自动发送评价(默认禁用，启用输入 /启用 nsfwauto)
 - **lolicon** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_lolicon"`
     - [x] 来份萝莉
 - **搜图** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_saucenao"`

go.mod

@@ -3,7 +3,7 @@ module github.com/FloatTech/ZeroBot-Plugin
 go 1.17
 
 require (
-	github.com/FloatTech/AnimeAPI v1.2.5-beta1
+	github.com/FloatTech/AnimeAPI v1.2.5-beta2
 	github.com/FloatTech/zbputils v1.2.5-beta1
 	github.com/antchfx/htmlquery v1.2.4
 	github.com/corona10/goimagehash v1.0.3

go.sum

@@ -1,5 +1,5 @@
-github.com/FloatTech/AnimeAPI v1.2.5-beta1 h1:Iv/Uy7ODu06Rd0n9CzUsivfHfI/Nmg4cHHdMU5j6FX4=
+github.com/FloatTech/AnimeAPI v1.2.5-beta2 h1:O8PGLzLLQNGNtJoq+gkpvkEe2JdaZ3vP9pD4+0ASPPk=
-github.com/FloatTech/AnimeAPI v1.2.5-beta1/go.mod h1:tWU3xDvkSqnFD8SCTS5oko1UhkVlrNLK6PmN0ea6n6Q=
+github.com/FloatTech/AnimeAPI v1.2.5-beta2/go.mod h1:tWU3xDvkSqnFD8SCTS5oko1UhkVlrNLK6PmN0ea6n6Q=
 github.com/FloatTech/bot-manager v1.0.0/go.mod h1:8YYRJ16oroGHQGD2En0oVnmcKJkxR9O/jd5BPSfWfOQ=
 github.com/FloatTech/zbputils v1.2.5-beta1 h1:wptmjPpsaoK77igPKFmdUB86AWPyNyuEtT7LNYt2t4U=
 github.com/FloatTech/zbputils v1.2.5-beta1/go.mod h1:1VE4gxnrr+uRz/TDPgx3GT8GmVayoF5e0wolBQU3t+k=

main.go

@@ -67,6 +67,7 @@ import (
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin_lolicon"       // lolicon 随机图片
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin_nativesetu"    // 本地涩图
 	_ "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_setutime"      // 来份涩图
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin_tracemoe"      // 搜番

order/prio.go

@@ -38,6 +38,8 @@ const (
 	PrioNativeWife
 	PrioNBNHHSH
 	PrioNovel
+	PrioNSFW
+	PrioNSFWAuto
 	PrioOmikuji
 	PrioReborn
 	PrioRuncode

plugin_nsfw/main.go

@@ -0,0 +1,69 @@
+package nsfw
+
+import (
+	"github.com/FloatTech/AnimeAPI/nsfw"
+	"github.com/FloatTech/zbputils/control"
+	"github.com/FloatTech/zbputils/ctxext"
+	zero "github.com/wdvxdr1123/ZeroBot"
+	"github.com/wdvxdr1123/ZeroBot/message"
+
+	"github.com/FloatTech/ZeroBot-Plugin/order"
+)
+
+func init() {
+	engine := control.Register("nsfw", order.PrioNSFW, &control.Options{
+		DisableOnDefault: false,
+		Help:             "nsfw图片识别\n- nsfw打分[图片]",
+	})
+	// 上传一张图进行评价
+	engine.OnKeywordGroup([]string{"nsfw打分"}, zero.OnlyPublic, ctxext.CmdMatch, ctxext.MustGiven).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...)
+				if err != nil {
+					ctx.SendChain(message.Text("ERROR:", err))
+					return
+				}
+				ctx.Send(message.ReplyWithMessage(ctx.Event.MessageID, message.Text(judge(p[0]))))
+			}
+		})
+	en := control.Register("nsfwauto", order.PrioNSFWAuto, &control.Options{
+		DisableOnDefault: true,
+		Help:             "nsfw图片自动识别\n- 当图片属于非 neutral 类别时自动发送评价",
+	})
+	en.OnMessage(ctxext.Exists).SetBlock(false).
+		Handle(func(ctx *zero.Ctx) {
+			url := ctx.State["image_url"].([]string)
+			if len(url) > 0 {
+				p, err := nsfw.Classify(url...)
+				if err != nil {
+					return
+				}
+				ctx.Send(message.ReplyWithMessage(ctx.Event.MessageID, message.Text(judge(p[0]))))
+			}
+		})
+}
+
+func judge(p nsfw.Picture) string {
+	if p.Neutral > 0.3 {
+		return "普通哦"
+	}
+	c := ""
+	if p.Drawings > 0.3 {
+		c = "二次元"
+	} else {
+		c = "三次元"
+	}
+	if p.Hentai > 0.3 {
+		c += " hentai"
+	}
+	if p.Porn > 0.3 {
+		c += " porn"
+	}
+	if p.Sexy > 0.3 {
+		c += " hso"
+	}
+	return c
+}