feat: 二次元识别功能 + 模型目录整理(nwfs/wd-tagger 子目录)，修复 wd-tagger 输入维度(NHWC)
This commit is contained in:
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user