feat(core): 添加PSD/PSB图片格式支持并重构媒体文件处理

- 新增ReadPsdImage函数支持PSD/PSB格式图片信息读取 - 修改文件扫描逻辑，将视频处理改为媒体候选处理，使用ffprobe统一判定 - 添加对PSB格式的支持，与PSD共享相同8BPS文件头处理 - 重构handle_image.go中的图片处理逻辑，优化错误处理和记录保存 - 更新文件扫描过滤规则，支持更多图片格式并跳过数据库文件 - 添加文件MD5计算功能用于文件唯一性标识 - 修改媒体处理流程，支持音频文件识别和分类处理 - 移除废弃的视频时长转换函数和相关依赖 - 更新依赖版本并添加SQLite支持用于媒体信息存储
2025-03-18 21:04:15 +08:00
parent f5e3f6b4ef
commit b276e8c56f
16 changed files with 770 additions and 311 deletions
@@ -3,6 +3,7 @@ package core
 import (
 	"encoding/binary"
 	"fmt"
+	"io"
 	"os"
 )

@@ -85,3 +86,29 @@ func absInt32(v int32) int32 {
 	}
 	return v
 }
+
+// ReadPsdImage 读取PSD/PSB格式的图片信息（两者共享相同的8BPS文件头）
+func ReadPsdImage(filePath string) (err error, width int, height int) {
+	file, err := os.Open(filePath)
+	if err != nil {
+		fmt.Printf("=== Failed to open file: %v\n", err)
+		return err, 0, 0
+	}
+	defer file.Close()
+
+	var header [26]byte
+	if _, err := io.ReadFull(file, header[:]); err != nil {
+		fmt.Printf("=== Failed to read PSD/PSB header: %v\n", err)
+		return err, 0, 0
+	}
+	if string(header[0:4]) != "8BPS" {
+		return fmt.Errorf("not a valid PSD/PSB file: magic %q, expected \"8BPS\"", string(header[0:4])), 0, 0
+	}
+	version := binary.BigEndian.Uint16(header[4:6])
+	if version != 1 && version != 2 {
+		return fmt.Errorf("unsupported PSD/PSB version: %d", version), 0, 0
+	}
+	height = int(binary.BigEndian.Uint32(header[14:18]))
+	width = int(binary.BigEndian.Uint32(header[18:22]))
+	return nil, width, height
+}
@@ -3,6 +3,7 @@ package core
 import (
 	"fmt"
 	"github.com/redmask-hb/GoSimplePrint/goPrint"
+	"kenaito-media-picker/store"
 	"kenaito-media-picker/util"
 	"kenaito-media-picker/vars"
 	"log"
@@ -55,29 +56,21 @@ func (FileScanner) DoFilter() {
 	for _, globalFilePath := range vars.GlobalFilePathList {
 		fileMime := vars.GlobalFilePath2MimeInfoMap[globalFilePath]
 		count = count + 1
-		if strings.Contains(fileMime, "video/") { // 视频
-			vars.GlobalVideoPathList = append(vars.GlobalVideoPathList, globalFilePath)
-			bar.PrintBar(util.CalcPercentage(count, total))
-			continue
-		}
-		// mime格式为application/octet-stream的视频
+		// 图片
 		ext := path.Ext(globalFilePath)
-		if util.IsSupportVideo(ext) {
-			vars.GlobalVideoPathList = append(vars.GlobalVideoPathList, globalFilePath)
-			bar.PrintBar(util.CalcPercentage(count, total))
-			continue
-		}
-		if strings.Contains(fileMime, "image/") { // 图片
+		if strings.Contains(fileMime, "image/") || util.IsSupportImage(ext) || strings.EqualFold(ext, ".psd") {
 			vars.GlobalImagePathList = append(vars.GlobalImagePathList, globalFilePath)
 			bar.PrintBar(util.CalcPercentage(count, total))
 			continue
 		}
-		if util.IsSupportImage(ext) {
-			vars.GlobalImagePathList = append(vars.GlobalImagePathList, globalFilePath)
+		// 其余作为媒体候选，由ffprobe最终判定（视频/音频/不支持），
+		// 只要ffmpeg支持的格式都会被识别，不再依赖扩展名白名单
+		if util.IsPotentialMedia(fileMime) {
+			vars.GlobalMediaPathList = append(vars.GlobalMediaPathList, globalFilePath)
 			bar.PrintBar(util.CalcPercentage(count, total))
 			continue
 		}
-		// 其他的文件不处理
+		// 明显的非媒体文件（代码/文档/压缩包等）直接不处理
 	}
 	bar.PrintEnd("=== Finish")
 }
@@ -91,6 +84,9 @@ func visit(currentPath string, info os.FileInfo, err error) error {
 		return filepath.SkipDir // 跳过白名单目录（已整理文件的输出目录）
 	}
 	if !info.IsDir() {
+		if filepath.Base(currentPath) == store.DBFileName {
+			return nil // 跳过媒体信息数据库文件
+		}
 		vars.GlobalFilePathList = append(vars.GlobalFilePathList, currentPath)
 		// filename, include ext
 		vars.GlobalFilePath2FileNameMap[currentPath] = filepath.Base(currentPath)
@@ -18,26 +18,15 @@ func DoHandleImage(pickerDir string) {
 	ignoreCount := 0                       // 忽略数
 	for _, imageFilePath := range vars.GlobalImagePathList {
 		suffix := vars.GlobalFilePath2FileExtMap[imageFilePath]
-		if util.IsSupportImage(suffix) {
-			err, width, height := util.ReadImageInfo(imageFilePath)
-			if err == nil {
-				successCount = successCount + 1
-				vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", width, height)
-				fmt.Printf("=== Image总数: %d, 已读取Info: %d, 成功数: %d, 失败数: %d \n", total, successCount+errorCount+ignoreCount, successCount, errorCount)
-			} else {
-				errorCount = errorCount + 1
-				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
-				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
-			}
-			continue
-		}
 		if strings.EqualFold(suffix, ".webp") { // 特殊文件处理, webp为网络常用图片格式
 			webpErr, webpWidth, webpHeight := ReadWebpImage(imageFilePath)
 			if webpErr == nil {
 				vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", webpWidth, webpHeight)
+				saveMediaRecord(imageFilePath, "image", webpWidth, webpHeight, 0, "", "", "")
 				successCount = successCount + 1
 			} else {
 				errorCount = errorCount + 1
+				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
 				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
 			}
 			continue
@@ -46,9 +35,11 @@ func DoHandleImage(pickerDir string) {
 			bpmErr, bmpWidth, bmpHeight := ReadBmpImage(imageFilePath)
 			if bpmErr == nil {
 				vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", bmpWidth, bmpHeight)
+				saveMediaRecord(imageFilePath, "image", int(bmpWidth), int(bmpHeight), 0, "", "", "")
 				successCount = successCount + 1
 			} else {
 				errorCount = errorCount + 1
+				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
 				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
 			}
 			continue
@@ -58,6 +49,46 @@ func DoHandleImage(pickerDir string) {
 			successCount = successCount + 1
 			continue
 		}
+		if strings.EqualFold(suffix, ".psb") { // 特殊文件处理, psb与psd共享文件头
+			psbErr, psbWidth, psbHeight := ReadPsdImage(imageFilePath)
+			if psbErr == nil {
+				vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", psbWidth, psbHeight)
+				saveMediaRecord(imageFilePath, "image", psbWidth, psbHeight, 0, "", "", "")
+				successCount = successCount + 1
+			} else {
+				errorCount = errorCount + 1
+				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
+				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
+			}
+			continue
+		}
+		if util.IsSupportImage(suffix) {
+			err, width, height := util.ReadImageInfo(imageFilePath)
+			if err == nil {
+				successCount = successCount + 1
+				vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", width, height)
+				saveMediaRecord(imageFilePath, "image", width, height, 0, "", "", "")
+				fmt.Printf("=== Image总数: %d, 已读取Info: %d, 成功数: %d, 失败数: %d \n", total, successCount+errorCount+ignoreCount, successCount, errorCount)
+			} else if util.CheckFileIsExist(util.FFprobeExePath()) {
+				// Go标准库解码失败，尝试用ffprobe兜底（ffmpeg支持的图片格式更广，如tiff/svg等）
+				if probe, perr := util.ProbeMediaFile(util.FFprobeExePath(), imageFilePath); perr == nil {
+					if w, h := probe.VideoResolution(); w > 0 && h > 0 {
+						successCount = successCount + 1
+						vars.ImagePath2WidthHeightMap[imageFilePath] = fmt.Sprintf("%d-%d", w, h)
+						saveMediaRecord(imageFilePath, "image", w, h, 0, probe.VideoCodec(), probe.FormatName(), probe.RawJSON())
+						continue
+					}
+				}
+				errorCount = errorCount + 1
+				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
+				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
+			} else {
+				errorCount = errorCount + 1
+				vars.ImageReadErrorPathList = append(vars.ImageReadErrorPathList, imageFilePath)
+				fmt.Printf("=== 异常图片: %s \n", imageFilePath)
+			}
+			continue
+		}
 		// 其他的直接先忽略吧, 爱改改, 不改拉倒
 		ignoreCount = ignoreCount + 1
 		vars.ImageIgnorePathList = append(vars.ImageIgnorePathList, imageFilePath)
@@ -6,57 +6,60 @@ import (
 	"kenaito-media-picker/vars"
 	"os"
 	"strings"
-	"time"
 )

-func DoHandleVideo(pickerDir string) {
-	// 释放ffprobe
-	readerFileName := "./ffprobe.exe"
-	if util.CheckFileIsExist(readerFileName) {
-		_ = os.Remove(readerFileName)
-	}
-	err := util.WriteByteArraysToFile(vars.VideoDecodeToolWin64, readerFileName)
-	if err != nil {
-		fmt.Println("=== 释放解码器失败, 5秒后将自动退出", err)
-		time.Sleep(time.Second * 5)
-		return
-	}
-	total := len(vars.GlobalVideoPathList) // 总数
-	successCount := 0                      // 成功数
-	errorCount := 0                        // 失败数
-	ignoreCount := 0                       // 忽略数
-	for _, videoFilePath := range vars.GlobalVideoPathList {
-		suffix := vars.GlobalFilePath2FileExtMap[videoFilePath]
-		if util.IsSupportVideo(suffix) {
-			width, height, err := util.ReadVideoWidthHeight(videoFilePath)
-			if err == nil {
-				successCount = successCount + 1
-				vars.VideoPath2WidthHeightMap[videoFilePath] = fmt.Sprintf("%d-%d", width, height)
-				vars.VideoPath2WidthHeightTagMap[videoFilePath] = fmt.Sprintf("[%dx%d]", width, height)
-				fmt.Printf("=== Video总数: %d, 已读取Info: %d, 成功数: %d, 失败数: %d \n", total, successCount+errorCount+ignoreCount, successCount, errorCount)
-				duration := util.ReadVideoDuration(videoFilePath)
-				if duration == 0 {
-					vars.VideoPath2DurationMap[videoFilePath] = "0H0M0S"
-				} else {
-					vars.VideoPath2DurationMap[videoFilePath] = util.SecondsToHms(duration)
-				}
-			} else {
+// DoHandleMedia 处理所有媒体候选文件：ffprobe判定视频/音频/不支持
+func DoHandleMedia(pickerDir string) {
+	total := len(vars.GlobalMediaPathList)
+	videoCount := 0 // 视频数
+	audioCount := 0 // 音频数
+	errorCount := 0 // 异常数
+	ignoreCount := 0
+	for _, mediaFilePath := range vars.GlobalMediaPathList {
+		fileMime := vars.GlobalFilePath2MimeInfoMap[mediaFilePath]
+		probe, err := util.ProbeMediaFile(util.FFprobeExePath(), mediaFilePath)
+		if err != nil {
+			if strings.Contains(fileMime, "video/") {
+				// mime是视频但ffprobe读不出来，视为损坏/异常视频
 				errorCount = errorCount + 1
-				vars.VideoReadErrorPathList = append(vars.VideoReadErrorPathList, videoFilePath)
-				fmt.Printf("=== 异常视频: %s \n", videoFilePath)
+				vars.VideoReadErrorPathList = append(vars.VideoReadErrorPathList, mediaFilePath)
+				fmt.Printf("=== 异常视频: %s \n", mediaFilePath)
+			} else {
+				ignoreCount = ignoreCount + 1
+				vars.VideoIgnorePathList = append(vars.VideoIgnorePathList, mediaFilePath)
 			}
 			continue
 		}
-		// 其他的直接先忽略吧, 爱改改, 不改拉倒
+		width, height := probe.VideoResolution()
+		if width > 0 && height > 0 { // 有视频流
+			videoCount = videoCount + 1
+			vars.VideoPath2WidthHeightMap[mediaFilePath] = fmt.Sprintf("%d-%d", width, height)
+			saveMediaRecord(mediaFilePath, "video", width, height, probe.DurationSeconds(), probe.VideoCodec(), probe.FormatName(), probe.RawJSON())
+			fmt.Printf("=== 媒体总数: %d, 已处理: %d, 视频: %d, 音频: %d, 异常: %d, 忽略: %d \n", total, videoCount+audioCount+errorCount+ignoreCount, videoCount, audioCount, errorCount, ignoreCount)
+			continue
+		}
+		if probe.HasAudio() { // 仅有音频流
+			audioCount = audioCount + 1
+			vars.GlobalAudioPathList = append(vars.GlobalAudioPathList, mediaFilePath)
+			saveMediaRecord(mediaFilePath, "audio", 0, 0, probe.DurationSeconds(), probe.AudioCodec(), probe.FormatName(), probe.RawJSON())
+			continue
+		}
+		// 既无视频也无音频（字幕流等）
 		ignoreCount = ignoreCount + 1
-		vars.VideoIgnorePathList = append(vars.VideoIgnorePathList, videoFilePath)
+		vars.VideoIgnorePathList = append(vars.VideoIgnorePathList, mediaFilePath)
 	}
 	if len(vars.VideoReadErrorPathList) > 0 {
-		readInfoErrorPath := pickerDir + string(os.PathSeparator) + "读取异常"
+		readInfoErrorPath := pickerDir + string(os.PathSeparator) + "视频_读取异常"
 		if util.CreateDir(readInfoErrorPath) {
 			util.DoMoveFileToDir(vars.VideoReadErrorPathList, readInfoErrorPath)
 		}
 	}
+	if len(vars.GlobalAudioPathList) > 0 {
+		audioPath := pickerDir + string(os.PathSeparator) + "音频"
+		if util.CreateDir(audioPath) {
+			util.DoMoveFileToDir(vars.GlobalAudioPathList, audioPath)
+		}
+	}
 	if len(vars.VideoIgnorePathList) > 0 {
 		ignorePath := pickerDir + string(os.PathSeparator) + "已忽略"
 		if util.CreateDir(ignorePath) {
@@ -64,21 +67,13 @@ func DoHandleVideo(pickerDir string) {
 		}
 	}
 	doPickVideoFile(pickerDir, vars.VideoPath2WidthHeightMap)
-	// 删除ffprobe
-	if util.CheckFileIsExist(readerFileName) {
-		_ = os.Remove(readerFileName)
-	}
-	fmt.Printf("=== 视频处理完毕 \n\n")
+	fmt.Printf("=== 视频/音频处理完毕 \n\n")
 }

 // 条件视频并分组存放
 func doPickVideoFile(pickerDir string, videoPath2WidthHeightMap map[string]string) {
 	if len(videoPath2WidthHeightMap) == 0 {
 		fmt.Printf("=== 当前目录下没有扫描到视频文件, %s \n", pickerDir)
-		readerFileName := "./ffprobe.exe"
-		if util.CheckFileIsExist(readerFileName) {
-			_ = os.Remove(readerFileName)
-		}
 		return
 	}
 	for currentVideoPath, infoStr := range videoPath2WidthHeightMap {
@@ -0,0 +1,38 @@
+package core
+
+import (
+	"fmt"
+	"kenaito-media-picker/store"
+	"kenaito-media-picker/util"
+	"kenaito-media-picker/vars"
+	"os"
+)
+
+// saveMediaRecord 计算MD5并将媒体信息写入数据库
+func saveMediaRecord(filePath string, mediaType string, width int, height int, duration float64, codec string, format string, rawInfo string) {
+	md5Str, err := util.CalcFileMd5(filePath)
+	if err != nil {
+		fmt.Printf("=== 计算MD5失败: %s, 错误: %v\n", filePath, err)
+		return
+	}
+	info := &store.MediaInfo{
+		FileMd5:   md5Str,
+		FilePath:  filePath,
+		FileName:  vars.GlobalFilePath2FileNameMap[filePath],
+		FileExt:   vars.GlobalFilePath2FileExtMap[filePath],
+		MimeType:  vars.GlobalFilePath2MimeInfoMap[filePath],
+		MediaType: mediaType,
+		Width:     width,
+		Height:    height,
+		Duration:  duration,
+		Codec:     codec,
+		Format:    format,
+		RawInfo:   rawInfo,
+	}
+	if st, err := os.Stat(filePath); err == nil {
+		info.FileSize = st.Size()
+	}
+	if err := store.SaveMedia(info); err != nil {
+		fmt.Printf("=== 保存媒体信息失败: %s, 错误: %v\n", filePath, err)
+	}
+}