diff --git a/README.md b/README.md index b9f8aec..a579ca0 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # kenaito-media-picker -媒体文件自动分类整理工具：一键扫描目录下的图片 / 视频 / 音频，按 **方向（横屏 / 竖屏 / 等比）+ 分辨率 + 明暗** 自动归入对应子目录，并提供 Web 管理界面（实时进度、运行日志、媒体库、统计看板、可选色情识别）。 +媒体文件自动分类整理工具：一键扫描目录下的图片 / 视频 / 音频，按 **方向（横屏 / 竖屏 / 等比）+ 分辨率 + 明暗** 自动归入对应子目录，并提供 Web 管理界面（实时进度、运行日志、媒体库、统计看板、可选色情识别、可选二次元识别）。 ## 功能特性 @@ -11,6 +11,7 @@ - **明暗分类**：基于 ITU-R BT.601 亮度公式，在分辨率目录下追加 `dark_N` / `mid` / `bright_N` 子目录（N 取 20/40/60/80/100 共 5 档） - **EXIF 旋转修正**：带旋转标记的竖拍图自动处理——有 ffmpeg 时物理转正像素，否则至少按显示方向（宽高对调）归类；可在配置中整体关闭 - **色情识别（可选）**：本机 Python + NudeNet 离线鉴黄，命中的图片移入 `图片_敏感_*` 目录并记录检测分，Web 端可逐张预览复核 +- **二次元识别（可选）**：本机 Python + wd-tagger 识别二次元/动漫风格图片，输出 safe/questionable/explicit 三级分级，命中的图片移入 `图片_二次元_*` 目录并记录分级标签与 top tags，Web 端可逐张预览复核 - **异常兜底**：读取失败、不支持格式自动归入「读取异常」「已忽略」目录；GIF、PSD 单独归类 - **断点续传**：每个文件归类结果实时写入续传表（`scan_log`），扫描中断/出错后重跑时，已归类过的文件直接复用上次分类，不再重新探测尺寸 / 算 MD5 / 分析明暗 / 色情识别 - **媒体信息入库**：自动把文件 MD5 与完整媒体信息（类型 / 宽高 / 时长 / 编码 / 容器等）写入 `media-info.db`（SQLite），支持筛选查询与去重 @@ -46,6 +47,7 @@ Web 管理界面（默认 http://localhost:8080）。 | 统计看板 | 媒体总数 / 总体积 / 平均·最大·最小文件 / 总时长，以及类型、方向、分辨率、编码、格式、扩展名等分布与体积 TOP10 | | 媒体库 | 按类型 / 方向 / 分辨率 / 编码 / 格式 / 文件名筛选，分页 + 列排序 | | 色情识别 | 展示最近一次开启识别扫描命中的图片：检测分 + 完整路径 + 缩略图预览，便于人工复核 | +| 二次元识别 | 展示最近一次开启识别扫描命中的二次元图片：分级（safe/questionable/explicit）+ 检测分 + top 标签 + 完整路径 + 缩略图预览，便于人工复核 | ## 配置（application.properties） @@ -62,6 +64,9 @@ Web 管理界面（默认 http://localhost:8080）。 | `nsfw.enabled` | `false` | 是否开启色情图片识别 | | `nsfw.python` | 空 | python 可执行文件路径；留空自动在 PATH 找 | | `nsfw.threshold` | `0.5` | 判定阈值（0~1），敏感部位检测分达到即判为色情 | +| `hentai.enabled` | `false` | 是否开启二次元图片识别 | +| `hentai.python` | 空 | python 可执行文件路径；留空自动在 PATH 找（可与 nsfw.python 相同） | +| `hentai.threshold` | `0.5` | 判定阈值（0~1），questionable 或 explicit 分数达到即判为二次元敏感 | ## 整理规则 @@ -80,6 +85,7 @@ KMP已整理/ ├── 图片_PSD/ ← PSD 设计稿原样收走 ├── 图片_读取异常/ 图片_已忽略/ ├── 图片_敏感_横屏_1080P/bright_20/ ← 判为色情的图片（明暗子目录照常） +├── 图片_二次元_横屏_1080P/bright_20/ ← 判为二次元敏感的图片（明暗子目录照常） ├── 视频_横屏_1080P/ 视频_竖屏_1K/ 视频_等比_原画质/ ├── 视频_横屏_GIF/ 视频_竖屏_GIF/ 视频_等比_GIF/ ├── 视频_读取异常/ 视频_已忽略/ @@ -121,11 +127,24 @@ lightScore = 亮度 > 175 的像素占比 - **开启**：`application.properties` 设 `nsfw.enabled=true` - **依赖**：本机安装 Python 并 `pip install nudenet` -- **模型**：优先加载程序运行目录下 `lib/320n.onnx`（本仓库已附带，约 11.6MB，离线可用）；本地缺失时自动下载到该位置，下载失败再回退用 nudenet 包自带模型 +- **模型**：优先加载程序运行目录下 `lib/nwfs/320n.onnx`（本仓库已附带，约 11.6MB，离线可用）；本地缺失时自动下载到该位置，下载失败再回退用 nudenet 包自带模型 - **可调**：`nsfw.python`（python 路径）、`nsfw.threshold`（判定阈值） - **敏感类**：`FEMALE_GENITALIA_EXPOSED` / `MALE_GENITALIA_EXPOSED` / `ANUS_EXPOSED` / `FEMALE_BREAST_EXPOSED`，任一分数达标即判定 - **降级**：未装 python / nudenet 时打印警告并跳过识别，图片按普通流程整理，不影响其它功能 +### 二次元图片识别（可选） + +基于本机 **Python + wd-tagger**（SmilingWolf/wd-vit-large-tagger-v3）的二次元图片分级识别，专门针对动漫/二次元风格训练，对 R15/R18 擦边图识别效果优于 NudeNet。判为敏感（questionable/explicit）的图片移入 `图片_二次元_*` 分类目录（明暗子目录照常保留），并可在 Web「二次元识别」页面查看每张图的分级、检测分、top 标签、路径与预览。 + +- **开启**：`application.properties` 设 `hentai.enabled=true` +- **依赖**：本机安装 Python 并 `pip install onnxruntime Pillow numpy` +- **模型**：首次启动自动从 HuggingFace 下载到 `lib/wd-tagger/`（约 1.5GB，需联网）；下载完成后离线可用 +- **可调**：`hentai.python`（python 路径）、`hentai.threshold`（判定阈值） +- **分级输出**：`rating:safe` / `rating:questionable` / `rating:explicit`，questionable 或 explicit 分数达到阈值即判定为敏感 +- **标签输出**：同时输出 top 10 高置信度标签（如 `1girl`, `solo`, `long_hair` 等），便于人工复核 +- **降级**：未装 python / 依赖或模型下载失败时打印警告并跳过识别，图片按普通流程整理，不影响其它功能 +- **与色情识别的关系**：两者独立运行，可同时启用；同一张图优先走色情识别（nsfw > hentai > normal） + ## 构建编译 ```shell @@ -167,8 +186,13 @@ Remove-Item Env:CGO_ENABLED,Env:GOOS,Env:GOARCH sudo yum install -y python3 pip3 install nudenet ``` - 模型文件 `lib/320n.onnx` 需与二进制放在同一目录（没有会自动下载）。 -3. 其余步骤与 Windows 一致：`application.properties`、`media-info.db` 随运行目录自动创建。 + 模型文件 `lib/nwfs/320n.onnx` 需与二进制放在同一目录（没有会自动下载）。 +3. **二次元识别（可选）**：Linux 上同样要装 Python 与依赖，并在 `application.properties` 里配 `hentai.python`（如 `/usr/bin/python3`）： + ```shell + pip3 install onnxruntime Pillow numpy + ``` + 模型文件首次启动会自动从 HuggingFace 下载到 `lib/wd-tagger/`（约 1.5GB，需联网）。 +4. 其余步骤与 Windows 一致：`application.properties`、`media-info.db` 随运行目录自动创建。 ### 前端开发模式（无需重复构建） @@ -188,7 +212,8 @@ npm install && npm run dev ```text ├── main.go # 程序入口（Gin + 内嵌前端产物） ├── application.properties # 运行配置（端口 / 数据库 / 旋转开关 / 色情识别） -├── lib/320n.onnx # 色情识别模型（随仓库分发，离线可用） +├── lib/nwfs/320n.onnx # 色情识别模型（随仓库分发，离线可用） +├── lib/wd-tagger/ # 二次元识别模型（首次启动自动下载，约 1.5GB） ├── build.sh # 编译脚本 ├── internal/ │ ├── scanner/ # 文件扫描、MIME 读取、候选过滤 @@ -196,10 +221,11 @@ npm install && npm run dev │ ├── classify/ # 纯函数归类规则（方向 + 标准分辨率 + K档 + 明暗等级） │ ├── organize/ # 整理编排：探测 → 归类 → 收集 → 移动 → 入库 │ ├── mover/ # 文件移动（rename / 跨盘 copy+delete） -│ ├── store/ # SQLite 媒体信息 / 色情记录持久化 + 查询/统计 +│ ├── store/ # SQLite 媒体信息 / 色情记录 / 二次元记录持久化 + 查询/统计 │ ├── service/ # 后台任务编排（扫描整理任务 + 状态/日志/WebSocket） │ ├── nsfw/ # 色情图片识别（内嵌 python 服务 + NudeNet 本地 HTTP） -│ ├── api/ # HTTP 接口层（扫描 / 媒体库 / 统计 / 色情识别 / 目录浏览） +│ ├── hentai/ # 二次元图片识别（内嵌 python 服务 + wd-tagger 本地 HTTP） +│ ├── api/ # HTTP 接口层（扫描 / 媒体库 / 统计 / 色情识别 / 二次元识别 / 目录浏览） │ └── digest/ # 文件 MD5 / 指纹采样 └── frontend/ # 前端工程（Vue 2 + ElementUI + Vite） ``` diff --git a/application.properties b/application.properties index 731c72d..554267b 100644 --- a/application.properties +++ b/application.properties @@ -14,11 +14,21 @@ organize.rotateImages=false organize.workers=8 # 色情图片识别（需本机安装 python 与 nudenet：pip install nudenet） # true=开启，判为色情的图片移入「图片_敏感_*」目录（明暗子目录照常）；false=关闭（默认） -# 模型：优先加载程序运行目录下 lib/320n.onnx（推荐随程序分发，离线可用）； +# 模型：优先加载程序运行目录下 lib/nwfs/320n.onnx（推荐随程序分发，离线可用）； # 本地没有会自动下载到该位置，下载失败再回退用 nudenet 自带模型 # NSFW是用来识别真人图片的 -nsfw.enabled=true +nsfw.enabled=false # python 可执行文件路径；留空自动在 PATH 里找 python/python3/py nsfw.python=C:\\Users\\Administrator\\AppData\\Local\\Programs\\Python\\Python312\\python.exe # 判定阈值 0~1，敏感部位检测分高于此才算色情（默认 0.2，实测） nsfw.threshold=0.2 + +# 二次元图片识别（需本机安装 python 与 onnxruntime、Pillow、numpy：pip install onnxruntime Pillow numpy） +# true=开启，判为二次元敏感（questionable/explicit）的图片移入「图片_二次元_*」目录；false=关闭（默认） +# 模型：首次启动会自动从 HuggingFace 下载 wd-vit-large-tagger-v3 到 lib/wd-tagger/（约 1.5GB） +# Hentai是用来识别二次元/动漫风格图片的 +hentai.enabled=true +# python 可执行文件路径；留空自动在 PATH 里找 python/python3/py（和 nsfw.python 共用同一个即可） +hentai.python=C:\\Users\\Administrator\\AppData\\Local\\Programs\\Python\\Python312\\python.exe +# 判定阈值 0~1，questionable 或 explicit 分数高于此才算二次元敏感（默认 0.5） +hentai.threshold=0.5 diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 37374b2..acef9b6 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -28,6 +28,10 @@ 色情识别 + + + 二次元识别 +
kenaito-media-picker
diff --git a/frontend/src/router/index.js b/frontend/src/router/index.js index 529f35c..50a21cc 100644 --- a/frontend/src/router/index.js +++ b/frontend/src/router/index.js @@ -3,6 +3,7 @@ import Router from 'vue-router' import Dashboard from '../views/Dashboard.vue' import MediaList from '../views/MediaList.vue' import NsfwList from '../views/NsfwList.vue' +import HentaiList from '../views/HentaiList.vue' import Scan from '../views/Scan.vue' Vue.use(Router) @@ -13,6 +14,7 @@ export default new Router({ {path: '/dashboard', component: Dashboard}, {path: '/media', component: MediaList}, {path: '/nsfw', component: NsfwList}, + {path: '/hentai', component: HentaiList}, {path: '/scan', component: Scan} ] }) diff --git a/frontend/src/views/HentaiList.vue b/frontend/src/views/HentaiList.vue new file mode 100644 index 0000000..5e2a194 --- /dev/null +++ b/frontend/src/views/HentaiList.vue @@ -0,0 +1,255 @@ + + + + + diff --git a/internal/api/hentai_controller.go b/internal/api/hentai_controller.go new file mode 100644 index 0000000..89c7b04 --- /dev/null +++ b/internal/api/hentai_controller.go @@ -0,0 +1,121 @@ +package api + +import ( + "encoding/json" + "net/http" + "os" + "strings" + + "github.com/gin-gonic/gin" + + "kenaito-media-picker/internal/config" + "kenaito-media-picker/internal/store" +) + +// hentaiDetectorReadyKey 记录最近一次扫描时二次元检测服务是否真的起来了。 +const hentaiDetectorReadyKey = "hentai.detectorReady" + +// HentaiController 二次元识别相关接口的控制器。 +type HentaiController struct { + store *store.Store +} + +// NewHentaiController 创建二次元识别控制器。 +func NewHentaiController(st *store.Store) *HentaiController { + return &HentaiController{store: st} +} + +// hentaiView 列表里展示的一条记录。 +type hentaiView struct { + Id int64 `json:"id"` + FileName string `json:"fileName"` + Score float64 `json:"score"` + Rating string `json:"rating"` + Threshold float64 `json:"threshold"` + Path string `json:"path"` + TopTags string `json:"topTags"` + CreatedAt string `json:"createdAt"` +} + +// List 处理 GET /api/hentai/list：查全部二次元识别记录（按检测分从高到低）。 +func (c *HentaiController) List(ctx *gin.Context) { + list, err := c.store.ListHentai() + if err != nil { + Error(ctx, http.StatusInternalServerError, "查询二次元识别记录失败："+err.Error()) + return + } + views := make([]hentaiView, 0, len(list)) + for _, r := range list { + path := r.MovedPath + if path == "" { + path = r.FilePath + } + views = append(views, hentaiView{ + Id: r.Id, + FileName: r.FileName, + Score: r.Score, + Rating: r.Rating, + Threshold: r.Threshold, + Path: path, + TopTags: r.TopTags, + CreatedAt: r.CreatedAt.Format("2006-01-02 15:04:05"), + }) + } + Success(ctx, map[string]interface{}{ + "list": views, + "total": len(views), + "hentaiEnabled": config.GetAppConfig().HentaiEnabled, + "hentaiDetector": c.store == nil || c.store.GetConfig(hentaiDetectorReadyKey) != "0", + }) +} + +// Preview 处理 GET /api/hentai/preview?path=xxx：把指定图片文件返回给浏览器做预览。 +func (c *HentaiController) Preview(ctx *gin.Context) { + path := strings.TrimSpace(ctx.Query("path")) + if path == "" { + Error(ctx, http.StatusBadRequest, "缺少 path 参数") + return + } + if _, err := os.Stat(path); err != nil { + Error(ctx, http.StatusNotFound, "图片不存在或已被移动: "+path) + return + } + ctx.File(path) +} + +// Tags 处理 GET /api/hentai/tags?id=xxx：查某条记录的标签详情。 +func (c *HentaiController) Tags(ctx *gin.Context) { + id := ctx.Query("id") + if id == "" { + Error(ctx, http.StatusBadRequest, "缺少 id 参数") + return + } + var idInt int64 + if _, err := json.Number(id).Int64(); err != nil { + Error(ctx, http.StatusBadRequest, "id 格式错误") + return + } + idInt, _ = json.Number(id).Int64() + + record, err := c.store.GetHentaiById(idInt) + if err != nil { + Error(ctx, http.StatusInternalServerError, "查询记录失败："+err.Error()) + return + } + if record == nil { + Error(ctx, http.StatusNotFound, "记录不存在") + return + } + + var tags []map[string]interface{} + if record.TopTags != "" { + if err := json.Unmarshal([]byte(record.TopTags), &tags); err != nil { + tags = nil + } + } + + Success(ctx, map[string]interface{}{ + "id": record.Id, + "topTags": tags, + }) +} diff --git a/internal/api/router.go b/internal/api/router.go index fb89fb6..8f5d5d7 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -38,6 +38,12 @@ func NewRouter(st *store.Store) *gin.Engine { apiGroup.GET("/nsfw/list", nc.List) apiGroup.GET("/nsfw/preview", nc.Preview) + // 二次元识别（查看分级结果 + 图片预览 + 标签详情） + hc := NewHentaiController(st) + apiGroup.GET("/hentai/list", hc.List) + apiGroup.GET("/hentai/preview", hc.Preview) + apiGroup.GET("/hentai/tags", hc.Tags) + // 扫描整理（开始 / 中止 / 查进度 / 实时进度） sc := NewScanController(service.NewScanService(st)) apiGroup.POST("/scan/start", sc.Start) diff --git a/internal/classify/classify.go b/internal/classify/classify.go index df14dbe..4835146 100644 --- a/internal/classify/classify.go +++ b/internal/classify/classify.go @@ -68,6 +68,13 @@ func ImageDirSensitive(ext string, width, height int, darkScore, lightScore floa return imageDir("图片_敏感", ext, width, height, darkScore, lightScore) } +// ImageDirHentai 返回一张被判为二次元敏感的图片该进的文件夹路径，比如： +// - 「图片_二次元_GIF」：GIF 单独放（不做明暗分类） +// - 「图片_二次元_横屏_1080P/bright_20」：普通分类前加「_二次元」前缀，明暗仍是子目录 +func ImageDirHentai(ext string, width, height int, darkScore, lightScore float64) string { + return imageDir("图片_二次元", ext, width, height, darkScore, lightScore) +} + // imageDir 拼出"类型 + 方向 + 清晰度 + 明暗子目录"这段路径。 func imageDir(kind, ext string, width, height int, darkScore, lightScore float64) string { if strings.EqualFold(ext, ".gif") { diff --git a/internal/config/config.go b/internal/config/config.go index 57eddb9..39d6196 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -22,6 +22,9 @@ import ( // - NSFWEnabled：是否开启色情图片识别（本地 python + NudeNet，默认 false） // - NSFWPython：python 可执行文件路径；留空自动在 PATH 里找 // - NSFWThreshold：判定阈值 0~1（默认 0.5） +// - HentaiEnabled：是否开启二次元图片识别（本地 python + wd-tagger，默认 false） +// - HentaiPython：python 可执行文件路径；留空自动在 PATH 里找 +// - HentaiThreshold：判定阈值 0~1（默认 0.5） type Config struct { Port int Dsn string @@ -32,6 +35,9 @@ type Config struct { NSFWEnabled bool NSFWPython string NSFWThreshold float64 + HentaiEnabled bool + HentaiPython string + HentaiThreshold float64 } // appConfig 当前生效的配置（程序运行期间一直不变）。 @@ -45,7 +51,7 @@ func GetAppConfig() Config { // LoadAppConfig 从 application.properties 读配置；文件不存在就用默认值。 func LoadAppConfig() { // 先给一套默认值 - appConfig = Config{Port: 8080, Dsn: "media-info.db", AutoOpen: true, RotateImages: true, Workers: 4, NSFWThreshold: 0.5} + appConfig = Config{Port: 8080, Dsn: "media-info.db", AutoOpen: true, RotateImages: true, Workers: 4, NSFWThreshold: 0.5, HentaiThreshold: 0.5} p, err := properties.LoadFile("application.properties", properties.UTF8) if err != nil { @@ -80,6 +86,16 @@ func LoadAppConfig() { appConfig.NSFWThreshold = t } } + // 二次元识别：默认关；python 路径留空自动在 PATH 找；阈值只认 0~1 之间的小数 + appConfig.HentaiEnabled = p.GetBool("hentai.enabled", false) + if v, ok := p.Get("hentai.python"); ok && v != "" { + appConfig.HentaiPython = v + } + if v, ok := p.Get("hentai.threshold"); ok && strings.TrimSpace(v) != "" { + if t, err := strconv.ParseFloat(strings.TrimSpace(v), 64); err == nil && t > 0 && t < 1 { + appConfig.HentaiThreshold = t + } + } } // WorkerCount 返回探测/MIME/移动阶段使用的并发 worker 数： diff --git a/internal/hentai/detector.go b/internal/hentai/detector.go new file mode 100644 index 0000000..ee7e607 --- /dev/null +++ b/internal/hentai/detector.go @@ -0,0 +1,359 @@ +// Package hentai 负责"二次元图片识别"：本机 Python + wd-tagger 起一个本地 HTTP 服务， +// Go 侧把图片路径发给它，拿到"是否二次元敏感"和分级信息。 +// +// 大白话理解：wd-tagger 是专门针对二次元训练的分级模型，能区分 safe/questionable/explicit。 +// 我们让 python 当"外援"： +// - 扫描开始时拉起一个常驻的 python 小服务（模型只加载一次）； +// - 探测图片时逐张把路径 POST 给它； +// - 扫描结束把服务停掉、清理临时脚本。 +// +// 机器上没有 python 或没装 onnxruntime 时，Start 会报错，调用方跳过识别、 +// 走普通整理流程——这个功能是纯附加的，不影响其它步骤。 +package hentai + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "time" + + _ "embed" +) + +// 内嵌的 python 服务脚本。 +//go:embed server.py +var serverScript []byte + +// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。 +// 首次启动若本地没有模型还要联网下载约 1.5GB，给足余量。 +const detectorStartTimeout = 600 * time.Second + +// requestTimeout 单张图片检测的超时时间。 +const requestTimeout = 30 * time.Second + +// Detector 封装一个本地 python 服务的客户端。 +type Detector struct { + url string // 本地服务地址，如 http://127.0.0.1:54322 + threshold float64 // 判定阈值 0~1 + client *http.Client + cmd *exec.Cmd // python 进程句柄（Stop 时杀掉） + scriptDir string // 临时脚本所在目录（Stop 时清理） + python string // python 可执行文件路径（服务重启时要用） + modelDir string // 模型目录 lib（服务重启时要用） + mu sync.Mutex // 保护 stopped / 进程句柄 + stopped bool // 已经 Stop 过了 + cancel context.CancelFunc + done chan struct{} // 进程退出时关闭（Wait 协程负责） + restartMu sync.Mutex // 串行化重启：同一时刻只一个 goroutine 重启，其余排队等 +} + +// DetectResult 二次元检测的结果。 +type DetectResult struct { + Hentai bool // 是否判定为二次元敏感（explicit 或 questionable 超过阈值） + Score float64 // 最高非 safe 分数（explicit 或 questionable 中较高的） + Rating string // 分级标签：safe / questionable / explicit + Ratings map[string]float64 // 三个分级的原始分数 + TopTags []TagScore // 高置信度标签（最多 10 个） +} + +// TagScore 一个标签及其分数。 +type TagScore struct { + Tag string `json:"tag"` + Score float64 `json:"score"` +} + +// Start 拉起本地 python 检测服务并等它就绪。 +// - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。 +// - threshold：判定阈值 0~1。 +// +// 找不到 python、装不了依赖、或服务起不来时返回错误（调用方跳过识别即可）。 +func Start(pythonPath string, threshold float64) (*Detector, error) { + if threshold <= 0 || threshold >= 1 { + threshold = 0.5 // 阈值传得不合理就用默认 + } + // 找 python + py := resolvePython(pythonPath) + if py == "" { + return nil, fmt.Errorf("未找到 python，二次元识别不可用") + } + + // 先把内嵌的 server.py 写到临时目录 + tmpDir, err := os.MkdirTemp("", "kmp-hentai-*") + if err != nil { + return nil, fmt.Errorf("创建临时目录失败: %v", err) + } + scriptPath := filepath.Join(tmpDir, "server.py") + if err := os.WriteFile(scriptPath, serverScript, 0644); err != nil { + _ = os.RemoveAll(tmpDir) + return nil, fmt.Errorf("写入检测脚本失败: %v", err) + } + + // 模型目录：程序运行目录下的 lib（和 application.properties / media-info.db 同一根目录）。 + modelDir := "" + if wd, err := os.Getwd(); err == nil { + modelDir = filepath.Join(wd, "lib") + } + + d := &Detector{ + threshold: threshold, + client: &http.Client{Timeout: requestTimeout}, + python: py, + modelDir: modelDir, + scriptDir: tmpDir, + } + + // 拉起服务（同样一套逻辑也用于崩溃后的重启） + if err := d.launch(); err != nil { + _ = os.RemoveAll(tmpDir) + return nil, err + } + return d, nil +} + +// launch 启动（或重启）python 检测服务，并等它就绪（模型加载完成）。 +func (d *Detector) launch() error { + // 先把上一次的进程收拾掉（首次启动时没有，直接跳过） + d.teardown() + + // 找一个空闲端口 + port, err := freePort() + if err != nil { + return fmt.Errorf("找空闲端口失败: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir) + // 标准输出/错误都接到程序自身的输出 + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stdout + if err := cmd.Start(); err != nil { + cancel() + return fmt.Errorf("启动二次元检测服务失败: %v", err) + } + + done := make(chan struct{}) + d.url = fmt.Sprintf("http://127.0.0.1:%d", port) + d.cmd = cmd + d.cancel = cancel + d.done = done + // 后台等进程退出 + go func() { + _ = cmd.Wait() + close(done) + }() + + // 等服务就绪（模型加载完成）。就绪前先探测 /health。 + if err := d.waitReady(); err != nil { + d.teardown() + return err + } + return nil +} + +// teardown 停掉当前 python 进程并清理句柄（不删临时目录）。 +func (d *Detector) teardown() { + if d.cancel != nil { + d.cancel() + d.cancel = nil + } + if d.cmd != nil && d.cmd.Process != nil { + _ = d.cmd.Process.Kill() + } + if d.done != nil { + select { + case <-d.done: + case <-time.After(3 * time.Second): + } + d.done = nil + } + d.cmd = nil +} + +// waitReady 轮询本地服务，直到能正常应答。 +func (d *Detector) waitReady() error { + deadline := time.Now().Add(detectorStartTimeout) + for time.Now().Before(deadline) { + // 服务进程提前退了直接报错 + select { + case <-d.done: + return fmt.Errorf("二次元检测服务提前退出，请确认已安装 onnxruntime 和 Pillow (pip install onnxruntime Pillow numpy)") + default: + } + req, err := http.NewRequest(http.MethodGet, d.url+"/health", nil) + if err != nil { + return err + } + resp, err := d.client.Do(req) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + } + time.Sleep(1 * time.Second) // 模型加载较慢，轮询间隔拉长 + } + return fmt.Errorf("二次元检测服务在 %v 内未就绪", detectorStartTimeout) +} + +// Detect 检测一张图片是否为二次元敏感内容。 +// 服务失联时会串行化地重启服务并重试一次。 +func (d *Detector) Detect(path string) (DetectResult, error) { + result, err := d.detectOnce(path) + if err == nil { + return result, nil + } + // 只有"发请求都失败"才走重启 + var ce *connError + if !errors.As(err, &ce) { + return DetectResult{}, err + } + + d.restartMu.Lock() + defer d.restartMu.Unlock() + + if d.isStopped() { + return DetectResult{}, err + } + // 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试 + if d.healthy() { + return d.detectOnce(path) + } + fmt.Printf("=== 二次元检测服务失联(%v), 正在重启…\n", ce.err) + if rerr := d.launch(); rerr != nil { + return DetectResult{}, fmt.Errorf("二次元检测服务重启失败: %v", rerr) + } + fmt.Println("=== 二次元检测服务已重启，继续识别") + return d.detectOnce(path) +} + +// healthy 快速探活：本地服务是否还能正常应答。 +func (d *Detector) healthy() bool { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.url+"/health", nil) + if err != nil { + return false + } + resp, err := d.client.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + return resp.StatusCode == http.StatusOK +} + +// isStopped 判断检测器是否已被 Stop（线程安全）。 +func (d *Detector) isStopped() bool { + d.mu.Lock() + defer d.mu.Unlock() + return d.stopped +} + +// detectOnce 发一次检测请求（不带重启逻辑）。 +func (d *Detector) detectOnce(path string) (DetectResult, error) { + body, _ := json.Marshal(map[string]interface{}{ + "path": path, + "threshold": d.threshold, + }) + resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body)) + if err != nil { + return DetectResult{}, &connError{err} + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err != nil { + return DetectResult{}, err + } + var out struct { + Ok bool `json:"ok"` + Error string `json:"error"` + Hentai bool `json:"hentai"` + Score float64 `json:"score"` + Rating string `json:"rating"` + Ratings map[string]float64 `json:"ratings"` + TopTags []TagScore `json:"topTags"` + } + if err := json.Unmarshal(data, &out); err != nil { + return DetectResult{}, fmt.Errorf("解析检测结果失败: %v", err) + } + if !out.Ok { + return DetectResult{}, fmt.Errorf("二次元检测服务返回错误: %s", out.Error) + } + return DetectResult{ + Hentai: out.Hentai, + Score: out.Score, + Rating: out.Rating, + Ratings: out.Ratings, + TopTags: out.TopTags, + }, nil +} + +// connError 标记"请求根本没发出去/连不上本地服务"的错误（服务失联）。 +type connError struct{ err error } + +func (e *connError) Error() string { return e.err.Error() } +func (e *connError) Unwrap() error { return e.err } + +// Stop 停掉检测服务并清理临时文件。可以安全地重复调用。 +func (d *Detector) Stop() { + d.mu.Lock() + defer d.mu.Unlock() + if d.stopped { + return + } + d.stopped = true + d.teardown() + if d.scriptDir != "" { + _ = os.RemoveAll(d.scriptDir) + } +} + +// Threshold 返回当前使用的判定阈值（0~1）。给上层记录结果用。 +func (d *Detector) Threshold() float64 { + return d.threshold +} + +// resolvePython 在配置值 / PATH 里找可用的 python 解释器。 +func resolvePython(configured string) string { + candidates := []string{} + if strings.TrimSpace(configured) != "" { + candidates = append(candidates, strings.TrimSpace(configured)) + } + if os.Getenv("KM_HENTAI_PYTHON") != "" { + candidates = append(candidates, os.Getenv("KM_HENTAI_PYTHON")) + } + candidates = append(candidates, "python", "python3", "py") + for _, c := range candidates { + if p, err := exec.LookPath(c); err == nil { + return p + } + // 配置给的是绝对路径：直接看文件在不在 + if strings.ContainsAny(c, `\/`) { + if _, err := os.Stat(c); err == nil { + return c + } + } + } + return "" +} + +// freePort 找一个当前空闲的 TCP 端口。 +func freePort() (int, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return 0, err + } + defer ln.Close() + return ln.Addr().(*net.TCPAddr).Port, nil +} diff --git a/internal/hentai/server.py b/internal/hentai/server.py new file mode 100644 index 0000000..cec90e3 --- /dev/null +++ b/internal/hentai/server.py @@ -0,0 +1,361 @@ +# -*- coding: utf-8 -*- +"""wd-tagger 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。 + +基于 SmilingWolf/wd-vit-large-tagger-v3 模型，专门用于二次元（anime/hentai）图片分级。 +输出 rating:safe / rating:questionable / rating:explicit 三个分级分数， +以及若干高置信度标签（用于辅助判断内容类型）。 + +只依赖 Python 标准库 + onnxruntime + Pillow。 +模型与标签文件放在 <程序运行目录>/lib/wd-tagger/ 下，首次启动自动从 HuggingFace 下载。 + +推理并发策略与 nsfw/server.py 相同：每个 worker 线程各持一份独立 session 实例， +避免 onnxruntime 同一实例并发推理崩溃的问题。 +""" +import concurrent.futures +import csv +import io +import json +import os +import shutil +import sys +import tempfile +import threading +import urllib.request +import zipfile +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +try: + import numpy as np +except ImportError: + print("=== 错误: 缺少 numpy，请执行 pip install numpy", flush=True) + sys.exit(1) + +try: + import onnxruntime as ort +except ImportError: + print("=== 错误: 缺少 onnxruntime，请执行 pip install onnxruntime", flush=True) + sys.exit(1) + +try: + from PIL import Image +except ImportError: + print("=== 错误: 缺少 Pillow，请执行 pip install Pillow", flush=True) + sys.exit(1) + +# 模型相关配置 +MODEL_DIR_NAME = "wd-tagger" +MODEL_FILENAME = "model.onnx" +TAGS_FILENAME = "selected_tags.csv" + +# HuggingFace 下载地址（wd-vit-large-tagger-v3） +HF_MODEL_URL = "https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3/resolve/main/model.onnx" +HF_TAGS_URL = "https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3/resolve/main/selected_tags.csv" + +# 模型输入尺寸 +INPUT_SIZE = 448 + +# ImageNet 归一化参数 +IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) + +# 我们关注的分级标签 +RATING_TAGS = { + "rating:safe": "safe", + "rating:questionable": "questionable", + "rating:explicit": "explicit", +} + +# 模型目录由 Go 侧作为第二个参数传入 +MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else "" + +# 推理并发数 +INFER_WORKERS = int(os.environ.get("KM_HENTAI_WORKERS", "4") or "4") +if INFER_WORKERS < 1: + INFER_WORKERS = 1 + +_executor = None +_thread_local = threading.local() + +# 标签表：启动时从 CSV 加载 +_tag_names = [] +_rating_indices = {} # tag名 -> 在输出数组中的索引 + + +def _looks_like_model(path): + """粗略校验文件是不是真 ONNX 模型。""" + try: + with open(path, "rb") as f: + head = f.read(4) + if not head or head.startswith(b"<"): + return False + return os.path.getsize(path) > 1_000_000 + except OSError: + return False + + +def _download_file(url, local): + """带超时地下载文件。""" + req = urllib.request.Request(url) + with urllib.request.urlopen(req, timeout=60) as resp, open(local, "wb") as f: + shutil.copyfileobj(resp, f) + + +def _ensure_model(model_dir): + """确保模型和标签文件就绪，返回 (model_path, tags_path) 或 (None, None) 表示用包自带。""" + if not model_dir: + return None, None + + wd_dir = os.path.join(model_dir, MODEL_DIR_NAME) + model_path = os.path.join(wd_dir, MODEL_FILENAME) + tags_path = os.path.join(wd_dir, TAGS_FILENAME) + + # 检查模型文件 + if os.path.isfile(model_path) and _looks_like_model(model_path): + print("=== 使用本地 wd-tagger 模型: %s" % model_path, flush=True) + else: + # 下载模型 + try: + os.makedirs(wd_dir, exist_ok=True) + print("=== 本地未找到 wd-tagger 模型, 正在下载: %s" % model_path, flush=True) + _download_file(HF_MODEL_URL, model_path) + if _looks_like_model(model_path): + print("=== wd-tagger 模型下载完成: %s" % model_path, flush=True) + else: + print("=== 下载的模型文件无效, 删除", flush=True) + try: + os.remove(model_path) + except OSError: + pass + return None, None + except Exception as exc: + print("=== wd-tagger 模型下载失败: %s" % exc, flush=True) + return None, None + + # 检查标签文件 + if os.path.isfile(tags_path): + print("=== 使用本地标签文件: %s" % tags_path, flush=True) + else: + try: + print("=== 下载标签文件: %s" % tags_path, flush=True) + _download_file(HF_TAGS_URL, tags_path) + print("=== 标签文件下载完成", flush=True) + except Exception as exc: + print("=== 标签文件下载失败: %s" % exc, flush=True) + return None, None + + return model_path, tags_path + + +def _load_tags(tags_path): + """加载标签 CSV，建立 tag名 -> 索引 的映射。""" + global _tag_names, _rating_indices + _tag_names = [] + _rating_indices = {} + + if not tags_path or not os.path.isfile(tags_path): + print("=== 警告: 标签文件不存在，无法解析输出", flush=True) + return + + with open(tags_path, "r", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + name = row.get("name", "") + _tag_names.append(name) + if name in RATING_TAGS: + _rating_indices[name] = len(_tag_names) - 1 + + print("=== 已加载 %d 个标签, 其中分级标签: %s" % (len(_tag_names), list(_rating_indices.keys())), flush=True) + + +def _new_session(model_path): + """创建一份独立的 ONNX 推理 session。""" + if model_path: + return ort.InferenceSession(model_path, providers=["CPUExecutionProvider"]) + # 没有本地模型，无法推理 + raise RuntimeError("wd-tagger 模型不可用") + + +def get_session(): + """确保模型已就绪，返回一个推理 session（启动时调用）。""" + global _session + if "_session" not in dir() or _session is None: + model_path, tags_path = _ensure_model(MODEL_DIR) + _load_tags(tags_path) + _session = _new_session(model_path) + return _session + + +_session = None + + +def _preprocess(img_bytes): + """把图片字节预处理成模型输入的 numpy 数组。 + + 步骤： + 1. 解码为 PIL Image + 2. 转为 RGB（如果有 alpha 通道） + 3. 缩放到 INPUT_SIZE x INPUT_SIZE（保持宽高比，用 padding 补齐） + 4. 归一化（ImageNet 参数） + 5. 转为 NHWC 格式的 float32 数组（模型期望 [batch, H, W, C]） + """ + img = Image.open(io.BytesIO(img_bytes)) + if img.mode != "RGB": + img = img.convert("RGB") + + # 保持宽高比缩放，短边补白色 + w, h = img.size + scale = INPUT_SIZE / max(w, h) + new_w = int(w * scale) + new_h = int(h * scale) + img = img.resize((new_w, new_h), Image.LANCZOS) + + # 创建白色背景，把缩放后的图贴到中间 + canvas = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE), (255, 255, 255)) + paste_x = (INPUT_SIZE - new_w) // 2 + paste_y = (INPUT_SIZE - new_h) // 2 + canvas.paste(img, (paste_x, paste_y)) + + # 转 numpy 数组并归一化 + arr = np.array(canvas, dtype=np.float32) / 255.0 + arr = (arr - IMAGENET_MEAN) / IMAGENET_STD + + # HWC -> NHWC (模型期望 [batch, height, width, channels]) + arr = np.expand_dims(arr, axis=0) + return arr + + +def _infer_bytes(session, img_bytes): + """用给定的 session 推理一张图片（字节），返回 (ratings_dict, top_tags_list)。""" + inp = _preprocess(img_bytes) + input_name = session.get_inputs()[0].name + outputs = session.run(None, {input_name: inp}) + + # 输出是 [1, num_tags] 的分数数组 + scores = outputs[0][0] if len(outputs[0].shape) == 2 else outputs[0] + + # 提取分级分数 + ratings = {} + for tag_name, idx in _rating_indices.items(): + if idx < len(scores): + rating_key = RATING_TAGS[tag_name] + ratings[rating_key] = float(scores[idx]) + + # 提取 top 标签（排除分级标签，取分数最高的 10 个） + top_tags = [] + non_rating_indices = [i for i in range(len(scores)) if i not in _rating_indices.values()] + if non_rating_indices: + top_indices = sorted(non_rating_indices, key=lambda i: scores[i], reverse=True)[:10] + for idx in top_indices: + if idx < len(_tag_names) and scores[idx] > 0.5: + top_tags.append({"tag": _tag_names[idx], "score": float(scores[idx])}) + + return ratings, top_tags + + +def _read_bytes(path): + """用 Python 读文件字节。""" + if os.name == "nt": + if path.startswith("\\\\?\\"): + lp = path + elif path.startswith("\\\\"): + lp = "\\\\?\\UNC\\" + path.lstrip("\\") + elif os.path.isabs(path): + lp = "\\\\?\\" + os.path.normpath(path) + else: + lp = path + else: + lp = path + + if not os.path.exists(lp): + raise OSError("文件不存在: %s" % path) + with open(lp, "rb") as f: + data = f.read() + if not data: + raise OSError("空文件: %s" % path) + return data + + +def _infer(path): + """整个推理任务：读文件 + 预处理 + 推理。""" + data = _read_bytes(path) + session = getattr(_thread_local, "session", None) + if session is None: + # 每个 worker 线程各持一份独立 session + model_path = os.path.join(MODEL_DIR, MODEL_DIR_NAME, MODEL_FILENAME) if MODEL_DIR else "" + session = _new_session(model_path) + _thread_local.session = session + return _infer_bytes(session, data) + + +class Handler(BaseHTTPRequestHandler): + def _send(self, obj, status=200): + data = json.dumps(obj, ensure_ascii=False).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def do_GET(self): + if self.path == "/health": + self._send({"ok": True}) + else: + self._send({"ok": False, "error": "not found"}, 404) + + def do_POST(self): + if self.path != "/detect": + return self._send({"ok": False, "error": "not found"}, 404) + length = int(self.headers.get("Content-Length", 0)) + try: + body = json.loads(self.rfile.read(length) or b"{}") + path = body.get("path", "") + threshold = float(body.get("threshold", 0.5)) + if not path: + return self._send({"ok": False, "error": "缺少图片路径"}) + # 丢给推理线程池 + ratings, top_tags = _executor.submit(_infer, path).result() + except Exception as exc: + return self._send({"ok": False, "error": str(exc)}) + + # 判断是否为 hentai：explicit 分数 >= 阈值 或 questionable 分数 >= 阈值 + explicit_score = ratings.get("explicit", 0.0) + questionable_score = ratings.get("questionable", 0.0) + safe_score = ratings.get("safe", 0.0) + + # hentai 判定：explicit 或 questionable 分数超过阈值 + is_hentai = (explicit_score >= threshold) or (questionable_score >= threshold) + + # 取最高的非 safe 分数作为"敏感分" + max_score = max(explicit_score, questionable_score) + + # 确定分级标签 + rating = "safe" + if explicit_score >= questionable_score and explicit_score >= threshold: + rating = "explicit" + elif questionable_score >= threshold: + rating = "questionable" + + self._send({ + "ok": True, + "hentai": is_hentai, + "score": max_score, + "rating": rating, + "ratings": ratings, + "topTags": top_tags, + }) + + def log_message(self, *args): + pass + + +class Server(ThreadingHTTPServer): + request_queue_size = 128 + + +if __name__ == "__main__": + port = int(sys.argv[1]) if len(sys.argv) > 1 else 0 + # 启动前先把模型准备好 + get_session() + # 固定大小推理线程池 + _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS) + Server(("127.0.0.1", port), Handler).serve_forever() diff --git a/internal/nsfw/detector.go b/internal/nsfw/detector.go index f1eb0f8..e4f0510 100644 --- a/internal/nsfw/detector.go +++ b/internal/nsfw/detector.go @@ -92,7 +92,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) { } // 模型目录：程序运行目录下的 lib（和 application.properties / media-info.db 同一根目录）。 - // 服务脚本优先加载 lib/320n.onnx；本地没有会自己下载到那里，下载失败再回退用 pip 自带模型。 + // 服务脚本优先加载 lib/nwfs/320n.onnx；本地没有会自己下载到那里，下载失败再回退用 pip 自带模型。 modelDir := "" if wd, err := os.Getwd(); err == nil { modelDir = filepath.Join(wd, "lib") diff --git a/internal/nsfw/server.py b/internal/nsfw/server.py index c322e98..c4472ec 100644 --- a/internal/nsfw/server.py +++ b/internal/nsfw/server.py @@ -6,9 +6,9 @@ onnxruntime 的 InferenceSession 不是线程安全的，同一个实例并发 所以这里用一个固定大小的推理线程池：每个 worker 线程各持一份独立模型实例， N 个请求可以真·并行推理。收到 /detect 就把图片喂给 NudeDetector，返回结果 JSON。 -模型加载优先级（模型放在 <程序运行目录>/lib/320n.onnx）： +模型加载优先级（模型放在 <程序运行目录>/lib/nwfs/320n.onnx）： 1. 本地已有该文件 → 直接加载（离线可用，适合随程序分发）； - 2. 本地没有 → 自动下载到 lib/320n.onnx 再加载（仅第一次需要网络）； + 2. 本地没有 → 自动下载到 lib/nwfs/320n.onnx 再加载（仅第一次需要网络）； 3. 下载失败 → 回退用 nudenet 包自带的模型。 下载地址是 NudeNet 官方 release 的 320n.onnx（约 11.6MB）。 @@ -46,7 +46,7 @@ SENSITIVE = { } # 模型文件与官方下载地址 -MODEL_FILENAME = "320n.onnx" +MODEL_FILENAME = "nwfs/320n.onnx" MODEL_URL = "https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/320n.onnx" # 模型目录由 Go 侧作为第二个参数传入（程序运行目录下的 lib）；没传则回退用包自带模型 @@ -265,4 +265,4 @@ if __name__ == "__main__": get_detector() # 固定大小推理线程池：每个 worker 线程持一份独立模型实例，可并行推理 _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS) - Server(("127.0.0.1", port), Handler).serve_forever() \ No newline at end of file + Server(("127.0.0.1", port), Handler).serve_forever() diff --git a/internal/organize/organize.go b/internal/organize/organize.go index 8168e0f..7f7d0b5 100644 --- a/internal/organize/organize.go +++ b/internal/organize/organize.go @@ -12,6 +12,7 @@ package organize import ( "context" + "encoding/json" "fmt" "os" "path/filepath" @@ -23,6 +24,7 @@ import ( "kenaito-media-picker/internal/classify" "kenaito-media-picker/internal/config" "kenaito-media-picker/internal/digest" + "kenaito-media-picker/internal/hentai" "kenaito-media-picker/internal/mover" "kenaito-media-picker/internal/nsfw" "kenaito-media-picker/internal/probe" @@ -71,22 +73,27 @@ type Organizer struct { mover *mover.Mover // 负责实际搬文件的家伙 ctx context.Context // 中止信号；ctx 被取消时探测/归类立即停下 nsfw *nsfw.Detector // 色情检测器；为 nil 表示没启用，图片不检测 + hentai *hentai.Detector // 二次元检测器；为 nil 表示没启用 rotateImages bool // 是否处理 EXIF 旋转（转正/按显示方向归类）；false 则完全不动，按存储方向归类 buckets map[string][]string // 目标目录 → 里面该放哪些文件（只记源路径） order []string // 目标目录的先后顺序（保证输出稳定、不来回跳） onProgress func(phase string, done, total int) pending []*store.MediaInfo // 待写入数据库的媒体记录（攒一批再批量写） pendingNsfw []store.NsfwRecord // 待写入数据库的色情识别记录（攒一批再批量写） + pendingHentai []store.HentaiRecord // 待写入数据库的二次元识别记录（攒一批再批量写） resume map[string]store.ScanLog // 续传表：源路径 → 上次归类结果（重扫跳过重新探测） resumeSkip int32 // 本次续传跳过的文件数（并行累计） pendingScan []store.ScanLog // 待写入续传表的记录（攒一批再批量写） nsfwChecked int32 // 色情识别已检查的图片数（并行累计，只涨不跌） nsfwHits int32 // 色情识别命中的图片数（并行累计） + hentaiChecked int32 // 二次元识别已检查的图片数（并行累计） + hentaiHits int32 // 二次元识别命中的图片数（并行累计） // 阶段耗时统计（并行 worker 累计，单位纳秒，只用于日志诊断瓶颈） tImageProbe int64 // 图片尺寸探测总耗时 tImgMd5 int64 // 图片 MD5/指纹计算总耗时 tBright int64 // 明暗分析总耗时 tNsfw int64 // 色情识别总耗时 + tHentai int64 // 二次元识别总耗时 tMediaProbe int64 // 媒体探测(ffprobe)总耗时 tMediaMd5 int64 // 媒体(视频/音频) MD5/指纹计算总耗时 } @@ -129,6 +136,12 @@ func (o *Organizer) SetNSFWDetector(d *nsfw.Detector) { o.nsfw = d } +// SetHentaiDetector 设置二次元检测器。传 nil 表示不检测（默认）。 +// 设置了检测器后，图片探测时会额外做二次元识别，命中的图片会归入「图片_二次元_*」目录。 +func (o *Organizer) SetHentaiDetector(d *hentai.Detector) { + o.hentai = d +} + // SetRotateImages 设置是否处理图片 EXIF 旋转（默认 true）。 // - true：带旋转标记的图片做转正处理（有 ffmpeg 就烧进像素，没有则按显示方向归类）； // - false：完全不动，图片按文件里存储的方向/尺寸归类，不做任何旋转相关处理。 @@ -392,6 +405,10 @@ type imageProbeResult struct { lightScore float64 // 亮色像素占比 nsfwScore float64 // 色情检测最高分（未检测为 0） nsfw bool // 是否判定为色情（检测可用时才可能为 true） + hentaiScore float64 // 二次元检测最高分（未检测为 0） + hentai bool // 是否判定为二次元敏感 + hentaiRating string // 二次元分级标签：safe / questionable / explicit + hentaiTags string // 高置信度标签（JSON 数组字符串） } // probeImage 给一张图片做体检（并行阶段执行）：探测尺寸 + 算 MD5 + 分析明暗。 @@ -467,6 +484,33 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult { } } } + // 二次元识别（GIF 不做；没启用检测器就直接跳过） + if o.hentai != nil && !strings.EqualFold(f.Ext, ".gif") { + hentaiStart := time.Now() + result, derr := o.hentai.Detect(f.Path) + atomic.AddInt64(&o.tHentai, int64(time.Since(hentaiStart))) + if derr != nil { + fmt.Printf("=== 二次元检测失败(按未命中归类): %s, 错误: %v \n", f.Path, derr) + } else { + r.hentaiScore = result.Score + r.hentai = result.Hentai + r.hentaiRating = result.Rating + // 把标签序列化成 JSON 字符串存起来 + if len(result.TopTags) > 0 { + if tagsJSON, err := json.Marshal(result.TopTags); err == nil { + r.hentaiTags = string(tagsJSON) + } + } + checked := atomic.AddInt32(&o.hentaiChecked, 1) + if result.Hentai { + atomic.AddInt32(&o.hentaiHits, 1) + fmt.Printf("=== 二次元命中: %s, 分级=%s, 检测分=%.2f (阈值=%.2f)\n", f.Path, result.Rating, result.Score, o.hentai.Threshold()) + } + if checked%nsfwProgressInterval == 0 { + fmt.Printf("=== 二次元识别进度: 已检查 %d 张, 命中 %d 张\n", checked, atomic.LoadInt32(&o.hentaiHits)) + } + } + } return r } @@ -511,12 +555,16 @@ func (o *Organizer) correctImageRotation(f scanner.File, res *probe.ProbeResult) func (o *Organizer) ProcessImages(images []scanner.File) { // 统计结果，最后打一行总结给用户看 success, fail, ignore, dup := 0, 0, 0, 0 - nsfwCount := 0 // 命中的色情图片数 + nsfwCount := 0 // 命中的色情图片数 + hentaiCount := 0 // 命中的二次元图片数 total := len(images) phaseStart := time.Now() if o.nsfw != nil { fmt.Printf("=== 色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval) } + if o.hentai != nil { + fmt.Printf("=== 二次元识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval) + } results := runParallel(o.ctx, images, o.probeImage, func(done int) { o.report(PhaseImages, done, total) }) @@ -552,6 +600,7 @@ func (o *Organizer) ProcessImages(images []scanner.File) { default: // 正常图片：按分辨率/方向/明暗归入对应的分类目录； // 被判为色情的改去「图片_敏感_*」目录，明暗子目录照常保留 + // 被判为二次元敏感的改去「图片_二次元_*」目录 o.record(f, "image", r.res, r.md5Str, r.headMd5, r.tailMd5, r.size) if r.nsfw { o.addClassified(f, classify.ImageDirSensitive(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5) @@ -563,6 +612,18 @@ func (o *Organizer) ProcessImages(images []scanner.File) { Threshold: o.nsfw.Threshold(), }) nsfwCount++ + } else if r.hentai { + o.addClassified(f, classify.ImageDirHentai(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5) + // 攒一条二次元识别记录 + o.pendingHentai = append(o.pendingHentai, store.HentaiRecord{ + FilePath: f.Path, + FileName: filepath.Base(f.Path), + Score: r.hentaiScore, + Rating: r.hentaiRating, + Threshold: o.hentai.Threshold(), + TopTags: r.hentaiTags, + }) + hentaiCount++ } else { o.addClassified(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5) } @@ -575,15 +636,19 @@ func (o *Organizer) ProcessImages(images []scanner.File) { fmt.Printf("=== 续传跳过: %d 个文件(上次已归类，直接复用)\n", skip) } fmt.Printf("=== 图片处理完毕，成功: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", success, fail, ignore, dup) - fmt.Printf("=== 图片归类耗时: 实际 %.1fs, 其中(并行累计) 探测 %.1fs, MD5 %.1fs, 明暗 %.1fs, 色情 %.1fs\n", + fmt.Printf("=== 图片归类耗时: 实际 %.1fs, 其中(并行累计) 探测 %.1fs, MD5 %.1fs, 明暗 %.1fs, 色情 %.1fs, 二次元 %.1fs\n", time.Since(phaseStart).Seconds(), time.Duration(atomic.LoadInt64(&o.tImageProbe)).Seconds(), time.Duration(atomic.LoadInt64(&o.tImgMd5)).Seconds(), time.Duration(atomic.LoadInt64(&o.tBright)).Seconds(), - time.Duration(atomic.LoadInt64(&o.tNsfw)).Seconds()) + time.Duration(atomic.LoadInt64(&o.tNsfw)).Seconds(), + time.Duration(atomic.LoadInt64(&o.tHentai)).Seconds()) if o.nsfw != nil && nsfwCount > 0 { fmt.Printf("=== 色情识别: 共命中 %d 张，已归入「图片_敏感_*」目录，可在【色情识别】页面查看并人工复核\n", nsfwCount) } + if o.hentai != nil && hentaiCount > 0 { + fmt.Printf("=== 二次元识别: 共命中 %d 张，已归入「图片_二次元_*」目录，可在【二次元识别】页面查看并人工复核\n", hentaiCount) + } } // ============================ 视频/音频整理 ============================ @@ -756,4 +821,23 @@ func (o *Organizer) Flush() { } o.pendingNsfw = nil } + + // 把二次元识别记录补上"整理后位置"（预览就用它），再批量写进数据库 + if len(o.pendingHentai) > 0 { + records := make([]*store.HentaiRecord, 0, len(o.pendingHentai)) + for i := range o.pendingHentai { + if target, ok := moved[o.pendingHentai[i].FilePath]; ok { + o.pendingHentai[i].MovedPath = target + } else { + o.pendingHentai[i].MovedPath = o.pendingHentai[i].FilePath // 没搬走就用源路径 + } + records = append(records, &o.pendingHentai[i]) + } + if o.store != nil { + if err := o.store.SaveHentaiBatch(records); err != nil { + fmt.Printf("=== 保存二次元识别记录失败: %v\n", err) + } + } + o.pendingHentai = nil + } } diff --git a/internal/service/scan.go b/internal/service/scan.go index ecff299..a86c4f4 100644 --- a/internal/service/scan.go +++ b/internal/service/scan.go @@ -24,6 +24,7 @@ import ( "time" "kenaito-media-picker/internal/config" + "kenaito-media-picker/internal/hentai" "kenaito-media-picker/internal/nsfw" "kenaito-media-picker/internal/organize" "kenaito-media-picker/internal/probe" @@ -76,6 +77,9 @@ const ConfigKeyDestDir = "scan.destDir" // 配置开了但检测器没起来（没装 python/模型下载失败等）也能区分出来，前端据此提示。 const nsfwDetectorReadyKey = "nsfw.detectorReady" +// hentaiDetectorReadyKey 记录最近一次扫描时二次元检测服务是否真的起来了。 +const hentaiDetectorReadyKey = "hentai.detectorReady" + // setNsfwDetectorReady 把"色情检测服务是否真的起来了"记进数据库（AppConfig）。 func (s *ScanService) setNsfwDetectorReady(ready bool) { if s.store == nil { @@ -88,6 +92,18 @@ func (s *ScanService) setNsfwDetectorReady(ready bool) { _ = s.store.SaveConfig(nsfwDetectorReadyKey, v) } +// setHentaiDetectorReady 把"二次元检测服务是否真的起来了"记进数据库（AppConfig）。 +func (s *ScanService) setHentaiDetectorReady(ready bool) { + if s.store == nil { + return + } + v := "0" + if ready { + v = "1" + } + _ = s.store.SaveConfig(hentaiDetectorReadyKey, v) +} + // DefaultPickerDirName 用户没指定存放目录时，默认在扫描目录下新建这个文件夹装整理好的文件。 const DefaultPickerDirName = "KMP已整理" @@ -533,8 +549,31 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro } } + // ---------- 第六步（可选）：二次元识别 ---------- + // 用本机 python + wd-tagger 起一个本地检测服务。没装依赖时跳过， + // 不影响图片整理——检测结果只决定走普通目录还是「图片_二次元」目录。 + hentaiDet := (*hentai.Detector)(nil) + s.setHentaiDetectorReady(false) + if cfg := config.GetAppConfig(); cfg.HentaiEnabled { + if d, err := hentai.Start(cfg.HentaiPython, cfg.HentaiThreshold); err != nil { + fmt.Printf("=== 二次元识别未启用(图片照常处理): %v\n", err) + } else { + hentaiDet = d + defer hentaiDet.Stop() + s.setHentaiDetectorReady(true) + fmt.Println("=== 二次元识别已启用(wd-tagger)") + // 清掉上一次的二次元识别记录，保证「二次元识别」页面显示的是本次扫描的结果 + if s.store != nil { + if err := s.store.ClearHentai(); err != nil { + fmt.Printf("=== 清空旧二次元识别记录失败: %v\n", err) + } + } + } + } + org := organize.New(pickerDir, ffprobePath, ffmpegPath, s.store, ctx) org.SetNSFWDetector(nsfwDet) + org.SetHentaiDetector(hentaiDet) org.SetRotateImages(config.GetAppConfig().RotateImages) // 整个"整理"分两个阶段，界面上共用一条进度条： diff --git a/internal/store/store.go b/internal/store/store.go index a3bcb88..9f2289d 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -75,6 +75,20 @@ type NsfwRecord struct { CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 检测时间 } +// HentaiRecord 一条"二次元图片识别"的记录：哪张图、分级结果、标签、现在在哪。 +// 前端「二次元识别」页面查这张表展示检测结果，方便人工复核。 +type HentaiRecord struct { + Id int64 `xorm:"pk autoincr 'id'" json:"id"` // 自增主键 + FilePath string `xorm:"text 'file_path'" json:"filePath"` // 检测时的源路径 + FileName string `xorm:"varchar(255) 'file_name'" json:"fileName"` // 文件名（含扩展名） + Score float64 `xorm:"real 'score'" json:"score"` // 最高非 safe 分数 0~1 + Rating string `xorm:"varchar(32) 'rating'" json:"rating"` // 分级标签：safe / questionable / explicit + Threshold float64 `xorm:"real 'threshold'" json:"threshold"` // 判定阈值 + TopTags string `xorm:"text 'top_tags'" json:"topTags"` // 高置信度标签（JSON 数组） + MovedPath string `xorm:"text 'moved_path'" json:"movedPath"` // 整理后实际位置（预览用） + CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 检测时间 +} + // ScanLog 一条"续传"记录：某个源文件上次被归类到了哪个目录。 // 用途：扫描中断/出错后重跑时，已经归类过的文件直接复用上次的目标目录， // 不再重新探测尺寸 / 算 MD5 / 分析明暗 / 色情识别，省下几小时的重扫时间。 @@ -106,8 +120,8 @@ func Open(dbPath string) (*Store, error) { _ = e.Close() return nil, err } - // 自动创建/补全四张表（媒体记录表、配置表、色情识别记录表、续传表） - if err := e.Sync(new(MediaInfo), new(AppConfig), new(NsfwRecord), new(ScanLog)); err != nil { + // 自动创建/补全五张表（媒体记录表、配置表、色情识别记录表、二次元识别记录表、续传表） + if err := e.Sync(new(MediaInfo), new(AppConfig), new(NsfwRecord), new(HentaiRecord), new(ScanLog)); err != nil { _ = e.Close() return nil, err } @@ -195,6 +209,67 @@ func (s *Store) ClearNsfw() error { return err } +// ============================ 二次元识别记录 ============================ + +// SaveHentaiBatch 批量保存二次元识别记录（一次扫描攒一批，最后统一写入）。 +func (s *Store) SaveHentaiBatch(records []*HentaiRecord) error { + if s == nil || s.engine == nil || len(records) == 0 { + return nil + } + const chunkSize = 100 + sess := s.engine.NewSession() + defer sess.Close() + if err := sess.Begin(); err != nil { + return err + } + for i := 0; i < len(records); i += chunkSize { + end := i + chunkSize + if end > len(records) { + end = len(records) + } + if _, err := sess.Insert(records[i:end]); err != nil { + _ = sess.Rollback() + return err + } + } + return sess.Commit() +} + +// ClearHentai 清空二次元识别记录。每次开启二次元识别的扫描开始时调用。 +func (s *Store) ClearHentai() error { + if s == nil || s.engine == nil { + return nil + } + _, err := s.engine.Where("1 = 1").Delete(new(HentaiRecord)) + return err +} + +// ListHentai 查全部二次元识别记录，按检测分从高到低排。 +func (s *Store) ListHentai() ([]HentaiRecord, error) { + if s == nil || s.engine == nil { + return nil, nil + } + list := make([]HentaiRecord, 0) + err := s.engine.Desc("score").Find(&list) + return list, err +} + +// GetHentaiById 按 ID 查一条二次元识别记录。 +func (s *Store) GetHentaiById(id int64) (*HentaiRecord, error) { + if s == nil || s.engine == nil { + return nil, nil + } + record := new(HentaiRecord) + has, err := s.engine.ID(id).Get(record) + if err != nil { + return nil, err + } + if !has { + return nil, nil + } + return record, nil +} + // ============================ 续传表 ============================ // LoadScanLog 把全部续传记录读出来，按源路径建索引，供扫描时快速判断"这个文件上次归到哪了"。 diff --git a/lib/320n.onnx b/lib/nwfs/320n.onnx similarity index 100% rename from lib/320n.onnx rename to lib/nwfs/320n.onnx diff --git a/x20260829.md b/x20260829.md index dbecb81..0dfc256 100644 --- a/x20260829.md +++ b/x20260829.md @@ -19,7 +19,7 @@ |------|------|------| | 真人 vs 动漫 | `dghs-imgutils` `anime_real` | 二分类，输出 real/anime 概率 | | 真人 vs 风景 | 真人检测器 | 检测到人脸/人物 → 真人，否则风景 | -| 真人鉴黄 | NudeNet（已有 `lib/320n.onnx`） | 保持现状 | +| 真人鉴黄 | NudeNet（已有 `lib/nwfs/320n.onnx`） | 保持现状 | | 动漫鉴黄 | `dghs-imgutils` `anime_rating` / `nsfw_pred` | danbooru 评级：safe/questionable/explicit 分数 | - `dghs-imgutils`（`pip install dghs-imgutils`）一个包覆盖「真人/动漫判定」+「动漫鉴黄」+「动漫分类」，