// Package organize 是"整理"这一步的总指挥,负责把一批文件分门别类、搬到该去的地方。 // // 整理一辆快递车的过程可以这样理解: // 1. 探测(并行):每个文件先做"体检"——是图片还是视频?多大?什么分辨率?会不会是重复的? // 2. 归类(串行):根据体检结果,决定这个文件该进哪个分类目录,或者直接判为重复/异常/忽略 // 3. 移动(并行):把每个分类目录攒齐的文件,批量搬进对应的文件夹,并把"搬去哪了"记进数据库 // // 为什么探测要并行、归类要串行? // - 探测很慢(读文件、算 MD5、调 ffprobe),多个一起干能明显提速; // - 归类要看"前面有没有见过同样的文件"(去重判断),必须按顺序一个个来,否则会乱。 package organize import ( "context" "fmt" "os" "path/filepath" "strings" "sync" "sync/atomic" "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 = "音频" // 音频统一放这一个文件夹 ) // ============================ 整理器主体 ============================ // 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 表示没启用,图片不检测 rotateImages bool // 是否处理 EXIF 旋转(转正/按显示方向归类);false 则完全不动,按存储方向归类 buckets map[string][]string // 目标目录 → 里面该放哪些文件(只记源路径) order []string // 目标目录的先后顺序(保证输出稳定、不来回跳) onProgress func(phase string, done, total int) pending []*store.MediaInfo // 待写入数据库的媒体记录(攒一批再批量写) pendingNsfw []store.NsfwRecord // 待写入数据库的色情识别记录(攒一批再批量写) } // New 创建整理器。store 可以为 nil——数据库打不开时,整理照常进行,只是不记入库信息。 // ctx 可以为 nil——不提供中止能力,任务跑到完为止。 func New(pickerDir, ffprobePath, ffmpegPath string, st *store.Store, ctx context.Context) *Organizer { return &Organizer{ pickerDir: pickerDir, ffprobePath: ffprobePath, ffmpegPath: ffmpegPath, store: st, mover: &mover.Mover{}, ctx: ctx, rotateImages: true, // 默认处理 EXIF 旋转;关掉用 SetRotateImages(false) buckets: make(map[string][]string), } } // 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 } // 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) { dir := filepath.Join(o.pickerDir, subDir) if _, ok := o.buckets[dir]; !ok { o.order = append(o.order, dir) } o.buckets[dir] = append(o.buckets[dir], f.Path) } // 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 把一份媒体信息记到"待入库"缓冲里(先攒着,最后统一批量写库)。 func (o *Organizer) record(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) { if o.store == nil || md5Str == "" { return } o.pending = append(o.pending, buildMediaInfo(f, mediaType, res, md5Str, headMd5, tailMd5, size)) } // 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 // 体检失败(打不开/读不出),归入"读取异常" res *probe.ProbeResult // 体检出的尺寸等信息(体检成功时才有) md5Str string // 文件 MD5 headMd5 string // 文件头采样 MD5(大文件指纹) tailMd5 string // 文件尾采样 MD5 size int64 // 文件大小 darkScore float64 // 暗色像素占比(用于亮度分类) lightScore float64 // 亮色像素占比 nsfwScore float64 // 色情检测最高分(未检测为 0) nsfw bool // 是否判定为色情(检测可用时才可能为 true) } // probeImage 给一张图片做体检(并行阶段执行):探测尺寸 + 算 MD5 + 分析明暗。 // 注意:这里不能碰 o.buckets 之类的共享可变状态,因为多个工人同时跑; // 只读数据库(查指纹),这是安全的。 func (o *Organizer) probeImage(f scanner.File) imageProbeResult { r := imageProbeResult{} // 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 } // 探测尺寸 res, err := probe.ProbeImage(f.Path, o.ffprobePath) if err != nil { r.probeErr = true return r } r.res = res // 先处理 EXIF 旋转:带旋转标记的图,要么用 ffmpeg 把像素转正(文件变了), // 要么退而求其次按"显示方向"归类(宽高对调)。 // 必须在算 MD5 / 分析明暗之前做——转正会改文件内容,得让 MD5 对得上最终文件。 o.correctImageRotation(f, r.res) // 算 MD5(顺便拿指纹,可能省掉全量计算) r.md5Str, r.headMd5, r.tailMd5, r.size = o.calcFileInfo(f) // 分析明暗(GIF 不做,动画图逐帧分析没意义) if !strings.EqualFold(f.Ext, ".gif") { dark, light, berr := probe.AnalyzeBrightness(f.Path) 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") { score, hit, derr := o.nsfw.Detect(f.Path) if derr != nil { // 检测失败按"未命中"处理,不影响整理流程 fmt.Printf("=== 色情检测失败(按未命中归类): %s, 错误: %v \n", f.Path, derr) } else { r.nsfwScore = score r.nsfw = hit // 命中的图片把"文件名 + 检测分"打进运行日志,方便用户盯着看是否判对 if hit { fmt.Printf("=== 色情命中: %s, 检测分=%.2f (阈值=%.2f)\n", f.Path, score, o.nsfw.Threshold()) } } } 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 nsfwCount := 0 // 命中的色情图片数 total := len(images) 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.psd: // PSD 设计稿:不做处理,原样收进「图片_PSD」 o.add(f, dirImagePSD) success++ case r.unsupported: // 不支持的格式:原样收进「图片_已忽略」 o.add(f, dirImageIgnored) ignore++ case r.probeErr: // 体检失败:收进「图片_读取异常」 o.add(f, dirImageErr) fail++ fmt.Printf("=== 异常图片: %s \n", f.Path) case o.isDuplicate(r.md5Str, f.Path): // 同一个文件又出现一次:直接忽略,不动它、不记录 dup++ fmt.Printf("=== 重复图片(同路径), 忽略: %s \n", f.Path) default: // 正常图片:按分辨率/方向/明暗归入对应的分类目录; // 被判为色情的改去「图片_敏感_*」目录,明暗子目录照常保留 o.record(f, "image", r.res, r.md5Str, r.headMd5, r.tailMd5, r.size) if r.nsfw { o.add(f, classify.ImageDirSensitive(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore)) // 攒一条色情识别记录,最后连同"搬去哪了"一起写进数据库 o.pendingNsfw = append(o.pendingNsfw, store.NsfwRecord{ FilePath: f.Path, FileName: filepath.Base(f.Path), Score: r.nsfwScore, Threshold: o.nsfw.Threshold(), }) nsfwCount++ } else { o.add(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore)) } success++ } } fmt.Printf("=== 图片处理完毕,成功: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", success, fail, ignore, dup) if o.nsfw != nil && nsfwCount > 0 { fmt.Printf("=== 色情识别: 共命中 %d 张,已归入「图片_敏感_*」目录,可在【色情识别】页面查看并人工复核\n", nsfwCount) } } // ============================ 视频/音频整理 ============================ // mediaProbeResult 是"给一个视频/音频做完体检"的结论(并行阶段产出)。 type mediaProbeResult struct { res *probe.ProbeResult // 体检结果(ffprobe 读出的信息) err error // 体检失败的原因 md5 string // 文件 MD5 head string // 文件头采样 MD5 tail string // 文件尾采样 MD5 size int64 // 文件大小 } // probeMedia 给一个媒体文件做体检(并行阶段执行):调 ffprobe 读信息 + 算 MD5。 // 不访问任何共享可变状态,多个工人可以同时跑。 func (o *Organizer) probeMedia(f scanner.File) mediaProbeResult { res, err := probe.ProbeMediaFile(o.ffprobePath, f.Path) if err != nil { return mediaProbeResult{err: err} } // 既没有视频流也没有音频流:说明根本不是媒体,不用算 MD5 了 if !res.HasVideo && !res.HasAudio { return mediaProbeResult{res: res} } md5, head, tail, size := o.calcFileInfo(f) 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) 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.err != nil { // 只有 MIME 明确说是视频/音频的才收进"异常/忽略"目录; // 完全没头绪的文件(MIME 为空或 octet-stream)原地不管,不搬走 if strings.HasPrefix(f.Mime, "video/") { fail++ o.add(f, dirVideoErr) fmt.Printf("=== 异常视频: %s \n", f.Path) } else if scanner.IsExplicitMediaMime(f.Mime) { ignore++ o.add(f, dirVideoIgnored) } 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.add(f, classify.VideoDir(f.Ext, r.res.Width, r.res.Height)) 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.add(f, dirAudio) default: // 没有视频流也没有音频流 → 假媒体。 // MIME 明确标注是媒体的,收进「已忽略」;否则原地不动。 if scanner.IsExplicitMediaMime(f.Mime) { ignore++ o.add(f, dirVideoIgnored) } else { ignore++ fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path) } } } fmt.Printf("=== 媒体处理完毕,视频: %d, 音频: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", video, audio, fail, ignore, dup) } // ============================ 移动(最后一步) ============================ // Flush 收尾两步: // 1. 把之前攒的媒体记录一次性写进数据库; // 2. 把每个分类目录攒的文件**全局并行**搬过去(所有目录的文件一起搬)。 // 不需要记录每个文件搬去哪了——文件就在分类目录里,一眼能看到。 func (o *Organizer) Flush() { o.flushRecords() // 把 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 收到的是全局累计已搬数,正好和总数对得上。 // moved 是"源路径 → 新位置"的对照表,色情记录要用它补上整理后的实际路径。 moved := o.mover.MoveAll(jobs, func(done int) { o.report(PhaseMove, done, len(jobs)) }) // 把色情识别记录补上"整理后位置"(预览就用它),再批量写进数据库 if len(o.pendingNsfw) > 0 { records := make([]*store.NsfwRecord, 0, len(o.pendingNsfw)) for i := range o.pendingNsfw { if target, ok := moved[o.pendingNsfw[i].FilePath]; ok { o.pendingNsfw[i].MovedPath = target } else { o.pendingNsfw[i].MovedPath = o.pendingNsfw[i].FilePath // 没搬走就用源路径 } records = append(records, &o.pendingNsfw[i]) } if o.store != nil { if err := o.store.SaveNsfwBatch(records); err != nil { fmt.Printf("=== 保存色情识别记录失败: %v\n", err) } } o.pendingNsfw = nil } }