Files
kenaito-media-picker/internal/organize/organize.go
T

505 lines
20 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Package organize 是"整理"这一步的总指挥,负责把一批文件分门别类、搬到该去的地方。
//
// 整理一辆快递车的过程可以这样理解:
// 1. 探测(并行):每个文件先做"体检"——是图片还是视频?多大?什么分辨率?会不会是重复的?
// 2. 归类(串行):根据体检结果,决定这个文件该进哪个分类目录,或者直接判为重复/异常/忽略
// 3. 移动(并行):把每个分类目录攒齐的文件,批量搬进对应的文件夹,并把"搬去哪了"记进数据库
//
// 为什么探测要并行、归类要串行?
// - 探测很慢(读文件、算 MD5、调 ffprobe),多个一起干能明显提速;
// - 归类要看"前面有没有见过同样的文件"(去重判断),必须按顺序一个个来,否则会乱。
package organize
import (
"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/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 // 负责实际搬文件的家伙
buckets map[string][]string // 目标目录 → 里面该放哪些文件(只记源路径)
order []string // 目标目录的先后顺序(保证输出稳定、不来回跳)
onProgress func(phase string, done, total int)
pending []*store.MediaInfo // 待写入数据库的媒体记录(攒一批再批量写)
}
// New 创建整理器。store 可以为 nil——数据库打不开时,整理照常进行,只是不记入库信息。
func New(pickerDir, ffprobePath, ffmpegPath string, st *store.Store) *Organizer {
return &Organizer{
pickerDir: pickerDir,
ffprobePath: ffprobePath,
ffmpegPath: ffmpegPath,
store: st,
mover: &mover.Mover{},
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
}
// 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 每干完一个活就报一次数(已完成几个),线程安全、只会涨不会退。
func runParallel[T any, R any](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 {
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 // 亮色像素占比
}
// 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
}
}
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) {
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)
results := runParallel(images, o.probeImage, func(done int) {
o.report(PhaseImages, done, total)
})
for i, r := range results {
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)
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)
}
// ============================ 视频/音频整理 ============================
// 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(media, o.probeMedia, func(done int) {
o.report(PhaseMedia, done, total)
})
for i, r := range results {
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 收到的是全局累计已搬数,正好和总数对得上。
o.mover.MoveAll(jobs, func(done int) {
o.report(PhaseMove, done, len(jobs))
})
}