feat: 添加动漫色情检测(WD Tagger v3)，媒体库展示优化
- 新增 anime_nsfw 模块：WD Tagger v3 模型检测动漫/二次元图片 - 修复误判：当 general 评分最高时判定为安全，避免非动漫图误判 - 数据库存储四个 rating 原始分数(JSON)，无需重扫即可查看 - 媒体库新增动漫色情检测列、文件路径列(支持一键复制) - 修复预览图加载占位符显示两个点的问题 - 移动文件按钮不再限制扫描完成后才显示
This commit is contained in:
@@ -12,6 +12,7 @@ package organize
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -20,6 +21,7 @@ import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"kenaito-media-picker/internal/anime_nsfw"
|
||||
"kenaito-media-picker/internal/classify"
|
||||
"kenaito-media-picker/internal/config"
|
||||
"kenaito-media-picker/internal/digest"
|
||||
@@ -71,21 +73,25 @@ type Organizer struct {
|
||||
mover *mover.Mover // 负责实际搬文件的家伙
|
||||
ctx context.Context // 中止信号；ctx 被取消时探测/归类立即停下
|
||||
nsfw *nsfw.Detector // 色情检测器；为 nil 表示没启用，图片不检测
|
||||
animeNsfw *anime_nsfw.Detector // 动漫色情检测器；为 nil 表示没启用
|
||||
rotateImages bool // 是否处理 EXIF 旋转（转正/按显示方向归类）；false 则完全不动，按存储方向归类
|
||||
buckets map[string][]string // 目标目录 → 里面该放哪些文件（只记源路径）
|
||||
order []string // 目标目录的先后顺序（保证输出稳定、不来回跳）
|
||||
onProgress func(phase string, done, total int)
|
||||
pending []*store.MediaInfo // 待写入数据库的媒体记录（攒一批再批量写）
|
||||
pending []*store.MediaInfo // 待写入数据库的媒体记录（攒一批再批量写）
|
||||
resume map[string]store.ScanLog // 续传表：源路径 → 上次归类结果（重扫跳过重新探测）
|
||||
resumeSkip int32 // 本次续传跳过的文件数（并行累计）
|
||||
pendingScan []store.ScanLog // 待写入续传表的记录（攒一批再批量写）
|
||||
nsfwChecked int32 // 色情识别已检查的图片数（并行累计，只涨不跌）
|
||||
nsfwHits int32 // 色情识别命中的图片数（并行累计）
|
||||
nsfwChecked int32 // 色情识别已检查的图片数（并行累计，只涨不跌）
|
||||
nsfwHits int32 // 色情识别命中的图片数（并行累计）
|
||||
animeNsfwChecked int32 // 动漫色情识别已检查的图片数
|
||||
animeNsfwHits int32 // 动漫色情识别命中的图片数
|
||||
// 阶段耗时统计（并行 worker 累计，单位纳秒，只用于日志诊断瓶颈）
|
||||
tImageProbe int64 // 图片尺寸探测总耗时
|
||||
tImgMd5 int64 // 图片 MD5/指纹计算总耗时
|
||||
tBright int64 // 明暗分析总耗时
|
||||
tNsfw int64 // 色情识别总耗时
|
||||
tAnimeNsfw int64 // 动漫色情识别总耗时
|
||||
tMediaProbe int64 // 媒体探测(ffprobe)总耗时
|
||||
tMediaMd5 int64 // 媒体(视频/音频) MD5/指纹计算总耗时
|
||||
}
|
||||
@@ -128,6 +134,12 @@ func (o *Organizer) SetNSFWDetector(d *nsfw.Detector) {
|
||||
o.nsfw = d
|
||||
}
|
||||
|
||||
// SetAnimeNSFWDetector 设置动漫色情检测器。传 nil 表示不检测（默认）。
|
||||
// 设置了检测器后，图片探测时会额外做动漫色情识别，命中的图片同样归入「图片_敏感_*」目录。
|
||||
func (o *Organizer) SetAnimeNSFWDetector(d *anime_nsfw.Detector) {
|
||||
o.animeNsfw = d
|
||||
}
|
||||
|
||||
// SetRotateImages 设置是否处理图片 EXIF 旋转（默认 true）。
|
||||
// - true：带旋转标记的图片做转正处理（有 ffmpeg 就烧进像素，没有则按显示方向归类）；
|
||||
// - false：完全不动，图片按文件里存储的方向/尺寸归类，不做任何旋转相关处理。
|
||||
@@ -394,6 +406,9 @@ type imageProbeResult struct {
|
||||
lightScore float64 // 亮色像素占比
|
||||
nsfwScore float64 // 色情检测最高分（未检测为 0）
|
||||
nsfw bool // 是否判定为色情（检测可用时才可能为 true）
|
||||
animeNsfwScore float64 // 动漫色情检测最高分（未检测为 0）
|
||||
animeNsfw bool // 是否判定为动漫色情
|
||||
animeNsfwRatings anime_nsfw.AnimeRatings // WD Tagger 四个 rating 的原始分数
|
||||
}
|
||||
|
||||
// probeImage 给一张图片做体检（并行阶段执行）：探测尺寸 + 算 MD5 + 分析明暗。
|
||||
@@ -470,6 +485,29 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
|
||||
}
|
||||
}
|
||||
}
|
||||
// 动漫色情识别（GIF 不做；没启用检测器就直接跳过）
|
||||
if o.animeNsfw != nil && !strings.EqualFold(f.Ext, ".gif") {
|
||||
animeStart := time.Now()
|
||||
score, hit, ratings, derr := o.animeNsfw.Detect(f.Path)
|
||||
atomic.AddInt64(&o.tAnimeNsfw, int64(time.Since(animeStart)))
|
||||
if derr != nil {
|
||||
fmt.Printf("=== 动漫色情检测失败(按未命中归类): %s, 错误: %v \n", f.Path, derr)
|
||||
} else {
|
||||
r.animeNsfwScore = score
|
||||
r.animeNsfw = hit
|
||||
r.animeNsfwRatings = ratings
|
||||
checked := atomic.AddInt32(&o.animeNsfwChecked, 1)
|
||||
if hit {
|
||||
atomic.AddInt32(&o.animeNsfwHits, 1)
|
||||
fmt.Printf("=== 动漫色情命中: %s, 检测分=%.4f (阈值=%.4f)\n", f.Path, score, o.animeNsfw.Threshold())
|
||||
} else {
|
||||
fmt.Printf("=== 动漫色情未命中: %s, 检测分=%.4f (阈值=%.4f)\n", f.Path, score, o.animeNsfw.Threshold())
|
||||
}
|
||||
if checked%nsfwProgressInterval == 0 {
|
||||
fmt.Printf("=== 动漫色情识别进度: 已检查 %d 张, 命中 %d 张\n", checked, atomic.LoadInt32(&o.animeNsfwHits))
|
||||
}
|
||||
}
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -514,12 +552,15 @@ 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 // 命中的色情图片数
|
||||
total := len(images)
|
||||
phaseStart := time.Now()
|
||||
if o.nsfw != nil {
|
||||
fmt.Printf("=== 色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval)
|
||||
}
|
||||
if o.animeNsfw != nil {
|
||||
fmt.Printf("=== 动漫色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval)
|
||||
}
|
||||
results := runParallel(o.ctx, images, o.probeImage, func(done int) {
|
||||
o.report(PhaseImages, done, total)
|
||||
})
|
||||
@@ -561,9 +602,19 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
info.NsfwScore = r.nsfwScore
|
||||
info.NsfwThreshold = o.nsfw.Threshold()
|
||||
}
|
||||
if r.nsfw {
|
||||
if info != nil && r.animeNsfwScore > 0 {
|
||||
info.AnimeNsfwScore = r.animeNsfwScore
|
||||
info.AnimeNsfwThreshold = o.animeNsfw.Threshold()
|
||||
if ratingsJSON, err := json.Marshal(r.animeNsfwRatings); err == nil {
|
||||
info.AnimeNsfwRatings = string(ratingsJSON)
|
||||
}
|
||||
}
|
||||
// 真人色情或动漫色情任一命中，都归入敏感目录
|
||||
if r.nsfw || r.animeNsfw {
|
||||
o.addClassified(f, classify.ImageDirSensitive(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5)
|
||||
nsfwCount++
|
||||
if r.nsfw {
|
||||
nsfwCount++
|
||||
}
|
||||
} else {
|
||||
o.addClassified(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5)
|
||||
}
|
||||
@@ -576,14 +627,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.tAnimeNsfw)).Seconds())
|
||||
animeNsfwHits := atomic.LoadInt32(&o.animeNsfwHits)
|
||||
if o.nsfw != nil && nsfwCount > 0 {
|
||||
fmt.Printf("=== 色情识别: 共命中 %d 张，已归入「图片_敏感_*」目录，可在【色情识别】页面查看并人工复核\n", nsfwCount)
|
||||
fmt.Printf("=== 色情识别: 共命中 %d 张，已归入「图片_敏感_*」目录，可在【媒体库】页面查看并通过<色情标记>人工复核\n", nsfwCount)
|
||||
}
|
||||
if o.animeNsfw != nil && animeNsfwHits > 0 {
|
||||
fmt.Printf("=== 动漫色情识别: 共命中 %d 张，已归入「图片_敏感_*」目录，可在【媒体库】页面查看并通过<色情标记>人工复核\n", animeNsfwHits)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user