780 lines
32 KiB
Go
780 lines
32 KiB
Go
// Package organize 是"整理"这一步的总指挥，负责把一批文件分门别类、搬到该去的地方。
|
||
//
|
||
// 整理一辆快递车的过程可以这样理解：
|
||
// 1. 探测（并行）：每个文件先做"体检"——是图片还是视频？多大？什么分辨率？会不会是重复的？
|
||
// 2. 归类（串行）：根据体检结果，决定这个文件该进哪个分类目录，或者直接判为重复/异常/忽略
|
||
// 3. 移动（并行）：把每个分类目录攒齐的文件，批量搬进对应的文件夹，并把"搬去哪了"记进数据库
|
||
//
|
||
// 为什么探测要并行、归类要串行？
|
||
// - 探测很慢（读文件、算 MD5、调 ffprobe），多个一起干能明显提速；
|
||
// - 归类要看"前面有没有见过同样的文件"（去重判断），必须按顺序一个个来，否则会乱。
|
||
package organize
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"kenaito-media-picker/internal/anime_nsfw"
|
||
"kenaito-media-picker/internal/classify"
|
||
"kenaito-media-picker/internal/config"
|
||
"kenaito-media-picker/internal/digest"
|
||
"kenaito-media-picker/internal/mover"
|
||
"kenaito-media-picker/internal/nsfw"
|
||
"kenaito-media-picker/internal/probe"
|
||
"kenaito-media-picker/internal/scanner"
|
||
"kenaito-media-picker/internal/store"
|
||
)
|
||
|
||
// ============================ 阶段与目录命名 ============================
|
||
|
||
// 整理内部的两个阶段名，上报进度时用，供上层（service）区分显示。
|
||
const (
|
||
PhaseImages = "images" // 正在处理图片
|
||
PhaseMedia = "media" // 正在处理视频/音频
|
||
PhaseMove = "move" // 正在移动文件
|
||
)
|
||
|
||
// 特殊归类目录的名字（都是大白话，用户一眼能看懂）。
|
||
const (
|
||
dirImagePSD = "图片_PSD" // 没做任何处理，原样收走的 PSD 设计稿
|
||
dirImageErr = "图片_读取异常" // 打不开、读不出信息的图片
|
||
dirImageIgnored = "图片_已忽略" // 不支持的图片格式，原样收走但不做处理
|
||
dirVideoErr = "视频_读取异常" // 打不开、读不出信息的视频
|
||
dirVideoIgnored = "视频_已忽略" // 不支持的视频格式，原样收走
|
||
dirAudio = "音频" // 音频统一放这一个文件夹
|
||
)
|
||
|
||
// nsfwProgressInterval 色情识别每隔多少张图打一行进度，让前端能看到识别在推进
|
||
// （逐张打太吵，日志还会被 trim 掉一半）。
|
||
const nsfwProgressInterval = 200
|
||
|
||
// scanLogFlushInterval 续传记录攒够多少条就落一次库（中断/出错时尽量少丢已归类的结果）。
|
||
const scanLogFlushInterval = 200
|
||
|
||
// ============================ 整理器主体 ============================
|
||
|
||
// Organizer 整理器。它的工作方式：
|
||
// - 先把文件按"目标目录"分堆攒起来（buckets），
|
||
// - 攒完后统一创建目录、批量搬文件。
|
||
//
|
||
// 这样设计是为了避免"边分类边搬"导致目录反复开关、文件搬来搬去，效率更高。
|
||
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 // 色情识别已检查的图片数（并行累计，只涨不跌）
|
||
animeNsfwChecked int32 // 动漫色情识别已检查的图片数
|
||
// 阶段耗时统计（并行 worker 累计，单位纳秒，只用于日志诊断瓶颈）
|
||
tImageProbe int64 // 图片尺寸探测总耗时
|
||
tImgMd5 int64 // 图片 MD5/指纹计算总耗时
|
||
tBright int64 // 明暗分析总耗时
|
||
tNsfw int64 // 色情识别总耗时
|
||
tAnimeNsfw int64 // 动漫色情识别总耗时
|
||
tMediaProbe int64 // 媒体探测(ffprobe)总耗时
|
||
tMediaMd5 int64 // 媒体(视频/音频) MD5/指纹计算总耗时
|
||
}
|
||
|
||
// New 创建整理器。store 可以为 nil——数据库打不开时，整理照常进行，只是不记入库信息。
|
||
// ctx 可以为 nil——不提供中止能力，任务跑到完为止。
|
||
func New(pickerDir, ffprobePath, ffmpegPath string, st *store.Store, ctx context.Context) *Organizer {
|
||
org := &Organizer{
|
||
pickerDir: pickerDir,
|
||
ffprobePath: ffprobePath,
|
||
ffmpegPath: ffmpegPath,
|
||
store: st,
|
||
mover: &mover.Mover{},
|
||
ctx: ctx,
|
||
rotateImages: true, // 默认处理 EXIF 旋转；关掉用 SetRotateImages(false)
|
||
buckets: make(map[string][]string),
|
||
}
|
||
// 加载续传表：上次归类过的文件，这次跳过重新探测，直接复用目标目录
|
||
if st != nil {
|
||
if m, err := st.LoadScanLog(); err != nil {
|
||
fmt.Printf("=== 加载续传记录失败: %v\n", err)
|
||
} else if len(m) > 0 {
|
||
org.resume = m
|
||
fmt.Printf("=== 续传: 上次已归类 %d 个文件，本次将跳过它们的重新探测/识别\n", len(m))
|
||
}
|
||
}
|
||
return org
|
||
}
|
||
|
||
// SetProgress 设置进度回调。回调参数：
|
||
// - phase：当前在干哪一步（PhaseImages / PhaseMedia / PhaseMove）
|
||
// - done / total：这一步已完成数和总数
|
||
func (o *Organizer) SetProgress(fn func(phase string, done, total int)) {
|
||
o.onProgress = fn
|
||
}
|
||
|
||
// SetNSFWDetector 设置色情检测器。传 nil 表示不检测（默认）。
|
||
// 设置了检测器后，图片探测时会额外做色情识别，命中的图片会归入「图片_敏感_*」目录。
|
||
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：完全不动，图片按文件里存储的方向/尺寸归类，不做任何旋转相关处理。
|
||
func (o *Organizer) SetRotateImages(b bool) {
|
||
o.rotateImages = b
|
||
}
|
||
|
||
// report 上报进度。total 为 0 表示这一步没活干，就不用报了。
|
||
func (o *Organizer) report(phase string, done, total int) {
|
||
if o.onProgress != nil && total > 0 {
|
||
o.onProgress(phase, done, total)
|
||
}
|
||
}
|
||
|
||
// ============================ 并行小工具 ============================
|
||
|
||
// workerCount 并行探测的工人数：默认 4 个（配置 organize.workers 可调，0 表示用满 CPU）。
|
||
// 多核机器并行探测能明显提速；但并发太大会吃满 CPU，给个配置让用户自行权衡。
|
||
func workerCount() int {
|
||
return config.WorkerCount()
|
||
}
|
||
|
||
// runParallel 用固定数量的工人并行处理一批东西，结果按原来的顺序放回。
|
||
//
|
||
// 通俗理解：一条流水线有 N 个工人，每个工人从传送带上拿一个活干完，
|
||
// 干完把结果按编号放回架子上（results[i] 对应 items[i] 的结果）。
|
||
// onProgress 每干完一个活就报一次数（已完成几个），线程安全、只会涨不会退。
|
||
// ctx 被取消时工人不再干活（任务照常消费，避免发送方堵死），
|
||
// 被跳过的位置留零值，调用方看到 ctx.Err() 会自行停下。
|
||
func runParallel[T any, R any](ctx context.Context, items []T, fn func(T) R, onProgress func(done int)) []R {
|
||
results := make([]R, len(items))
|
||
if len(items) == 0 {
|
||
return results
|
||
}
|
||
workers := workerCount()
|
||
if workers > len(items) {
|
||
workers = len(items)
|
||
}
|
||
var wg sync.WaitGroup
|
||
var doneCount int32
|
||
jobs := make(chan int)
|
||
for w := 0; w < workers; w++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
for i := range jobs {
|
||
// 用户点了中止：不再干活，但任务还要继续消费，免得主 goroutine 堵死在发送上
|
||
if ctx != nil && ctx.Err() != nil {
|
||
continue
|
||
}
|
||
results[i] = fn(items[i])
|
||
if onProgress != nil {
|
||
onProgress(int(atomic.AddInt32(&doneCount, 1)))
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
for i := range items {
|
||
jobs <- i
|
||
}
|
||
close(jobs)
|
||
wg.Wait()
|
||
return results
|
||
}
|
||
|
||
// ============================ 分类与入库小工具 ============================
|
||
|
||
// add 把一个文件记到某个目标目录（相对分类名）的名下（先攒着，最后统一搬）。
|
||
func (o *Organizer) add(f scanner.File, subDir string) {
|
||
o.addToDir(f, filepath.Join(o.pickerDir, subDir))
|
||
}
|
||
|
||
// addToDir 把一个文件记到某个目标目录（完整路径）的名下。
|
||
func (o *Organizer) addToDir(f scanner.File, dir string) {
|
||
if _, ok := o.buckets[dir]; !ok {
|
||
o.order = append(o.order, dir)
|
||
}
|
||
o.buckets[dir] = append(o.buckets[dir], f.Path)
|
||
}
|
||
|
||
// addClassified 归档一个文件并记进续传表：把"源路径 → 目标目录"写进 scan_log，
|
||
// 这样下次重扫时这个文件直接复用目标目录，不再重新探测/识别。
|
||
func (o *Organizer) addClassified(f scanner.File, subDir string, size int64, headMd5, tailMd5 string) {
|
||
o.add(f, subDir)
|
||
if o.store == nil {
|
||
return
|
||
}
|
||
o.pendingScan = append(o.pendingScan, store.ScanLog{
|
||
SourcePath: f.Path,
|
||
TargetDir: filepath.Join(o.pickerDir, subDir),
|
||
FileSize: size,
|
||
HeadMd5: headMd5,
|
||
TailMd5: tailMd5,
|
||
})
|
||
if len(o.pendingScan) >= scanLogFlushInterval {
|
||
o.flushScanLog()
|
||
}
|
||
}
|
||
|
||
// matchResume 判断一个文件上次是否已归类过，且内容没变。
|
||
// 命中返回上次的目标目录和 ok=true；内容变了/没记录则 ok=false（要重新探测）。
|
||
func (o *Organizer) matchResume(f scanner.File) (store.ScanLog, bool) {
|
||
if o.resume == nil {
|
||
return store.ScanLog{}, false
|
||
}
|
||
sl, ok := o.resume[f.Path]
|
||
if !ok {
|
||
return sl, false
|
||
}
|
||
st, err := os.Stat(f.Path)
|
||
if err != nil {
|
||
return sl, false
|
||
}
|
||
if st.Size() != sl.FileSize {
|
||
return sl, false
|
||
}
|
||
// 大文件再核对头尾指纹，确认内容真没变才放心跳过
|
||
if sl.FileSize > digest.FingerprintThreshold && sl.HeadMd5 != "" && sl.TailMd5 != "" {
|
||
head, tail, ferr := digest.CalcFileFingerprint(f.Path, st.Size())
|
||
if ferr != nil || head != sl.HeadMd5 || tail != sl.TailMd5 {
|
||
return sl, false
|
||
}
|
||
}
|
||
return sl, true
|
||
}
|
||
|
||
// flushScanLog 把攒下的续传记录一次性写进数据库（中断时尽量少丢已归类的结果）。
|
||
func (o *Organizer) flushScanLog() {
|
||
if o.store == nil || len(o.pendingScan) == 0 {
|
||
return
|
||
}
|
||
if err := o.store.SaveScanLogs(o.pendingScan); err != nil {
|
||
fmt.Printf("=== 保存续传记录失败: %v\n", err)
|
||
}
|
||
o.pendingScan = nil
|
||
}
|
||
|
||
// calcFileInfo 读取文件大小并算 MD5（文件的"身份证号"）。
|
||
//
|
||
// 大文件算全量 MD5 很费时间，所以先用"指纹"（文件大小 + 头尾各 256KB 的 MD5）
|
||
// 去数据库里找找看有没有同款——有的话直接借它的 MD5 用，省得再全量算一遍。
|
||
// store 为 nil 时直接返回空，因为没有数据库可查，也不需要去重/入库。
|
||
func (o *Organizer) calcFileInfo(f scanner.File) (md5Str, headMd5, tailMd5 string, size int64) {
|
||
if o.store == nil {
|
||
return "", "", "", 0
|
||
}
|
||
st, err := os.Stat(f.Path)
|
||
if err != nil {
|
||
fmt.Printf("=== 读取文件信息失败: %s, 错误: %v\n", f.Path, err)
|
||
return "", "", "", 0
|
||
}
|
||
size = st.Size()
|
||
|
||
if size > digest.FingerprintThreshold {
|
||
if head, tail, ferr := digest.CalcFileFingerprint(f.Path, size); ferr == nil {
|
||
headMd5, tailMd5 = head, tail
|
||
if existing, ok := o.store.FindMediaByFingerprint(size, head, tail); ok && existing.FileMd5 != "" {
|
||
md5Str = existing.FileMd5 // 指纹命中，复用已入库的 MD5，跳过全量计算
|
||
}
|
||
}
|
||
}
|
||
if md5Str == "" {
|
||
if md5Str, err = digest.CalcFileMd5(f.Path); err != nil {
|
||
fmt.Printf("=== 计算MD5失败: %s, 错误: %v\n", f.Path, err)
|
||
return "", "", "", size
|
||
}
|
||
}
|
||
return md5Str, headMd5, tailMd5, size
|
||
}
|
||
|
||
// isDuplicate 判断一个文件是不是"同一个文件又出现了一次"。
|
||
//
|
||
// 判据：数据库里已经有一条记录，**MD5 相同、源路径也相同**。
|
||
// 只有这种情况才认为是"同一个文件被重复扫到"，直接忽略：不动它、不记录、不搬走。
|
||
//
|
||
// 特别注意：MD5 一样但**路径不一样**，那是另一个文件（碰巧内容一样），
|
||
// 不算重复，照样要搬进媒体库。
|
||
func (o *Organizer) isDuplicate(md5Str, path string) bool {
|
||
if md5Str == "" {
|
||
return false
|
||
}
|
||
if o.store == nil {
|
||
return false
|
||
}
|
||
existing, ok := o.store.FindMediaByMd5(md5Str)
|
||
if !ok {
|
||
return false
|
||
}
|
||
return existing.FilePath == path
|
||
}
|
||
|
||
// record 把一份媒体信息记到"待入库"缓冲里（先攒着，最后统一批量写库）。
|
||
// 返回创建的 MediaInfo 指针，调用方可继续设置额外字段（如 NsfwScore）。
|
||
func (o *Organizer) record(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) *store.MediaInfo {
|
||
if o.store == nil || md5Str == "" {
|
||
return nil
|
||
}
|
||
info := buildMediaInfo(f, mediaType, res, md5Str, headMd5, tailMd5, size)
|
||
o.pending = append(o.pending, info)
|
||
return info
|
||
}
|
||
|
||
// buildMediaInfo 把探测结果整理成一条要入库的记录。
|
||
func buildMediaInfo(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) *store.MediaInfo {
|
||
// 视频记视频编码，音频记音频编码
|
||
codec := res.VideoCodec
|
||
if mediaType == "audio" {
|
||
codec = res.AudioCodec
|
||
}
|
||
// 格式名：ffprobe 没读出来就用扩展名顶替（去掉前面的点）
|
||
format := res.Format
|
||
if format == "" && f.Ext != "" {
|
||
format = strings.TrimPrefix(f.Ext, ".")
|
||
}
|
||
return &store.MediaInfo{
|
||
FileMd5: md5Str,
|
||
FilePath: f.Path,
|
||
FileName: f.Name,
|
||
FileExt: f.Ext,
|
||
FileSize: size,
|
||
MimeType: f.Mime,
|
||
MediaType: mediaType,
|
||
Width: res.Width,
|
||
Height: res.Height,
|
||
Duration: res.Duration,
|
||
Codec: codec,
|
||
Format: format,
|
||
RawInfo: res.RawJSON,
|
||
HeadMd5: headMd5,
|
||
TailMd5: tailMd5,
|
||
}
|
||
}
|
||
|
||
// flushRecords 把缓冲的媒体记录批量写进数据库（一次事务，快得多）。
|
||
func (o *Organizer) flushRecords() {
|
||
if o.store == nil {
|
||
return
|
||
}
|
||
if len(o.pending) > 0 {
|
||
if err := o.store.SaveMediaBatch(o.pending); err != nil {
|
||
fmt.Printf("=== 批量保存媒体信息失败: %v\n", err)
|
||
}
|
||
o.pending = nil
|
||
}
|
||
}
|
||
|
||
// ============================ 图片整理 ============================
|
||
|
||
// imageProbeResult 是"给一张图片做完体检"的结论，体检在并行阶段完成，
|
||
// 后面的归类（判断重复、分目录）在串行阶段根据这个结论来定。
|
||
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）
|
||
nsfwDetections []nsfw.NsfwDetection // NudeNet 检测详情
|
||
animeNsfwScore float64 // 动漫色情检测最高分（未检测为 0）
|
||
animeNsfwRatings anime_nsfw.AnimeRatings // WD Tagger 四个 rating 的原始分数
|
||
}
|
||
|
||
// probeImage 给一张图片做体检（并行阶段执行）：探测尺寸 + 算 MD5 + 分析明暗。
|
||
// 注意：这里不能碰 o.buckets 之类的共享可变状态，因为多个工人同时跑；
|
||
// 只读数据库（查指纹），这是安全的。
|
||
func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
|
||
r := imageProbeResult{}
|
||
// 续传：上次已归类过、内容没变 → 跳过全部体检，直接复用上次的目标目录
|
||
if sl, ok := o.matchResume(f); ok {
|
||
atomic.AddInt32(&o.resumeSkip, 1)
|
||
r.resume = true
|
||
r.resumeDir = sl.TargetDir
|
||
return r
|
||
}
|
||
// PSD/PSB 设计稿不体检，直接标记"原样收走"
|
||
if strings.EqualFold(f.Ext, ".psd") || strings.EqualFold(f.Ext, ".psb") {
|
||
r.psd = true
|
||
return r
|
||
}
|
||
// 格式不认识，又没有 ffprobe 兜底 → 标记"忽略"
|
||
if !probe.IsSupportImage(f.Ext) && o.ffprobePath == "" {
|
||
r.unsupported = true
|
||
return r
|
||
}
|
||
// 探测尺寸
|
||
probeStart := time.Now()
|
||
res, err := probe.ProbeImage(f.Path, o.ffprobePath)
|
||
atomic.AddInt64(&o.tImageProbe, int64(time.Since(probeStart)))
|
||
if err != nil {
|
||
r.probeErr = true
|
||
return r
|
||
}
|
||
r.res = res
|
||
// 先处理 EXIF 旋转：带旋转标记的图，要么用 ffmpeg 把像素转正（文件变了），
|
||
// 要么退而求其次按"显示方向"归类（宽高对调）。
|
||
// 必须在算 MD5 / 分析明暗之前做——转正会改文件内容，得让 MD5 对得上最终文件。
|
||
o.correctImageRotation(f, r.res)
|
||
// 算 MD5（顺便拿指纹，可能省掉全量计算）
|
||
md5Start := time.Now()
|
||
r.md5Str, r.headMd5, r.tailMd5, r.size = o.calcFileInfo(f)
|
||
atomic.AddInt64(&o.tImgMd5, int64(time.Since(md5Start)))
|
||
// 分析明暗（GIF 不做，动画图逐帧分析没意义）
|
||
if !strings.EqualFold(f.Ext, ".gif") {
|
||
brightStart := time.Now()
|
||
dark, light, berr := probe.AnalyzeBrightness(f.Path)
|
||
atomic.AddInt64(&o.tBright, int64(time.Since(brightStart)))
|
||
if berr != nil {
|
||
fmt.Printf("=== 明暗分析失败(按中性mid归类): %s \n", f.Path)
|
||
} else {
|
||
r.darkScore, r.lightScore = dark, light
|
||
}
|
||
}
|
||
// 色情识别（GIF 不做；没启用检测器就直接跳过）
|
||
if o.nsfw != nil && !strings.EqualFold(f.Ext, ".gif") {
|
||
nsfwStart := time.Now()
|
||
score, detections, derr := o.nsfw.Detect(f.Path)
|
||
atomic.AddInt64(&o.tNsfw, int64(time.Since(nsfwStart)))
|
||
if derr != nil {
|
||
// 检测失败不影响整理流程
|
||
fmt.Printf("=== 色情检测失败: %s, 错误: %v \n", f.Path, derr)
|
||
} else {
|
||
r.nsfwScore = score
|
||
r.nsfwDetections = detections
|
||
checked := atomic.AddInt32(&o.nsfwChecked, 1)
|
||
fmt.Printf("=== 色情检测: %s, 评分=%.4f\n", f.Path, score)
|
||
// 每固定张数打一行进度，让前端能看到识别在推进（不逐张刷屏）
|
||
if checked%nsfwProgressInterval == 0 {
|
||
fmt.Printf("=== 色情识别进度: 已检查 %d 张\n", checked)
|
||
}
|
||
}
|
||
}
|
||
// 动漫色情识别（GIF 不做；没启用检测器就直接跳过）
|
||
if o.animeNsfw != nil && !strings.EqualFold(f.Ext, ".gif") {
|
||
animeStart := time.Now()
|
||
score, 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.animeNsfwRatings = ratings
|
||
checked := atomic.AddInt32(&o.animeNsfwChecked, 1)
|
||
fmt.Printf("=== 动漫色情检测: %s, 评分=%.4f\n", f.Path, score)
|
||
if checked%nsfwProgressInterval == 0 {
|
||
fmt.Printf("=== 动漫色情识别进度: 已检查 %d 张\n", checked)
|
||
}
|
||
}
|
||
}
|
||
return r
|
||
}
|
||
|
||
// correctImageRotation 处理图片的 EXIF 旋转标记，保证分类方向正确。
|
||
//
|
||
// 手机/相机竖拍的图，文件里常横着存像素、再打一个 EXIF 旋转标记（Orientation=5~8），
|
||
// 查看器会自动转正显示，但我们读到的宽高是"横着存"的，会把竖图错判成横图。
|
||
//
|
||
// 两种处理方式（按优先级）：
|
||
// 1. 配置了 ffmpeg 且允许转正：用 ffmpeg 把像素真正转正（烧进旋转），文件内容会变；
|
||
// 2. 没有 ffmpeg：不改文件，只把宽高对调，按"显示方向"归类。
|
||
//
|
||
// 只有方向 5/6/7/8（90°/270° 旋转）才需要宽高对调；2/3/4 是镜像或 180°，横竖不变。
|
||
func (o *Organizer) correctImageRotation(f scanner.File, res *probe.ProbeResult) {
|
||
if !o.rotateImages {
|
||
return // 配置关闭：完全不处理旋转，按存储方向/尺寸归类
|
||
}
|
||
ori := probe.ReadExifOrientation(f.Path)
|
||
if ori == 0 {
|
||
return // 没有旋转标记，不用管
|
||
}
|
||
// 有 ffmpeg 且允许转正：真把图转过来
|
||
if o.ffmpegPath != "" {
|
||
if err := probe.RotateImage(o.ffmpegPath, f.Path); err != nil {
|
||
fmt.Printf("=== 图片转正失败(按显示方向归类): %s, 错误: %v \n", f.Path, err)
|
||
} else {
|
||
// 转正后文件变了，重新读一遍尺寸，保证分类和实际文件一致
|
||
if res2, err := probe.ProbeImage(f.Path, o.ffprobePath); err == nil {
|
||
res.Width, res.Height = res2.Width, res2.Height
|
||
}
|
||
fmt.Printf("=== 已转正图片: %s (Orientation=%d)\n", f.Path, ori)
|
||
return
|
||
}
|
||
}
|
||
// 没有 ffmpeg（或转正失败）：退而求其次，按显示方向对调宽高归类
|
||
if probe.NeedsDimensionSwap(ori) {
|
||
res.Width, res.Height = res.Height, res.Width
|
||
}
|
||
}
|
||
|
||
// ProcessImages 处理全部图片：先并行体检，再按顺序一个个归类。
|
||
func (o *Organizer) ProcessImages(images []scanner.File) {
|
||
// 统计结果，最后打一行总结给用户看
|
||
success, fail, ignore, dup := 0, 0, 0, 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)
|
||
})
|
||
for i, r := range results {
|
||
// 用户点了中止：停下归类，已分堆的交给 runScan 决定要不要搬
|
||
if o.ctx != nil && o.ctx.Err() != nil {
|
||
fmt.Println("=== 扫描已中止，跳过剩余图片归类")
|
||
break
|
||
}
|
||
f := images[i]
|
||
switch {
|
||
case r.resume:
|
||
// 续传命中：上次已归类过，直接把文件归回原目标目录（没搬走的这次搬）
|
||
o.addToDir(f, r.resumeDir)
|
||
success++
|
||
case r.psd:
|
||
// PSD 设计稿：不做处理，原样收进「图片_PSD」
|
||
o.addClassified(f, dirImagePSD, r.size, r.headMd5, r.tailMd5)
|
||
success++
|
||
case r.unsupported:
|
||
// 不支持的格式：原样收进「图片_已忽略」
|
||
o.addClassified(f, dirImageIgnored, r.size, r.headMd5, r.tailMd5)
|
||
ignore++
|
||
case r.probeErr:
|
||
// 体检失败：收进「图片_读取异常」
|
||
o.addClassified(f, dirImageErr, r.size, r.headMd5, r.tailMd5)
|
||
fail++
|
||
fmt.Printf("=== 异常图片: %s \n", f.Path)
|
||
case o.isDuplicate(r.md5Str, f.Path):
|
||
// 同一个文件又出现一次：直接忽略，不动它、不记录
|
||
dup++
|
||
fmt.Printf("=== 重复图片(同路径), 忽略: %s \n", f.Path)
|
||
default:
|
||
// 正常图片：按分辨率/方向/明暗/色情评分归入对应的分类目录
|
||
info := o.record(f, "image", r.res, r.md5Str, r.headMd5, r.tailMd5, r.size)
|
||
// 检测过的图片都记录分数
|
||
if info != nil && r.nsfwScore > 0 {
|
||
info.NsfwScore = r.nsfwScore
|
||
if len(r.nsfwDetections) > 0 {
|
||
if detJSON, err := json.Marshal(r.nsfwDetections); err == nil {
|
||
info.NsfwDetections = string(detJSON)
|
||
}
|
||
}
|
||
}
|
||
if info != nil && (r.animeNsfwScore > 0 || r.animeNsfwRatings.General > 0) {
|
||
info.AnimeNsfwScore = r.animeNsfwScore
|
||
if ratingsJSON, err := json.Marshal(r.animeNsfwRatings); err == nil {
|
||
info.AnimeNsfwRatings = string(ratingsJSON)
|
||
}
|
||
}
|
||
nsfwEnabled := o.nsfw != nil || o.animeNsfw != nil
|
||
nsfwScore := r.nsfwScore
|
||
if r.animeNsfwScore > nsfwScore {
|
||
nsfwScore = r.animeNsfwScore
|
||
}
|
||
o.addClassified(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore, nsfwScore, nsfwEnabled), r.size, r.headMd5, r.tailMd5)
|
||
success++
|
||
}
|
||
}
|
||
// 把续传记录落库（中断/出错时尽量少丢已归类的结果）
|
||
o.flushScanLog()
|
||
if skip := atomic.LoadInt32(&o.resumeSkip); skip > 0 {
|
||
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, 动漫色情 %.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.tAnimeNsfw)).Seconds())
|
||
}
|
||
|
||
// ============================ 视频/音频整理 ============================
|
||
|
||
// mediaProbeResult 是"给一个视频/音频做完体检"的结论（并行阶段产出）。
|
||
type mediaProbeResult struct {
|
||
res *probe.ProbeResult // 体检结果（ffprobe 读出的信息）
|
||
err error // 体检失败的原因
|
||
md5 string // 文件 MD5
|
||
head string // 文件头采样 MD5
|
||
tail string // 文件尾采样 MD5
|
||
size int64 // 文件大小
|
||
resume bool // 续传命中：上次已归类过，跳过全部体检
|
||
resumeDir string // 续传命中的目标目录
|
||
}
|
||
|
||
// probeMedia 给一个媒体文件做体检（并行阶段执行）：调 ffprobe 读信息 + 算 MD5。
|
||
// 不访问任何共享可变状态，多个工人可以同时跑。
|
||
func (o *Organizer) probeMedia(f scanner.File) mediaProbeResult {
|
||
// 续传：上次已归类过、内容没变 → 跳过 ffprobe/MD5，直接复用上次的目标目录
|
||
if sl, ok := o.matchResume(f); ok {
|
||
atomic.AddInt32(&o.resumeSkip, 1)
|
||
return mediaProbeResult{resume: true, resumeDir: sl.TargetDir}
|
||
}
|
||
probeStart := time.Now()
|
||
res, err := probe.ProbeMediaFile(o.ffprobePath, f.Path)
|
||
atomic.AddInt64(&o.tMediaProbe, int64(time.Since(probeStart)))
|
||
if err != nil {
|
||
return mediaProbeResult{err: err}
|
||
}
|
||
// 既没有视频流也没有音频流：说明根本不是媒体，不用算 MD5 了
|
||
if !res.HasVideo && !res.HasAudio {
|
||
return mediaProbeResult{res: res}
|
||
}
|
||
md5Start := time.Now()
|
||
md5, head, tail, size := o.calcFileInfo(f)
|
||
atomic.AddInt64(&o.tMediaMd5, int64(time.Since(md5Start)))
|
||
return mediaProbeResult{res: res, md5: md5, head: head, tail: tail, size: size}
|
||
}
|
||
|
||
// ProcessMedia 处理全部媒体候选：先并行体检，再按顺序一个个归类。
|
||
//
|
||
// 特别说明：文件虽然进了"媒体候选"名单，但很多其实是"不知道是什么"的文件
|
||
// （比如扩展名被改过、或者 MIME 没认出来）。ffprobe 一探便知真假——
|
||
// 探出有视频/音频流的就是真媒体，探不出来的就看它的 MIME 是不是明确标注了媒体类型。
|
||
func (o *Organizer) ProcessMedia(media []scanner.File) {
|
||
video, audio, fail, ignore, dup := 0, 0, 0, 0, 0
|
||
total := len(media)
|
||
phaseStart := time.Now()
|
||
results := runParallel(o.ctx, media, o.probeMedia, func(done int) {
|
||
o.report(PhaseMedia, done, total)
|
||
})
|
||
for i, r := range results {
|
||
// 用户点了中止：停下归类，已分堆的交给 runScan 决定要不要搬
|
||
if o.ctx != nil && o.ctx.Err() != nil {
|
||
fmt.Println("=== 扫描已中止，跳过剩余媒体归类")
|
||
break
|
||
}
|
||
f := media[i]
|
||
// 续传命中：上次已归类过，直接把文件归回原目标目录（没搬走的这次搬）
|
||
if r.resume {
|
||
o.addToDir(f, r.resumeDir)
|
||
continue
|
||
}
|
||
// ---------- 体检失败 ----------
|
||
if r.err != nil {
|
||
// 只有 MIME 明确说是视频/音频的才收进"异常/忽略"目录；
|
||
// 完全没头绪的文件（MIME 为空或 octet-stream）原地不管，不搬走
|
||
if strings.HasPrefix(f.Mime, "video/") {
|
||
fail++
|
||
o.addClassified(f, dirVideoErr, r.size, r.head, r.tail)
|
||
fmt.Printf("=== 异常视频: %s \n", f.Path)
|
||
} else if scanner.IsExplicitMediaMime(f.Mime) {
|
||
ignore++
|
||
o.addClassified(f, dirVideoIgnored, r.size, r.head, r.tail)
|
||
} else {
|
||
ignore++
|
||
fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
|
||
}
|
||
continue
|
||
}
|
||
// ---------- 体检成功：看它到底是什么 ----------
|
||
switch {
|
||
case r.res.HasVideo:
|
||
// 有视频流 → 视频
|
||
if o.isDuplicate(r.md5, f.Path) {
|
||
// 同一个文件又出现一次：直接忽略，不动它、不记录
|
||
dup++
|
||
fmt.Printf("=== 重复媒体(同路径), 忽略: %s \n", f.Path)
|
||
continue
|
||
}
|
||
video++
|
||
o.record(f, "video", r.res, r.md5, r.head, r.tail, r.size)
|
||
o.addClassified(f, classify.VideoDir(f.Ext, r.res.Width, r.res.Height), r.size, r.head, r.tail)
|
||
case r.res.HasAudio:
|
||
// 有音频流 → 音频
|
||
if o.isDuplicate(r.md5, f.Path) {
|
||
// 同一个文件又出现一次：直接忽略，不动它、不记录
|
||
dup++
|
||
fmt.Printf("=== 重复媒体(同路径), 忽略: %s \n", f.Path)
|
||
continue
|
||
}
|
||
audio++
|
||
o.record(f, "audio", r.res, r.md5, r.head, r.tail, r.size)
|
||
o.addClassified(f, dirAudio, r.size, r.head, r.tail)
|
||
default:
|
||
// 没有视频流也没有音频流 → 假媒体。
|
||
// MIME 明确标注是媒体的，收进「已忽略」；否则原地不动。
|
||
if scanner.IsExplicitMediaMime(f.Mime) {
|
||
ignore++
|
||
o.addClassified(f, dirVideoIgnored, r.size, r.head, r.tail)
|
||
} else {
|
||
ignore++
|
||
fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
|
||
}
|
||
}
|
||
}
|
||
// 把续传记录落库（中断/出错时尽量少丢已归类的结果）
|
||
o.flushScanLog()
|
||
if skip := atomic.LoadInt32(&o.resumeSkip); skip > 0 {
|
||
fmt.Printf("=== 续传跳过: %d 个文件(上次已归类，直接复用)\n", skip)
|
||
}
|
||
fmt.Printf("=== 媒体处理完毕，视频: %d, 音频: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", video, audio, fail, ignore, dup)
|
||
fmt.Printf("=== 媒体归类耗时: 实际 %.1fs, 其中(并行累计) ffprobe %.1fs, MD5 %.1fs\n",
|
||
time.Since(phaseStart).Seconds(),
|
||
time.Duration(atomic.LoadInt64(&o.tMediaProbe)).Seconds(),
|
||
time.Duration(atomic.LoadInt64(&o.tMediaMd5)).Seconds())
|
||
}
|
||
|
||
// ============================ 移动（最后一步） ============================
|
||
|
||
// FlushRecordsOnly 只写入数据库记录，不移动文件。
|
||
// 用于 moveFiles=false 时，仍然保存媒体信息和续传记录。
|
||
func (o *Organizer) FlushRecordsOnly() {
|
||
o.flushRecords()
|
||
o.flushScanLog()
|
||
}
|
||
|
||
// Flush 收尾两步：
|
||
// 1. 把之前攒的媒体记录一次性写进数据库；
|
||
// 2. 把每个分类目录攒的文件**全局并行**搬过去（所有目录的文件一起搬）。
|
||
// 不需要记录每个文件搬去哪了——文件就在分类目录里，一眼能看到。
|
||
func (o *Organizer) Flush() {
|
||
o.FlushRecordsOnly()
|
||
|
||
// 把 buckets 里所有文件拼成一张"搬运清单"
|
||
jobs := make([]mover.MoveJob, 0)
|
||
for _, dir := range o.order {
|
||
for _, src := range o.buckets[dir] {
|
||
jobs = append(jobs, mover.MoveJob{Source: src, TargetDir: dir})
|
||
}
|
||
}
|
||
|
||
// 全局并行搬完：不管文件分散在几个目录，所有文件同时开搬。
|
||
// onProgress 收到的是全局累计已搬数，正好和总数对得上。
|
||
o.mover.MoveAll(jobs, func(done int) {
|
||
o.report(PhaseMove, done, len(jobs))
|
||
})
|
||
}
|