odboy/kenaito-media-picker
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

upd

Browse Source
This commit is contained in:
odboy
2025-08-29 21:16:19 +08:00
parent 61e7ffc03b
commit 79156aa8ae
4 changed files with 119 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/nsfw/detector.go
+37 -17
View File
@@ -39,6 +39,13 @@ var serverScript []byte
// 注意：哪些敏感类才算色情，定义在 server.py 的 SENSITIVE 集合里（单一事实来源），
// Go 侧只负责把路径发给服务、读取判定结果，不重复维护这份清单。
// NsfwDetection NudeNet 检测到的一个目标（类别+置信度+位置）。
type NsfwDetection struct {
Class string `json:"class"`
Score float64 `json:"score"`
Box [4]int `json:"box"` // [x, y, w, h]
}
// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。
// 首次启动若本地没有模型还要联网下载约 11.6MB，给足余量。
const detectorStartTimeout = 180 * time.Second
@@ -210,28 +217,29 @@ func (d *Detector) waitReady() error {
// Detect 检测一张图片是否色情。返回：
// - score：命中的敏感类里最高分数（没有命中为 0）
// - nsfw：是否判定为色情（有敏感类且分数 >= 阈值）
// - detections：所有检测到的目标（类别+分数+位置）
// - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
//
// 服务失联（连接被拒，多半是 python 进程在处理某张图时崩了）时，会串行化地
// 重启服务并重试一次：多个并发 worker 里只有一个负责重启，其余排队，等重启完
// 也重试自己的图——保证服务短暂挂掉不会漏掉一整批图片。
func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
score, nsfw, err = d.detectOnce(path)
func (d *Detector) Detect(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
score, nsfw, detections, err = d.detectOnce(path)
if err == nil {
return score, nsfw, nil
return score, nsfw, detections, nil
}
// 只有"发请求都失败"（连接被拒/重置，说明服务失联）才走重启；
// 服务正常返回的明确错误（如某张图读不了）不重启。
var ce *connError
if !errors.As(err, &ce) {
return score, nsfw, err
return score, nsfw, detections, err
}
d.restartMu.Lock()
defer d.restartMu.Unlock()
if d.isStopped() {
return 0, false, err
return 0, false, nil, err
}
// 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试，别反复重启
if d.healthy() {
@@ -239,7 +247,7 @@ func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
}
fmt.Printf("=== 色情检测服务失联(%v), 正在重启…\n", ce.err)
if rerr := d.launch(); rerr != nil {
return 0, false, fmt.Errorf("检测服务重启失败: %v", rerr)
return 0, false, nil, fmt.Errorf("检测服务重启失败: %v", rerr)
}
fmt.Println("=== 色情检测服务已重启，继续识别")
return d.detectOnce(path)
@@ -269,7 +277,7 @@ func (d *Detector) isStopped() bool {
}
// detectOnce 发一次检测请求（不带重启逻辑）。
func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error) {
func (d *Detector) detectOnce(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
body, _ := json.Marshal(map[string]interface{}{
"path": path,
"threshold": d.threshold,
@@ -277,12 +285,12 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error)
resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
if err != nil {
// 打上"连接失败"标记，供 Detect 判断是否要重启
return 0, false, &connError{err}
return 0, false, nil, &connError{err}
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
return 0, false, err
return 0, false, nil, err
}
var out struct {
Ok bool `json:"ok"`
@@ -292,24 +300,36 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error)
Detections []map[string]interface{} `json:"detections"`
}
if err := json.Unmarshal(data, &out); err != nil {
return 0, false, fmt.Errorf("解析检测结果失败: %v", err)
return 0, false, nil, fmt.Errorf("解析检测结果失败: %v", err)
}
if !out.Ok {
return 0, false, fmt.Errorf("检测服务返回错误: %s", out.Error)
return 0, false, nil, fmt.Errorf("检测服务返回错误: %s", out.Error)
}
// 解析检测详情
detections = make([]NsfwDetection, 0, len(out.Detections))
for _, det := range out.Detections {
cls, _ := det["class"].(string)
sc, _ := det["score"].(float64)
boxArr, _ := det["box"].([]interface{})
var box [4]int
for i := 0; i < 4 && i < len(boxArr); i++ {
if v, ok := boxArr[i].(float64); ok {
box[i] = int(v)
}
}
detections = append(detections, NsfwDetection{Class: cls, Score: sc, Box: box})
}
// 打印所有检测到的类别和分数（调试用）
if len(out.Detections) > 0 {
if len(detections) > 0 {
var classes []string
for _, det := range out.Detections {
cls, _ := det["class"].(string)
sc, _ := det["score"].(float64)
classes = append(classes, fmt.Sprintf("%s:%.3f", cls, sc))
for _, det := range detections {
classes = append(classes, fmt.Sprintf("%s:%.3f", det.Class, det.Score))
}
fmt.Printf("=== NSFW 检测详情 %s: [%s]\n", path, strings.Join(classes, ", "))
} else {
fmt.Printf("=== NSFW 检测详情 %s: [无检测结果]\n", path)
}
return out.Score, out.Nsfw, nil
return out.Score, out.Nsfw, detections, nil
}
// connError 标记"请求根本没发出去/连不上本地服务"的错误（服务失联）。
internal/organize/organize.go
+44 -37
View File
@@ -66,26 +66,26 @@ const scanLogFlushInterval = 200
//
// 这样设计是为了避免"边分类边搬"导致目录反复开关、文件搬来搬去，效率更高。
type Organizer struct {
pickerDir string // 整理后的文件都放到这个根目录下
ffprobePath string // ffprobe 程序路径；为空表示没有可用解码器
ffmpegPath string // ffmpeg 程序路径；为空表示没有转正 EXIF 旋转图片的能力
store *store.Store // 数据库；为 nil 表示不记录媒体信息
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 // 待写入数据库的媒体记录（攒一批再批量写）
resume map[string]store.ScanLog // 续传表：源路径 → 上次归类结果（重扫跳过重新探测）
resumeSkip int32 // 本次续传跳过的文件数（并行累计）
pendingScan []store.ScanLog // 待写入续传表的记录（攒一批再批量写）
nsfwChecked int32 // 色情识别已检查的图片数（并行累计，只涨不跌）
nsfwHits int32 // 色情识别命中的图片数（并行累计）
animeNsfwChecked int32 // 动漫色情识别已检查的图片数
animeNsfwHits int32 // 动漫色情识别命中的图片数
pickerDir string // 整理后的文件都放到这个根目录下
ffprobePath string // ffprobe 程序路径；为空表示没有可用解码器
ffmpegPath string // ffmpeg 程序路径；为空表示没有转正 EXIF 旋转图片的能力
store *store.Store // 数据库；为 nil 表示不记录媒体信息
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 // 待写入数据库的媒体记录（攒一批再批量写）
resume map[string]store.ScanLog // 续传表：源路径 → 上次归类结果（重扫跳过重新探测）
resumeSkip int32 // 本次续传跳过的文件数（并行累计）
pendingScan []store.ScanLog // 待写入续传表的记录（攒一批再批量写）
nsfwChecked int32 // 色情识别已检查的图片数（并行累计，只涨不跌）
nsfwHits int32 // 色情识别命中的图片数（并行累计）
animeNsfwChecked int32 // 动漫色情识别已检查的图片数
animeNsfwHits int32 // 动漫色情识别命中的图片数
// 阶段耗时统计（并行 worker 累计，单位纳秒，只用于日志诊断瓶颈）
tImageProbe int64 // 图片尺寸探测总耗时
tImgMd5 int64 // 图片 MD5/指纹计算总耗时
@@ -392,22 +392,23 @@ func (o *Organizer) flushRecords() {
// 后面的归类（判断重复、分目录）在串行阶段根据这个结论来定。
type imageProbeResult struct {
// 下面几个是"体检结论"，最多只有一个为 true：
psd bool // 是 PSD 设计稿，不做处理原样收走
unsupported bool // 不支持的格式，归入"已忽略"
probeErr bool // 体检失败（打不开/读不出），归入"读取异常"
resume bool // 续传命中：上次已归类过，跳过全部体检，直接复用目标目录
resumeDir string // 续传命中的目标目录
res *probe.ProbeResult // 体检出的尺寸等信息（体检成功时才有）
md5Str string // 文件 MD5
headMd5 string // 文件头采样 MD5（大文件指纹）
tailMd5 string // 文件尾采样 MD5
size int64 // 文件大小
darkScore float64 // 暗色像素占比（用于亮度分类）
lightScore float64 // 亮色像素占比
nsfwScore float64 // 色情检测最高分（未检测为 0）
nsfw bool // 是否判定为色情（检测可用时才可能为 true）
animeNsfwScore float64 // 动漫色情检测最高分（未检测为 0）
animeNsfw bool // 是否判定为动漫色情
psd bool // 是 PSD 设计稿，不做处理原样收走
unsupported bool // 不支持的格式，归入"已忽略"
probeErr bool // 体检失败（打不开/读不出），归入"读取异常"
resume bool // 续传命中：上次已归类过，跳过全部体检，直接复用目标目录
resumeDir string // 续传命中的目标目录
res *probe.ProbeResult // 体检出的尺寸等信息（体检成功时才有）
md5Str string // 文件 MD5
headMd5 string // 文件头采样 MD5（大文件指纹）
tailMd5 string // 文件尾采样 MD5
size int64 // 文件大小
darkScore float64 // 暗色像素占比（用于亮度分类）
lightScore float64 // 亮色像素占比
nsfwScore float64 // 色情检测最高分（未检测为 0）
nsfw bool // 是否判定为色情（检测可用时才可能为 true）
nsfwDetections []nsfw.NsfwDetection // NudeNet 检测详情
animeNsfwScore float64 // 动漫色情检测最高分（未检测为 0）
animeNsfw bool // 是否判定为动漫色情
animeNsfwRatings anime_nsfw.AnimeRatings // WD Tagger 四个 rating 的原始分数
}
@@ -464,7 +465,7 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
// 色情识别（GIF 不做；没启用检测器就直接跳过）
if o.nsfw != nil && !strings.EqualFold(f.Ext, ".gif") {
nsfwStart := time.Now()
score, hit, derr := o.nsfw.Detect(f.Path)
score, hit, detections, derr := o.nsfw.Detect(f.Path)
atomic.AddInt64(&o.tNsfw, int64(time.Since(nsfwStart)))
if derr != nil {
// 检测失败按"未命中"处理，不影响整理流程
@@ -472,6 +473,7 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
} else {
r.nsfwScore = score
r.nsfw = hit
r.nsfwDetections = detections
checked := atomic.AddInt32(&o.nsfwChecked, 1)
if hit {
atomic.AddInt32(&o.nsfwHits, 1)
@@ -601,6 +603,11 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
if info != nil && r.nsfwScore > 0 {
info.NsfwScore = r.nsfwScore
info.NsfwThreshold = o.nsfw.Threshold()
if len(r.nsfwDetections) > 0 {
if detJSON, err := json.Marshal(r.nsfwDetections); err == nil {
info.NsfwDetections = string(detJSON)
}
}
}
if info != nil && r.animeNsfwScore > 0 {
info.AnimeNsfwScore = r.animeNsfwScore
internal/store/store.go
+24 -23
View File
@@ -31,30 +31,31 @@ func init() {
// MediaInfo 一条媒体文件的记录：文件本身的信息——路径、名字、大小、类型、宽高、编码、格式等。
type MediaInfo struct {
Id int64 `xorm:"pk autoincr 'id'" json:"id"`
FileMd5 string `xorm:"char(32) index 'file_md5'" json:"fileMd5"` // 文件 MD5，内容相同则相同，用来判断是不是同一个文件
FilePath string `xorm:"text 'file_path'" json:"filePath"` // 扫描时的源路径（在哪发现的）
FileName string `xorm:"varchar(255) 'file_name'" json:"fileName"` // 文件名（含扩展名）
FileExt string `xorm:"varchar(32) 'file_ext'" json:"fileExt"` // 扩展名（如 .jpg）
FileSize int64 `xorm:"bigint 'file_size'" json:"fileSize"` // 文件大小（字节）
MimeType string `xorm:"varchar(128) 'mime_type'" json:"mimeType"` // MIME 类型（如 image/jpeg）
MediaType string `xorm:"varchar(16) index 'media_type'" json:"mediaType"` // 大类型：image / video / audio
Width int `xorm:"int index 'width'" json:"width"` // 宽（像素）
Height int `xorm:"int index 'height'" json:"height"` // 高（像素）
Duration float64 `xorm:"real 'duration'" json:"duration"` // 时长（秒）
Codec string `xorm:"varchar(64) 'codec'" json:"codec"` // 编码，如 h264 / hevc
Format string `xorm:"varchar(64) 'format'" json:"format"` // 容器格式，如 mp4 / mkv
RawInfo string `xorm:"text 'raw_info'" json:"rawInfo"` // ffprobe 读出的原始信息（JSON），留底备用
HeadMd5 string `xorm:"char(32) index('idx_fp') 'head_md5'" json:"headMd5"` // 文件头 256KB 的 MD5（指纹用）
TailMd5 string `xorm:"char(32) index('idx_fp') 'tail_md5'" json:"tailMd5"` // 文件尾 256KB 的 MD5（指纹用）
NsfwScore float64 `xorm:"real 'nsfw_score'" json:"nsfwScore"` // NSFW 检测分 0~1（0 表示未检测或非图片）
NsfwThreshold float64 `xorm:"real 'nsfw_threshold'" json:"nsfwThreshold"` // NSFW 判定阈值
AnimeNsfwScore float64 `xorm:"real 'anime_nsfw_score'" json:"animeNsfwScore"` // 动漫 NSFW 检测分 0~1（0 表示未检测或非动漫）
Id int64 `xorm:"pk autoincr 'id'" json:"id"`
FileMd5 string `xorm:"char(32) index 'file_md5'" json:"fileMd5"` // 文件 MD5，内容相同则相同，用来判断是不是同一个文件
FilePath string `xorm:"text 'file_path'" json:"filePath"` // 扫描时的源路径（在哪发现的）
FileName string `xorm:"varchar(255) 'file_name'" json:"fileName"` // 文件名（含扩展名）
FileExt string `xorm:"varchar(32) 'file_ext'" json:"fileExt"` // 扩展名（如 .jpg）
FileSize int64 `xorm:"bigint 'file_size'" json:"fileSize"` // 文件大小（字节）
MimeType string `xorm:"varchar(128) 'mime_type'" json:"mimeType"` // MIME 类型（如 image/jpeg）
MediaType string `xorm:"varchar(16) index 'media_type'" json:"mediaType"` // 大类型：image / video / audio
Width int `xorm:"int index 'width'" json:"width"` // 宽（像素）
Height int `xorm:"int index 'height'" json:"height"` // 高（像素）
Duration float64 `xorm:"real 'duration'" json:"duration"` // 时长（秒）
Codec string `xorm:"varchar(64) 'codec'" json:"codec"` // 编码，如 h264 / hevc
Format string `xorm:"varchar(64) 'format'" json:"format"` // 容器格式，如 mp4 / mkv
RawInfo string `xorm:"text 'raw_info'" json:"rawInfo"` // ffprobe 读出的原始信息（JSON），留底备用
HeadMd5 string `xorm:"char(32) index('idx_fp') 'head_md5'" json:"headMd5"` // 文件头 256KB 的 MD5（指纹用）
TailMd5 string `xorm:"char(32) index('idx_fp') 'tail_md5'" json:"tailMd5"` // 文件尾 256KB 的 MD5（指纹用）
NsfwScore float64 `xorm:"real 'nsfw_score'" json:"nsfwScore"` // NSFW 检测分 0~1（0 表示未检测或非图片）
NsfwThreshold float64 `xorm:"real 'nsfw_threshold'" json:"nsfwThreshold"` // NSFW 判定阈值
NsfwDetections string `xorm:"text 'nsfw_detections'" json:"nsfwDetections"` // NudeNet 检测详情 JSON（类别+分数+位置）
AnimeNsfwScore float64 `xorm:"real 'anime_nsfw_score'" json:"animeNsfwScore"` // 动漫 NSFW 检测分 0~1（0 表示未检测或非动漫）
AnimeNsfwThreshold float64 `xorm:"real 'anime_nsfw_threshold'" json:"animeNsfwThreshold"` // 动漫 NSFW 判定阈值
AnimeNsfwRatings string `xorm:"text 'anime_nsfw_ratings'" json:"animeNsfwRatings"` // WD Tagger 四个 rating 的 JSON（G/S/Q/E）
ManualNsfw string `xorm:"varchar(1) 'manual_nsfw'" json:"manualNsfw"` // 手动标记是否色情：Y/N，默认 N
CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 首次入库时间
UpdatedAt time.Time `xorm:"updated 'updated_at'" json:"updatedAt"` // 最近更新时间
AnimeNsfwRatings string `xorm:"text 'anime_nsfw_ratings'" json:"animeNsfwRatings"` // WD Tagger 四个 rating 的 JSON（G/S/Q/E）
ManualNsfw string `xorm:"varchar(1) 'manual_nsfw'" json:"manualNsfw"` // 手动标记是否色情：Y/N，默认 N
CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 首次入库时间
UpdatedAt time.Time `xorm:"updated 'updated_at'" json:"updatedAt"` // 最近更新时间
}
// Store 数据库访问的封装。整个程序共用一个实例，通过它读写数据库。