feat: 断点续传与色情识别提速/修复,支持 Linux 交叉编译
- feat: 新增续传表 scan_log,重扫跳过已归类文件(尺寸/头尾指纹校验) - fix: 色情识别记录批量入库超 SQLite 999 变量上限,改每 100 条一批 - perf: NudeNet 检测改为多线程推理池(每线程独立模型实例) - perf: 明暗分析 JPEG(YCbCr) 快路径;SQLite 读并发放开;加阶段耗时统计 - fix: 前端色情识别状态由后端如实上报(配置开关+检测服务是否起来) - feat: ffprobe 内嵌按平台拆分,Linux 构建回退系统 ffprobe,支持交叉编译 - docs: README 补充断点续传与 CentOS 交叉编译步骤 - chore: organize.workers 默认 8
This commit is contained in:
@@ -10,9 +10,15 @@ import (
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"kenaito-media-picker/internal/config"
|
||||
"kenaito-media-picker/internal/store"
|
||||
)
|
||||
|
||||
// nsfwDetectorReadyKey 记录"最近一次扫描时色情检测服务是否真的起来了"(存 AppConfig)。
|
||||
// 配置开了但检测器没起来(没装 python/模型下载失败等)也能区分出来,方便前端提示。
|
||||
// 由 scan 流程在服务启动成功/失败时写 1/0。
|
||||
const nsfwDetectorReadyKey = "nsfw.detectorReady"
|
||||
|
||||
// NsfwController 色情识别相关接口的"窗口"。
|
||||
type NsfwController struct {
|
||||
store *store.Store // 数据库(只读)
|
||||
@@ -55,7 +61,13 @@ func (c *NsfwController) List(ctx *gin.Context) {
|
||||
CreatedAt: r.CreatedAt.Format("2006-01-02 15:04:05"),
|
||||
})
|
||||
}
|
||||
Success(ctx, map[string]interface{}{"list": views, "total": len(views)})
|
||||
Success(ctx, map[string]interface{}{
|
||||
"list": views,
|
||||
"total": len(views),
|
||||
"nsfwEnabled": config.GetAppConfig().NSFWEnabled,
|
||||
// 没扫描过(空)或上次成功(1)都算可用;只有明确记过失败(0)才提示检测服务没起来
|
||||
"nsfwDetector": c.store == nil || c.store.GetConfig(nsfwDetectorReadyKey) != "0",
|
||||
})
|
||||
}
|
||||
|
||||
// Preview 处理 GET /api/nsfw/preview?path=xxx:把指定图片文件返回给浏览器做预览。
|
||||
|
||||
+64
-4
@@ -13,6 +13,7 @@ import (
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -22,6 +23,14 @@ import (
|
||||
uuid "github.com/satori/go.uuid"
|
||||
)
|
||||
|
||||
// uuidPrefixRe 匹配"一长串 32 位十六进制 UUID 前缀"的文件名(如 3f2a...9c1d_photo.jpg)。
|
||||
// 用它可以剥掉之前搬移留下的 UUID 前缀,还原视频的原始文件名。
|
||||
var uuidPrefixRe = regexp.MustCompile(`^[0-9a-fA-F]{32}[_\.]`)
|
||||
|
||||
// nameMu 保护"选名 + 独占创建",保证并发搬多个同名文件时不互相踩到。
|
||||
// 只锁"选名"这一小段(独占创建完就放锁),真正的复制/改名在锁外进行,不影响并行效率。
|
||||
var nameMu sync.Mutex
|
||||
|
||||
// MoveJob 一次"搬文件"的任务:把 Source 这个文件搬进 TargetDir 目录。
|
||||
type MoveJob struct {
|
||||
Source string // 文件现在在哪
|
||||
@@ -40,7 +49,10 @@ func workerCount() int {
|
||||
// MoveAll 把一批"搬文件任务"全局并行搬完——不管任务分散在几个目录,
|
||||
// 所有文件同时开搬(最多 8 个工人),这才是真正的并行。
|
||||
//
|
||||
// 目标文件名会加一长串随机前缀(uuid),这样即使两个文件同名也不会互相覆盖。
|
||||
// 命名规则看目标目录类型:
|
||||
// - 图片目录(名字带「图片」):新文件名 = 一串随机 UUID + 扩展名,不保留原名;
|
||||
// - 视频/音频等目录:保留原文件名,重名时自动追加 (1)、(2)… 序号。
|
||||
//
|
||||
// 返回"移动成功的 源路径 → 新位置完整路径"的对照表,方便调用方记录每个文件去哪了。
|
||||
// onProgress 每搬完一个文件调一次,参数是"已累计搬完多少个"(只会涨不会退,并发安全)。
|
||||
func (m *Mover) MoveAll(jobs []MoveJob, onProgress func(done int)) map[string]string {
|
||||
@@ -103,9 +115,21 @@ func (m *Mover) MoveAll(jobs []MoveJob, onProgress func(done int)) map[string]st
|
||||
// 注意一种特殊情况:复制成功了但删源文件失败(比如源文件被占用),
|
||||
// 这时源和目标会同时存在,但依然算"移动成功",返回新路径。
|
||||
func moveOne(src, targetDir string) (string, bool) {
|
||||
// 新文件名 = 随机前缀 + 原文件名,防止重名覆盖
|
||||
fileName := strings.ReplaceAll(uuid.NewV4().String(), "-", "") + "_" + filepath.Base(src)
|
||||
target := filepath.Join(targetDir, fileName)
|
||||
// 命名规则:图片不保留原名(直接 UUID + 扩展名),视频/音频保留原文件名
|
||||
keepName := !strings.Contains(targetDir, "图片")
|
||||
base := filepath.Base(src)
|
||||
if keepName {
|
||||
// 视频/音频:剥掉之前搬移留下的 UUID 前缀,还原真正的原文件名
|
||||
base = stripUuidPrefix(base)
|
||||
} else {
|
||||
base = strings.ReplaceAll(uuid.NewV4().String(), "-", "") + filepath.Ext(base)
|
||||
}
|
||||
// 文件本来就在目标目录、名字也一样(重扫时分类没变)→ 不用搬
|
||||
if filepath.Clean(filepath.Join(targetDir, base)) == filepath.Clean(src) {
|
||||
return src, true
|
||||
}
|
||||
// 选一个不冲突的目标名(重名自动加 (1)、(2)…),并独占占住,防止并发踩踏
|
||||
target := reservePath(targetDir, base)
|
||||
// 先在同一个磁盘里直接改名(秒完成)
|
||||
if err := os.Rename(src, target); err == nil {
|
||||
return target, true
|
||||
@@ -121,6 +145,42 @@ func moveOne(src, targetDir string) (string, bool) {
|
||||
return target, true
|
||||
}
|
||||
|
||||
// stripUuidPrefix 剥掉文件名开头的 UUID 前缀,还原原始文件名。
|
||||
// 比如 c3f2...7b_电影.mp4 → 电影.mp4。
|
||||
// 纯 UUID + 扩展名(原始名早已丢失)时保持原样不动。
|
||||
func stripUuidPrefix(base string) string {
|
||||
if uuidPrefixRe.MatchString(base) && base[32] == '_' {
|
||||
return base[33:]
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
// reservePath 在 targetDir 里找一个不重名的文件名并"独占创建"出来占住坑位,
|
||||
// 返回完整路径。重名时自动在扩展名前加序号:原名 (1).ext、原名 (2).ext……
|
||||
// 独占创建保证并发搬同名文件时,各自挑到不同的名字,不会互相覆盖。
|
||||
func reservePath(targetDir, base string) string {
|
||||
nameMu.Lock()
|
||||
defer nameMu.Unlock()
|
||||
ext := filepath.Ext(base)
|
||||
stem := strings.TrimSuffix(base, ext)
|
||||
for i := 0; ; i++ {
|
||||
name := base
|
||||
if i > 0 {
|
||||
name = fmt.Sprintf("%s (%d)%s", stem, i, ext)
|
||||
}
|
||||
p := filepath.Join(targetDir, name)
|
||||
f, err := os.OpenFile(p, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
|
||||
if err == nil {
|
||||
f.Close()
|
||||
return p
|
||||
}
|
||||
// 文件已存在 → 换个序号再试;其他错误(权限等)就当占住失败,原样返回兜底
|
||||
if !os.IsExist(err) {
|
||||
return p
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// copyFile 把一个文件的内容完整复制到另一个位置。
|
||||
// 复制到一半失败时,会把已经写出来的半截文件删掉,不给磁盘留垃圾。
|
||||
func copyFile(src, dst string) error {
|
||||
|
||||
+137
-41
@@ -15,6 +15,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
@@ -48,17 +49,20 @@ const requestTimeout = 30 * time.Second
|
||||
|
||||
// Detector 封装一个本地 python 服务的客户端。
|
||||
// 多个探测协程会并发调 Detect,底层 HTTP 客户端线程安全;
|
||||
// 服务端单线程处理,请求自己排队,不用额外加锁。
|
||||
// 服务端用固定大小线程池并行推理(每个 worker 一份独立模型实例),请求不必排队。
|
||||
type Detector struct {
|
||||
url string // 本地服务地址,如 http://127.0.0.1:54321
|
||||
threshold float64 // 判定阈值 0~1
|
||||
client *http.Client
|
||||
cmd *exec.Cmd // python 进程句柄(Stop 时杀掉)
|
||||
scriptDir string // 临时脚本所在目录(Stop 时清理)
|
||||
mu sync.Mutex // 保护下面的关闭状态
|
||||
python string // python 可执行文件路径(服务重启时要用)
|
||||
modelDir string // 模型目录 lib(服务重启时要用)
|
||||
mu sync.Mutex // 保护 stopped / 进程句柄
|
||||
stopped bool // 已经 Stop 过了
|
||||
cancel context.CancelFunc
|
||||
done chan struct{} // 进程退出时关闭(Wait 协程负责)
|
||||
restartMu sync.Mutex // 串行化重启:同一时刻只一个 goroutine 重启,其余排队等
|
||||
}
|
||||
|
||||
// Start 拉起本地 python 检测服务并等它就绪。
|
||||
@@ -87,13 +91,6 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
return nil, fmt.Errorf("写入检测脚本失败: %v", err)
|
||||
}
|
||||
|
||||
// 找一个空闲端口
|
||||
port, err := freePort()
|
||||
if err != nil {
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
return nil, fmt.Errorf("找空闲端口失败: %v", err)
|
||||
}
|
||||
|
||||
// 模型目录:程序运行目录下的 lib(和 application.properties / media-info.db 同一根目录)。
|
||||
// 服务脚本优先加载 lib/320n.onnx;本地没有会自己下载到那里,下载失败再回退用 pip 自带模型。
|
||||
modelDir := ""
|
||||
@@ -101,27 +98,51 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
modelDir = filepath.Join(wd, "lib")
|
||||
}
|
||||
|
||||
// 拉起服务。标准输出/错误接到程序自身的输出,方便看 python 报错。
|
||||
d := &Detector{
|
||||
threshold: threshold,
|
||||
client: &http.Client{Timeout: requestTimeout},
|
||||
python: py,
|
||||
modelDir: modelDir,
|
||||
scriptDir: tmpDir,
|
||||
}
|
||||
|
||||
// 拉起服务(同样一套逻辑也用于崩溃后的重启)
|
||||
if err := d.launch(); err != nil {
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
return nil, err
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// launch 启动(或重启)python 检测服务,并等它就绪(模型加载完成)。
|
||||
// 调用方需要保证持有 mu,且 stopped 为 false。
|
||||
func (d *Detector) launch() error {
|
||||
// 先把上一次的进程收拾掉(首次启动时没有,直接跳过)
|
||||
d.teardown()
|
||||
|
||||
// 找一个空闲端口
|
||||
port, err := freePort()
|
||||
if err != nil {
|
||||
return fmt.Errorf("找空闲端口失败: %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port), modelDir)
|
||||
cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
|
||||
// 标准输出/错误都接到程序自身的输出:
|
||||
// - 扫描期间 os.Stdout 会被接管进"运行日志",这样 python 的报错(崩溃原因)也能在前端看到
|
||||
// - 方便排查服务为什么挂掉
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
cmd.Stderr = os.Stdout
|
||||
if err := cmd.Start(); err != nil {
|
||||
cancel()
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
return nil, fmt.Errorf("启动检测服务失败: %v", err)
|
||||
return fmt.Errorf("启动检测服务失败: %v", err)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
d := &Detector{
|
||||
url: fmt.Sprintf("http://127.0.0.1:%d", port),
|
||||
threshold: threshold,
|
||||
client: &http.Client{Timeout: requestTimeout},
|
||||
cmd: cmd,
|
||||
scriptDir: tmpDir,
|
||||
cancel: cancel,
|
||||
done: done,
|
||||
}
|
||||
d.url = fmt.Sprintf("http://127.0.0.1:%d", port)
|
||||
d.cmd = cmd
|
||||
d.cancel = cancel
|
||||
d.done = done
|
||||
// 后台等进程退出:ProcessState 要 Wait 之后才有值,这里用 done 通知,
|
||||
// 这样 waitReady 能及时发现"python 没装 nudenet 直接退掉"的情况。
|
||||
go func() {
|
||||
@@ -131,10 +152,32 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
|
||||
// 等服务就绪(模型加载完成)。就绪前先探测 /health。
|
||||
if err := d.waitReady(); err != nil {
|
||||
d.Stop()
|
||||
return nil, err
|
||||
d.teardown()
|
||||
return err
|
||||
}
|
||||
return d, nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// teardown 停掉当前 python 进程并清理句柄(不删临时目录)。
|
||||
// 调用方需要保证持有 mu。进程可能已经死了(崩溃),kill 失败也无所谓。
|
||||
func (d *Detector) teardown() {
|
||||
if d.cancel != nil {
|
||||
d.cancel()
|
||||
d.cancel = nil
|
||||
}
|
||||
if d.cmd != nil && d.cmd.Process != nil {
|
||||
_ = d.cmd.Process.Kill()
|
||||
}
|
||||
if d.done != nil {
|
||||
// 等 Wait 协程回收进程,避免僵尸进程;进程已死则立刻返回。
|
||||
// 等超过 3 秒就放弃,防止极端情况下卡死重启流程。
|
||||
select {
|
||||
case <-d.done:
|
||||
case <-time.After(3 * time.Second):
|
||||
}
|
||||
d.done = nil
|
||||
}
|
||||
d.cmd = nil
|
||||
}
|
||||
|
||||
// waitReady 轮询本地服务,直到能正常应答。
|
||||
@@ -168,14 +211,73 @@ func (d *Detector) waitReady() error {
|
||||
// - score:命中的敏感类里最高分数(没有命中为 0)
|
||||
// - nsfw:是否判定为色情(有敏感类且分数 >= 阈值)
|
||||
// - err:检测失败(服务挂了/图读不出等),由调用方兜底按未命中处理
|
||||
//
|
||||
// 服务失联(连接被拒,多半是 python 进程在处理某张图时崩了)时,会串行化地
|
||||
// 重启服务并重试一次:多个并发 worker 里只有一个负责重启,其余排队,等重启完
|
||||
// 也重试自己的图——保证服务短暂挂掉不会漏掉一整批图片。
|
||||
func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
|
||||
score, nsfw, err = d.detectOnce(path)
|
||||
if err == nil {
|
||||
return score, nsfw, nil
|
||||
}
|
||||
// 只有"发请求都失败"(连接被拒/重置,说明服务失联)才走重启;
|
||||
// 服务正常返回的明确错误(如某张图读不了)不重启。
|
||||
var ce *connError
|
||||
if !errors.As(err, &ce) {
|
||||
return score, nsfw, err
|
||||
}
|
||||
|
||||
d.restartMu.Lock()
|
||||
defer d.restartMu.Unlock()
|
||||
|
||||
if d.isStopped() {
|
||||
return 0, false, err
|
||||
}
|
||||
// 可能排在前面的 goroutine 已经把服务重启好了:先探活,活着就直接重试,别反复重启
|
||||
if d.healthy() {
|
||||
return d.detectOnce(path)
|
||||
}
|
||||
fmt.Printf("=== 色情检测服务失联(%v), 正在重启…\n", ce.err)
|
||||
if rerr := d.launch(); rerr != nil {
|
||||
return 0, false, fmt.Errorf("检测服务重启失败: %v", rerr)
|
||||
}
|
||||
fmt.Println("=== 色情检测服务已重启,继续识别")
|
||||
return d.detectOnce(path)
|
||||
}
|
||||
|
||||
// healthy 快速探活:本地服务是否还能正常应答(带短超时,避免拖慢扫描)。
|
||||
func (d *Detector) healthy() bool {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.url+"/health", nil)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
resp, err := d.client.Do(req)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
return resp.StatusCode == http.StatusOK
|
||||
}
|
||||
|
||||
// isStopped 判断检测器是否已被 Stop(线程安全)。
|
||||
func (d *Detector) isStopped() bool {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return d.stopped
|
||||
}
|
||||
|
||||
// detectOnce 发一次检测请求(不带重启逻辑)。
|
||||
func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error) {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"path": path,
|
||||
"threshold": d.threshold,
|
||||
})
|
||||
resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return 0, false, err
|
||||
// 打上"连接失败"标记,供 Detect 判断是否要重启
|
||||
return 0, false, &connError{err}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
@@ -197,27 +299,21 @@ func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
|
||||
return out.Score, out.Nsfw, nil
|
||||
}
|
||||
|
||||
// connError 标记"请求根本没发出去/连不上本地服务"的错误(服务失联)。
|
||||
type connError struct{ err error }
|
||||
|
||||
func (e *connError) Error() string { return e.err.Error() }
|
||||
func (e *connError) Unwrap() error { return e.err }
|
||||
|
||||
// Stop 停掉检测服务并清理临时文件。可以安全地重复调用。
|
||||
func (d *Detector) Stop() {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
if d.stopped {
|
||||
d.mu.Unlock()
|
||||
return
|
||||
}
|
||||
d.stopped = true
|
||||
d.mu.Unlock()
|
||||
|
||||
if d.cancel != nil {
|
||||
d.cancel()
|
||||
}
|
||||
if d.cmd != nil && d.cmd.Process != nil {
|
||||
// 进程活着才 kill;kill 失败也无所谓,服务很快就没了
|
||||
_ = d.cmd.Process.Kill()
|
||||
}
|
||||
if d.done != nil {
|
||||
// 等 Wait 协程回收进程,避免僵尸进程;进程已死则立刻返回
|
||||
<-d.done
|
||||
}
|
||||
d.teardown()
|
||||
if d.scriptDir != "" {
|
||||
_ = os.RemoveAll(d.scriptDir)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package nsfw
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 1x1 透明 PNG,当测试图片用(NudeNet 会自己缩放)。
|
||||
const testPNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
||||
|
||||
// TestDetectRestart 验证:python 检测服务中途崩溃后,Detect 会自动重启并恢复正常。
|
||||
func TestDetectRestart(t *testing.T) {
|
||||
py := resolvePython("")
|
||||
if py == "" {
|
||||
t.Skip("无 python")
|
||||
}
|
||||
img := filepath.Join(t.TempDir(), "test.png")
|
||||
if data, err := base64.StdEncoding.DecodeString(testPNG); err != nil {
|
||||
t.Fatalf("解码测试图片失败: %v", err)
|
||||
} else if err := os.WriteFile(img, data, 0644); err != nil {
|
||||
t.Fatalf("写测试图片失败: %v", err)
|
||||
}
|
||||
|
||||
d, err := Start(py, 0.5)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
defer d.Stop()
|
||||
|
||||
if _, _, err := d.Detect(img); err != nil {
|
||||
t.Fatalf("崩溃前检测失败: %v", err)
|
||||
}
|
||||
|
||||
// 直接杀掉 python 进程,模拟处理某张图时崩溃
|
||||
if d.cmd != nil && d.cmd.Process != nil {
|
||||
_ = d.cmd.Process.Kill()
|
||||
}
|
||||
time.Sleep(2 * time.Second) // 等进程退出、端口释放
|
||||
|
||||
if _, _, err := d.Detect(img); err != nil {
|
||||
t.Fatalf("崩溃后检测失败(应自动重启): %v", err)
|
||||
}
|
||||
t.Log("重启后检测成功")
|
||||
}
|
||||
+81
-21
@@ -1,8 +1,10 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""NudeNet 本地检测服务:由 Go 侧启动,常驻提供 /health 和 /detect。
|
||||
|
||||
只依赖 Python 标准库 + nudenet 包,模型在进程里只加载一次。
|
||||
收到 /detect 就把图片喂给 NudeDetector,返回结果 JSON。
|
||||
只依赖 Python 标准库 + nudenet 包。
|
||||
onnxruntime 的 InferenceSession 不是线程安全的,同一个实例并发推理会崩,
|
||||
所以这里用一个固定大小的推理线程池:每个 worker 线程各持一份独立模型实例,
|
||||
N 个请求可以真·并行推理。收到 /detect 就把图片喂给 NudeDetector,返回结果 JSON。
|
||||
|
||||
模型加载优先级(模型放在 <程序运行目录>/lib/320n.onnx):
|
||||
1. 本地已有该文件 → 直接加载(离线可用,适合随程序分发);
|
||||
@@ -15,12 +17,15 @@ Windows 超长路径说明:OpenCV 的 imread 走老式 C API,路径超过 26
|
||||
再把字节直接交给 NudeNet(v3 的 detect 支持 bytes)——从内存解码,彻底
|
||||
绕开路径限制,同时兼容老版 v2(只收路径时落到临时短路径文件再试)。
|
||||
"""
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import urllib.request
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
from nudenet import NudeDetector
|
||||
|
||||
@@ -49,6 +54,17 @@ MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""
|
||||
|
||||
_detector = None
|
||||
|
||||
# 推理并发数:onnxruntime 的 InferenceSession.run() 不是线程安全的,
|
||||
# 同一个实例并发推理会崩,所以每个 worker 线程各持有一个独立实例(NudeDetector),
|
||||
# N 个 worker 即可真·并行推理。默认 4,可用环境变量 KM_NSFW_WORKERS 调整。
|
||||
INFER_WORKERS = int(os.environ.get("KM_NSFW_WORKERS", "4") or "4")
|
||||
if INFER_WORKERS < 1:
|
||||
INFER_WORKERS = 1
|
||||
|
||||
_executor = None # 推理线程池(固定 N 个 worker,模型各持一份)
|
||||
_thread_local = threading.local() # 每个 worker 线程自己的检测器
|
||||
_MODEL_PATH = None # 启动时解析好的模型路径(None 表示用包自带模型)
|
||||
|
||||
|
||||
def _looks_like_model(path):
|
||||
"""粗略校验文件是不是真 ONNX 模型:开头不是 HTML、大小不小于 1MB(320n 模型约 11.6MB)。
|
||||
@@ -65,6 +81,13 @@ def _looks_like_model(path):
|
||||
return False
|
||||
|
||||
|
||||
def _download_model(local):
|
||||
"""带超时地把模型下载到 local。urlretrieve 没有超时,网络挂起会卡死整个服务。"""
|
||||
req = urllib.request.Request(MODEL_URL)
|
||||
with urllib.request.urlopen(req, timeout=30) as resp, open(local, "wb") as f:
|
||||
shutil.copyfileobj(resp, f)
|
||||
|
||||
|
||||
def _ensure_model(model_dir):
|
||||
"""返回要加载的模型路径;返回 None 表示用 nudenet 包自带的模型。"""
|
||||
if not model_dir:
|
||||
@@ -73,11 +96,11 @@ def _ensure_model(model_dir):
|
||||
if os.path.isfile(local) and _looks_like_model(local):
|
||||
print("=== 使用本地模型: %s" % local, flush=True)
|
||||
return local
|
||||
# 本地没有(或文件无效)→ 下载到 lib 目录(仅第一次需要网络)
|
||||
# 本地没有(或文件无效)→ 下载到 lib 目录(仅第一次需要网络,带 30 秒超时)
|
||||
try:
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
print("=== 本地未找到有效模型, 正在下载: %s" % local, flush=True)
|
||||
urllib.request.urlretrieve(MODEL_URL, local)
|
||||
_download_model(local)
|
||||
if _looks_like_model(local):
|
||||
print("=== 模型下载完成: %s" % local, flush=True)
|
||||
return local
|
||||
@@ -93,22 +116,33 @@ def _ensure_model(model_dir):
|
||||
|
||||
|
||||
def get_detector():
|
||||
global _detector
|
||||
"""确保模型已就绪(下载/加载),返回一个检测器实例。
|
||||
|
||||
启动时调用一次:把模型下载/解析好,/health 就绪即代表模型可用。
|
||||
真正推理用的独立实例由 _new_detector() 按线程各建一份,共享这份模型路径。
|
||||
"""
|
||||
global _detector, _MODEL_PATH
|
||||
if _detector is None:
|
||||
mp = _ensure_model(MODEL_DIR)
|
||||
try:
|
||||
_detector = NudeDetector(model_path=mp)
|
||||
except Exception as exc: # noqa: BLE001 - 指定路径加载失败就回退包自带模型,别崩
|
||||
print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True)
|
||||
if mp:
|
||||
try:
|
||||
os.remove(mp)
|
||||
except OSError:
|
||||
pass
|
||||
_detector = NudeDetector()
|
||||
_MODEL_PATH = _ensure_model(MODEL_DIR)
|
||||
_detector = _new_detector()
|
||||
return _detector
|
||||
|
||||
|
||||
def _new_detector():
|
||||
"""创建一份独立检测器实例(每个推理 worker 各持一份,可并行推理)。"""
|
||||
mp = _MODEL_PATH
|
||||
try:
|
||||
return NudeDetector(model_path=mp) if mp else NudeDetector()
|
||||
except Exception as exc: # noqa: BLE001 - 指定路径加载失败就回退包自带模型,别崩
|
||||
print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True)
|
||||
if mp:
|
||||
try:
|
||||
os.remove(mp)
|
||||
except OSError:
|
||||
pass
|
||||
return NudeDetector()
|
||||
|
||||
|
||||
def _longpath(path):
|
||||
"""Windows 下给绝对路径加 \\\\?\\ 前缀,让 Python 原生 open 也能处理超长路径。"""
|
||||
if os.name != "nt":
|
||||
@@ -136,7 +170,10 @@ def _read_bytes(path):
|
||||
|
||||
|
||||
def _detect_bytes(det, data, path):
|
||||
"""把图片字节喂给 NudeNet。v3 直接吃 bytes;老版 v2 只收路径时落到临时短路径重试。"""
|
||||
"""把图片字节喂给 NudeNet。v3 直接吃 bytes;老版 v2 只收路径时落到临时短路径重试。
|
||||
|
||||
用线程自己的检测器实例推理,各实例互不干扰,可以真·并行。
|
||||
"""
|
||||
try:
|
||||
return det.detect(data)
|
||||
except (TypeError, AttributeError):
|
||||
@@ -154,6 +191,19 @@ def _detect_bytes(det, data, path):
|
||||
pass
|
||||
|
||||
|
||||
def _infer(path):
|
||||
"""整个推理任务(读文件 + 检测),提交到固定线程池执行。
|
||||
|
||||
读文件在池外并行进行;推理用线程各自的独立检测器,N 个 worker 真·并行。
|
||||
"""
|
||||
data = _read_bytes(path)
|
||||
det = getattr(_thread_local, "detector", None)
|
||||
if det is None:
|
||||
det = _new_detector()
|
||||
_thread_local.detector = det
|
||||
return _detect_bytes(det, data, path)
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def _send(self, obj, status=200):
|
||||
data = json.dumps(obj, ensure_ascii=False).encode("utf-8")
|
||||
@@ -179,8 +229,8 @@ class Handler(BaseHTTPRequestHandler):
|
||||
threshold = float(body.get("threshold", 0.5))
|
||||
if not path:
|
||||
return self._send({"ok": False, "error": "缺少图片路径"})
|
||||
data = _read_bytes(path)
|
||||
detections = _detect_bytes(get_detector(), data, path)
|
||||
# 丢给推理线程池:读文件并行、推理由 N 个 worker 并行处理
|
||||
detections = _executor.submit(_infer, path).result()
|
||||
except Exception as exc: # noqa: BLE001 - 服务端任何异常都要回给 Go,别让进程崩
|
||||
return self._send({"ok": False, "error": str(exc)})
|
||||
|
||||
@@ -200,9 +250,19 @@ class Handler(BaseHTTPRequestHandler):
|
||||
pass
|
||||
|
||||
|
||||
# 多线程接收连接 + 大监听队列:Go 侧会并发发请求(worker 数可到 8),
|
||||
# 单线程 HTTPServer 默认 backlog 只有 5,并发一高连接就被内核拒绝(报 10061 连接拒绝)。
|
||||
# ThreadingHTTPServer 让连接即时被 accept;推理再交给固定线程池并行执行
|
||||
# (每个 worker 一份独立模型实例,onnxruntime 同一实例并发推理会崩)。
|
||||
class Server(ThreadingHTTPServer):
|
||||
request_queue_size = 128
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
|
||||
# 启动前先把模型准备好(本地加载 / 下载 / 回退),/health 就绪即代表模型可用,
|
||||
# 这样 Go 侧 waitReady 等到的就是"模型已就绪"的服务。
|
||||
get_detector()
|
||||
HTTPServer(("127.0.0.1", port), Handler).serve_forever()
|
||||
# 固定大小推理线程池:每个 worker 线程持一份独立模型实例,可并行推理
|
||||
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
|
||||
Server(("127.0.0.1", port), Handler).serve_forever()
|
||||
+180
-20
@@ -18,6 +18,7 @@ import (
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"kenaito-media-picker/internal/classify"
|
||||
"kenaito-media-picker/internal/config"
|
||||
@@ -48,6 +49,13 @@ const (
|
||||
dirAudio = "音频" // 音频统一放这一个文件夹
|
||||
)
|
||||
|
||||
// nsfwProgressInterval 色情识别每隔多少张图打一行进度,让前端能看到识别在推进
|
||||
// (逐张打太吵,日志还会被 trim 掉一半)。
|
||||
const nsfwProgressInterval = 200
|
||||
|
||||
// scanLogFlushInterval 续传记录攒够多少条就落一次库(中断/出错时尽量少丢已归类的结果)。
|
||||
const scanLogFlushInterval = 200
|
||||
|
||||
// ============================ 整理器主体 ============================
|
||||
|
||||
// Organizer 整理器。它的工作方式:
|
||||
@@ -69,12 +77,24 @@ type Organizer struct {
|
||||
onProgress func(phase string, done, total int)
|
||||
pending []*store.MediaInfo // 待写入数据库的媒体记录(攒一批再批量写)
|
||||
pendingNsfw []store.NsfwRecord // 待写入数据库的色情识别记录(攒一批再批量写)
|
||||
resume map[string]store.ScanLog // 续传表:源路径 → 上次归类结果(重扫跳过重新探测)
|
||||
resumeSkip int32 // 本次续传跳过的文件数(并行累计)
|
||||
pendingScan []store.ScanLog // 待写入续传表的记录(攒一批再批量写)
|
||||
nsfwChecked int32 // 色情识别已检查的图片数(并行累计,只涨不跌)
|
||||
nsfwHits int32 // 色情识别命中的图片数(并行累计)
|
||||
// 阶段耗时统计(并行 worker 累计,单位纳秒,只用于日志诊断瓶颈)
|
||||
tImageProbe int64 // 图片尺寸探测总耗时
|
||||
tImgMd5 int64 // 图片 MD5/指纹计算总耗时
|
||||
tBright int64 // 明暗分析总耗时
|
||||
tNsfw 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 {
|
||||
return &Organizer{
|
||||
org := &Organizer{
|
||||
pickerDir: pickerDir,
|
||||
ffprobePath: ffprobePath,
|
||||
ffmpegPath: ffmpegPath,
|
||||
@@ -84,6 +104,16 @@ func New(pickerDir, ffprobePath, ffmpegPath string, st *store.Store, ctx context
|
||||
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 设置进度回调。回调参数:
|
||||
@@ -166,15 +196,76 @@ func runParallel[T any, R any](ctx context.Context, items []T, fn func(T) R, onP
|
||||
|
||||
// ============================ 分类与入库小工具 ============================
|
||||
|
||||
// add 把一个文件记到某个目标目录的名下(先攒着,最后统一搬)。
|
||||
// add 把一个文件记到某个目标目录(相对分类名)的名下(先攒着,最后统一搬)。
|
||||
func (o *Organizer) add(f scanner.File, subDir string) {
|
||||
dir := filepath.Join(o.pickerDir, subDir)
|
||||
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)
|
||||
@@ -290,6 +381,8 @@ type imageProbeResult struct {
|
||||
psd bool // 是 PSD 设计稿,不做处理原样收走
|
||||
unsupported bool // 不支持的格式,归入"已忽略"
|
||||
probeErr bool // 体检失败(打不开/读不出),归入"读取异常"
|
||||
resume bool // 续传命中:上次已归类过,跳过全部体检,直接复用目标目录
|
||||
resumeDir string // 续传命中的目标目录
|
||||
res *probe.ProbeResult // 体检出的尺寸等信息(体检成功时才有)
|
||||
md5Str string // 文件 MD5
|
||||
headMd5 string // 文件头采样 MD5(大文件指纹)
|
||||
@@ -306,6 +399,13 @@ type imageProbeResult struct {
|
||||
// 只读数据库(查指纹),这是安全的。
|
||||
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
|
||||
@@ -317,7 +417,9 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
|
||||
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
|
||||
@@ -328,10 +430,14 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
|
||||
// 必须在算 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 {
|
||||
@@ -340,17 +446,25 @@ 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)
|
||||
atomic.AddInt64(&o.tNsfw, int64(time.Since(nsfwStart)))
|
||||
if derr != nil {
|
||||
// 检测失败按"未命中"处理,不影响整理流程
|
||||
fmt.Printf("=== 色情检测失败(按未命中归类): %s, 错误: %v \n", f.Path, derr)
|
||||
} else {
|
||||
r.nsfwScore = score
|
||||
r.nsfw = hit
|
||||
// 命中的图片把"文件名 + 检测分"打进运行日志,方便用户盯着看是否判对
|
||||
checked := atomic.AddInt32(&o.nsfwChecked, 1)
|
||||
if hit {
|
||||
atomic.AddInt32(&o.nsfwHits, 1)
|
||||
// 命中的图片把"文件名 + 检测分"打进运行日志,方便用户盯着看是否判对
|
||||
fmt.Printf("=== 色情命中: %s, 检测分=%.2f (阈值=%.2f)\n", f.Path, score, o.nsfw.Threshold())
|
||||
}
|
||||
// 每固定张数打一行进度,让前端能看到识别在推进(不逐张刷屏)
|
||||
if checked%nsfwProgressInterval == 0 {
|
||||
fmt.Printf("=== 色情识别进度: 已检查 %d 张, 命中 %d 张\n", checked, atomic.LoadInt32(&o.nsfwHits))
|
||||
}
|
||||
}
|
||||
}
|
||||
return r
|
||||
@@ -399,6 +513,10 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
success, fail, ignore, dup := 0, 0, 0, 0
|
||||
nsfwCount := 0 // 命中的色情图片数
|
||||
total := len(images)
|
||||
phaseStart := time.Now()
|
||||
if o.nsfw != nil {
|
||||
fmt.Printf("=== 色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval)
|
||||
}
|
||||
results := runParallel(o.ctx, images, o.probeImage, func(done int) {
|
||||
o.report(PhaseImages, done, total)
|
||||
})
|
||||
@@ -410,17 +528,21 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
}
|
||||
f := images[i]
|
||||
switch {
|
||||
case r.resume:
|
||||
// 续传命中:上次已归类过,直接把文件归回原目标目录(没搬走的这次搬)
|
||||
o.addToDir(f, r.resumeDir)
|
||||
success++
|
||||
case r.psd:
|
||||
// PSD 设计稿:不做处理,原样收进「图片_PSD」
|
||||
o.add(f, dirImagePSD)
|
||||
o.addClassified(f, dirImagePSD, r.size, r.headMd5, r.tailMd5)
|
||||
success++
|
||||
case r.unsupported:
|
||||
// 不支持的格式:原样收进「图片_已忽略」
|
||||
o.add(f, dirImageIgnored)
|
||||
o.addClassified(f, dirImageIgnored, r.size, r.headMd5, r.tailMd5)
|
||||
ignore++
|
||||
case r.probeErr:
|
||||
// 体检失败:收进「图片_读取异常」
|
||||
o.add(f, dirImageErr)
|
||||
o.addClassified(f, dirImageErr, r.size, r.headMd5, r.tailMd5)
|
||||
fail++
|
||||
fmt.Printf("=== 异常图片: %s \n", f.Path)
|
||||
case o.isDuplicate(r.md5Str, f.Path):
|
||||
@@ -432,7 +554,7 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
// 被判为色情的改去「图片_敏感_*」目录,明暗子目录照常保留
|
||||
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.addClassified(f, classify.ImageDirSensitive(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5)
|
||||
// 攒一条色情识别记录,最后连同"搬去哪了"一起写进数据库
|
||||
o.pendingNsfw = append(o.pendingNsfw, store.NsfwRecord{
|
||||
FilePath: f.Path,
|
||||
@@ -442,12 +564,23 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
})
|
||||
nsfwCount++
|
||||
} else {
|
||||
o.add(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore))
|
||||
o.addClassified(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), 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\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())
|
||||
if o.nsfw != nil && nsfwCount > 0 {
|
||||
fmt.Printf("=== 色情识别: 共命中 %d 张,已归入「图片_敏感_*」目录,可在【色情识别】页面查看并人工复核\n", nsfwCount)
|
||||
}
|
||||
@@ -457,18 +590,27 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
|
||||
|
||||
// mediaProbeResult 是"给一个视频/音频做完体检"的结论(并行阶段产出)。
|
||||
type mediaProbeResult struct {
|
||||
res *probe.ProbeResult // 体检结果(ffprobe 读出的信息)
|
||||
err error // 体检失败的原因
|
||||
md5 string // 文件 MD5
|
||||
head string // 文件头采样 MD5
|
||||
tail string // 文件尾采样 MD5
|
||||
size int64 // 文件大小
|
||||
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}
|
||||
}
|
||||
@@ -476,7 +618,9 @@ func (o *Organizer) probeMedia(f scanner.File) mediaProbeResult {
|
||||
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}
|
||||
}
|
||||
|
||||
@@ -488,6 +632,7 @@ func (o *Organizer) probeMedia(f scanner.File) mediaProbeResult {
|
||||
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)
|
||||
})
|
||||
@@ -498,17 +643,22 @@ func (o *Organizer) ProcessMedia(media []scanner.File) {
|
||||
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.add(f, dirVideoErr)
|
||||
o.addClassified(f, dirVideoErr, r.size, r.head, r.tail)
|
||||
fmt.Printf("=== 异常视频: %s \n", f.Path)
|
||||
} else if scanner.IsExplicitMediaMime(f.Mime) {
|
||||
ignore++
|
||||
o.add(f, dirVideoIgnored)
|
||||
o.addClassified(f, dirVideoIgnored, r.size, r.head, r.tail)
|
||||
} else {
|
||||
ignore++
|
||||
fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
|
||||
@@ -527,7 +677,7 @@ func (o *Organizer) ProcessMedia(media []scanner.File) {
|
||||
}
|
||||
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))
|
||||
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) {
|
||||
@@ -538,20 +688,29 @@ func (o *Organizer) ProcessMedia(media []scanner.File) {
|
||||
}
|
||||
audio++
|
||||
o.record(f, "audio", r.res, r.md5, r.head, r.tail, r.size)
|
||||
o.add(f, dirAudio)
|
||||
o.addClassified(f, dirAudio, r.size, r.head, r.tail)
|
||||
default:
|
||||
// 没有视频流也没有音频流 → 假媒体。
|
||||
// MIME 明确标注是媒体的,收进「已忽略」;否则原地不动。
|
||||
if scanner.IsExplicitMediaMime(f.Mime) {
|
||||
ignore++
|
||||
o.add(f, dirVideoIgnored)
|
||||
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())
|
||||
}
|
||||
|
||||
// ============================ 移动(最后一步) ============================
|
||||
@@ -562,6 +721,7 @@ func (o *Organizer) ProcessMedia(media []scanner.File) {
|
||||
// 不需要记录每个文件搬去哪了——文件就在分类目录里,一眼能看到。
|
||||
func (o *Organizer) Flush() {
|
||||
o.flushRecords()
|
||||
o.flushScanLog()
|
||||
|
||||
// 把 buckets 里所有文件拼成一张"搬运清单"
|
||||
jobs := make([]mover.MoveJob, 0)
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
package organize
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"kenaito-media-picker/internal/digest"
|
||||
"kenaito-media-picker/internal/scanner"
|
||||
"kenaito-media-picker/internal/store"
|
||||
)
|
||||
|
||||
func TestResumeMatchSize(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "a.jpg")
|
||||
os.WriteFile(p, []byte("hello world hello world"), 0644)
|
||||
st, _ := os.Stat(p)
|
||||
|
||||
org := New(filepath.Join(dir, "out"), "", "", nil, nil)
|
||||
org.resume = map[string]store.ScanLog{
|
||||
p: {SourcePath: p, TargetDir: filepath.Join(dir, "out", "图片_测试"), FileSize: st.Size()},
|
||||
}
|
||||
f := scanner.File{Path: p, Name: "a.jpg", Ext: ".jpg", Mime: "image/jpeg"}
|
||||
|
||||
sl, ok := org.matchResume(f)
|
||||
if !ok || sl.TargetDir != filepath.Join(dir, "out", "图片_测试") {
|
||||
t.Fatalf("内容未变应命中续传, ok=%v", ok)
|
||||
}
|
||||
|
||||
// 内容变大 → 大小不同 → 不命中
|
||||
os.WriteFile(p, []byte("hello world hello world extra bytes"), 0644)
|
||||
if _, ok := org.matchResume(f); ok {
|
||||
t.Fatal("内容变了(大小不同)应不命中续传")
|
||||
}
|
||||
|
||||
// 没记录 → 不命中
|
||||
org.resume = nil
|
||||
if _, ok := org.matchResume(f); ok {
|
||||
t.Fatal("无续传记录应不命中")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResumeMatchFingerprint(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "big.bin")
|
||||
data := make([]byte, 2<<20)
|
||||
for i := range data {
|
||||
data[i] = byte(i)
|
||||
}
|
||||
os.WriteFile(p, data, 0644)
|
||||
st, _ := os.Stat(p)
|
||||
head, tail, err := digest.CalcFileFingerprint(p, st.Size())
|
||||
if err != nil {
|
||||
t.Fatalf("fingerprint: %v", err)
|
||||
}
|
||||
|
||||
org := New(filepath.Join(dir, "out"), "", "", nil, nil)
|
||||
org.resume = map[string]store.ScanLog{
|
||||
p: {SourcePath: p, TargetDir: filepath.Join(dir, "out", "视频_测试"), FileSize: st.Size(), HeadMd5: head, TailMd5: tail},
|
||||
}
|
||||
f := scanner.File{Path: p, Name: "big.bin", Ext: ".bin", Mime: "video/mp4"}
|
||||
|
||||
if _, ok := org.matchResume(f); !ok {
|
||||
t.Fatal("指纹一致应命中续传")
|
||||
}
|
||||
|
||||
// 改文件头里一个字节、大小不变 → 指纹不同 → 不命中
|
||||
mod := append([]byte(nil), data...)
|
||||
mod[1024] ^= 0xFF
|
||||
os.WriteFile(p, mod, 0644)
|
||||
if _, ok := org.matchResume(f); ok {
|
||||
t.Fatal("内容变了(大小相同)应不命中续传")
|
||||
}
|
||||
}
|
||||
@@ -152,6 +152,23 @@ func analyzeBrightness(img image.Image) (darkScore float64, lightScore float64)
|
||||
offset++
|
||||
}
|
||||
}
|
||||
case *image.YCbCr:
|
||||
// JPEG 解码出来是 YCbCr:Y 平面就是亮度,直接翻数组。
|
||||
// 这是最常见的照片格式,走这分支能避开 default 里逐像素 img.At().RGBA() 的慢路径。
|
||||
pix := m.Y
|
||||
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
||||
offset := y*m.YStride + bounds.Min.X
|
||||
for x := bounds.Min.X; x < bounds.Max.X; x++ {
|
||||
v := float64(pix[offset])
|
||||
if v < 80 {
|
||||
darkCount++
|
||||
} else if v > 175 {
|
||||
lightCount++
|
||||
}
|
||||
sampleCount++
|
||||
offset++
|
||||
}
|
||||
}
|
||||
default:
|
||||
// 其他格式:用通用的读法
|
||||
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
||||
|
||||
@@ -3,15 +3,24 @@ package probe
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
)
|
||||
|
||||
// FFprobeExeName 释放出来的解码器文件名。程序启动时扫到它要跳过,
|
||||
// 免得把自己人(ffprobe)当成媒体文件收走了。
|
||||
// 只在 Windows(有内嵌 ffprobe.exe)时才会真正释放出来。
|
||||
const FFprobeExeName = "ffprobe.exe"
|
||||
|
||||
// FFprobePath 返回当前目录下 ffprobe 的完整路径。
|
||||
// systemFFprobe 回退用的系统 ffprobe 路径(非 Windows 平台、没有内嵌时启用)。
|
||||
var systemFFprobe string
|
||||
|
||||
// FFprobePath 返回当前要用的 ffprobe 完整路径:
|
||||
// 有内嵌就先释放到当前目录用内嵌的;没有内嵌则用系统 PATH 里的那个。
|
||||
func FFprobePath() string {
|
||||
return "." + string(os.PathSeparator) + FFprobeExeName
|
||||
if len(embeddedFFprobe) > 0 {
|
||||
return "." + string(os.PathSeparator) + FFprobeExeName
|
||||
}
|
||||
return systemFFprobe
|
||||
}
|
||||
|
||||
// HasEmbeddedFFprobe 当前编译出的程序里有没有内嵌 ffprobe。
|
||||
@@ -20,15 +29,23 @@ func HasEmbeddedFFprobe() bool {
|
||||
return len(embeddedFFprobe) > 0
|
||||
}
|
||||
|
||||
// ReleaseFFprobe 把内嵌的 ffprobe 释放(写)到当前目录,才能当程序调用它。
|
||||
// ReleaseFFprobe 准备出一个可用的 ffprobe:
|
||||
// - 有内嵌 → 释放(写)到当前目录,才能当程序调用它;
|
||||
// - 没有内嵌(如 Linux 构建)→ 找系统 PATH 里的 ffprobe(如 /usr/bin/ffprobe)。
|
||||
//
|
||||
// 两者都拿不到才报错:调用方会跳过视频/音频整理,图片照常处理。
|
||||
func ReleaseFFprobe() error {
|
||||
if len(embeddedFFprobe) == 0 {
|
||||
return fmt.Errorf("当前系统没有内嵌的ffprobe")
|
||||
if len(embeddedFFprobe) > 0 {
|
||||
return os.WriteFile(FFprobeExeName, embeddedFFprobe, 0777)
|
||||
}
|
||||
return os.WriteFile(FFprobeExeName, embeddedFFprobe, 0777)
|
||||
if p, err := exec.LookPath("ffprobe"); err == nil {
|
||||
systemFFprobe = p
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("没有内嵌 ffprobe,系统 PATH 里也找不到 ffprobe(Linux 请先安装 ffmpeg)")
|
||||
}
|
||||
|
||||
// RemoveFFprobe 整理完把临时释放的 ffprobe 删掉,不留在目录里。
|
||||
func RemoveFFprobe() {
|
||||
_ = os.Remove(FFprobeExeName)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
//go:build !windows
|
||||
|
||||
package probe
|
||||
|
||||
// 非 Windows 平台不在二进制里内嵌 ffprobe(仓库只带了 Windows 版)。
|
||||
// 运行时回退用系统 PATH 里的 ffprobe:如 CentOS 上 yum install ffmpeg 后
|
||||
// 就有 /usr/bin/ffprobe。找不到系统 ffprobe 时,视频/音频整理自动跳过,图片照常。
|
||||
var embeddedFFprobe []byte
|
||||
@@ -1,6 +1,8 @@
|
||||
//go:build windows
|
||||
|
||||
package probe
|
||||
|
||||
import _ "embed"
|
||||
|
||||
//go:embed bin/ffprobe.exe
|
||||
var embeddedFFprobe []byte // Windows 版 ffprobe(未入库,见 .gitignore)
|
||||
var embeddedFFprobe []byte // Windows 版 ffprobe(未入库,见 .gitignore)
|
||||
@@ -69,9 +69,25 @@ type ScanStatus struct {
|
||||
|
||||
// ============================ 应用级常量 ============================
|
||||
|
||||
// 应用级配置键。目前只存了一个:上次使用的存放目录,方便下次打开页面自动回填。
|
||||
// 应用级配置键。目前只存了两个:上次使用的存放目录、色情检测服务是否真的起来了。
|
||||
const ConfigKeyDestDir = "scan.destDir"
|
||||
|
||||
// nsfwDetectorReadyKey 记录最近一次扫描时色情检测服务是否真的起来了(python+nudenet 可用)。
|
||||
// 配置开了但检测器没起来(没装 python/模型下载失败等)也能区分出来,前端据此提示。
|
||||
const nsfwDetectorReadyKey = "nsfw.detectorReady"
|
||||
|
||||
// setNsfwDetectorReady 把"色情检测服务是否真的起来了"记进数据库(AppConfig)。
|
||||
func (s *ScanService) setNsfwDetectorReady(ready bool) {
|
||||
if s.store == nil {
|
||||
return
|
||||
}
|
||||
v := "0"
|
||||
if ready {
|
||||
v = "1"
|
||||
}
|
||||
_ = s.store.SaveConfig(nsfwDetectorReadyKey, v)
|
||||
}
|
||||
|
||||
// DefaultPickerDirName 用户没指定存放目录时,默认在扫描目录下新建这个文件夹装整理好的文件。
|
||||
const DefaultPickerDirName = "KMP已整理"
|
||||
|
||||
@@ -469,10 +485,11 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro
|
||||
res.Filter()
|
||||
|
||||
// ---------- 第三步:准备解码器 ----------
|
||||
// 读取视频/音频信息需要 ffprobe。程序里内嵌了一份,用之前先释放到当前目录。
|
||||
// 释放失败就跳过视频/音频整理,图片照常处理。
|
||||
// 读取视频/音频信息需要 ffprobe。Windows 构建内嵌了一份,先释放到当前目录;
|
||||
// Linux 构建没有内嵌,自动回退用系统 PATH 里的 ffprobe(如 /usr/bin/ffprobe)。
|
||||
// 两者都没有就跳过视频/音频整理,图片照常处理。
|
||||
ffprobePath := ""
|
||||
if (len(res.Images) > 0 || len(res.Media) > 0) && probe.HasEmbeddedFFprobe() {
|
||||
if len(res.Images) > 0 || len(res.Media) > 0 {
|
||||
if err := probe.ReleaseFFprobe(); err != nil {
|
||||
fmt.Printf("=== 释放解码器失败, 视频/音频整理将跳过(图片照常处理): %v\n", err)
|
||||
} else {
|
||||
@@ -498,12 +515,14 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro
|
||||
// 用本机 python + NudeNet 起一个本地检测服务。没装 python/nudenet 时跳过,
|
||||
// 不影响图片整理——检测结果只决定走普通目录还是「图片_敏感」目录。
|
||||
nsfwDet := (*nsfw.Detector)(nil)
|
||||
s.setNsfwDetectorReady(false)
|
||||
if cfg := config.GetAppConfig(); cfg.NSFWEnabled {
|
||||
if d, err := nsfw.Start(cfg.NSFWPython, cfg.NSFWThreshold); err != nil {
|
||||
fmt.Printf("=== 色情识别未启用(图片照常处理): %v\n", err)
|
||||
} else {
|
||||
nsfwDet = d
|
||||
defer nsfwDet.Stop()
|
||||
s.setNsfwDetectorReady(true)
|
||||
fmt.Println("=== 色情识别已启用(NudeNet)")
|
||||
// 清掉上一次的色情识别记录,保证「色情识别」页面显示的是本次扫描的结果
|
||||
if s.store != nil {
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSaveNsfwBatchLarge(t *testing.T) {
|
||||
db := filepath.Join(t.TempDir(), "test.db")
|
||||
s, err := Open(db)
|
||||
if err != nil {
|
||||
t.Fatalf("Open: %v", err)
|
||||
}
|
||||
defer func() {
|
||||
s.Close()
|
||||
os.Remove(db)
|
||||
}()
|
||||
|
||||
// 3000 条,远超 SQLite 999 变量上限(旧代码会报 too many SQL variables)
|
||||
records := make([]*NsfwRecord, 0, 3000)
|
||||
for i := 0; i < 3000; i++ {
|
||||
records = append(records, &NsfwRecord{
|
||||
FilePath: fmt.Sprintf("/path/%d.jpg", i),
|
||||
FileName: fmt.Sprintf("%d.jpg", i),
|
||||
Score: 0.5,
|
||||
Threshold: 0.2,
|
||||
MovedPath: fmt.Sprintf("/moved/%d.jpg", i),
|
||||
CreatedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
if err := s.SaveNsfwBatch(records); err != nil {
|
||||
t.Fatalf("SaveNsfwBatch: %v", err)
|
||||
}
|
||||
list, err := s.ListNsfw()
|
||||
if err != nil {
|
||||
t.Fatalf("ListNsfw: %v", err)
|
||||
}
|
||||
if len(list) != 3000 {
|
||||
t.Fatalf("got %d records, want 3000", len(list))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestScanLogRoundTrip(t *testing.T) {
|
||||
s, err := Open(filepath.Join(t.TempDir(), "t.db"))
|
||||
if err != nil {
|
||||
t.Fatalf("Open: %v", err)
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
if err := s.SaveScanLogs([]ScanLog{
|
||||
{SourcePath: "/a.jpg", TargetDir: "/out/图片_A", FileSize: 100},
|
||||
{SourcePath: "/b.mp4", TargetDir: "/out/视频_B", FileSize: 200},
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveScanLogs: %v", err)
|
||||
}
|
||||
|
||||
m, err := s.LoadScanLog()
|
||||
if err != nil {
|
||||
t.Fatalf("LoadScanLog: %v", err)
|
||||
}
|
||||
if len(m) != 2 || m["/a.jpg"].TargetDir != "/out/图片_A" {
|
||||
t.Fatalf("加载结果不对: %v", m)
|
||||
}
|
||||
|
||||
// 同路径 upsert:覆盖旧目标目录与大小,不新增行
|
||||
if err := s.SaveScanLogs([]ScanLog{
|
||||
{SourcePath: "/a.jpg", TargetDir: "/out/图片_A2", FileSize: 99},
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveScanLogs upsert: %v", err)
|
||||
}
|
||||
m2, _ := s.LoadScanLog()
|
||||
if len(m2) != 2 || m2["/a.jpg"].TargetDir != "/out/图片_A2" || m2["/a.jpg"].FileSize != 99 {
|
||||
t.Fatalf("upsert 结果不对: %v", m2)
|
||||
}
|
||||
}
|
||||
+87
-6
@@ -75,6 +75,20 @@ type NsfwRecord struct {
|
||||
CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 检测时间
|
||||
}
|
||||
|
||||
// ScanLog 一条"续传"记录:某个源文件上次被归类到了哪个目录。
|
||||
// 用途:扫描中断/出错后重跑时,已经归类过的文件直接复用上次的目标目录,
|
||||
// 不再重新探测尺寸 / 算 MD5 / 分析明暗 / 色情识别,省下几小时的重扫时间。
|
||||
// 大文件记录头尾指纹,重扫时核对内容没变才放心跳过。
|
||||
type ScanLog struct {
|
||||
Id int64 `xorm:"pk autoincr 'id'" json:"id"`
|
||||
SourcePath string `xorm:"varchar(1024) unique 'source_path'" json:"sourcePath"` // 源文件路径(唯一)
|
||||
TargetDir string `xorm:"text 'target_dir'" json:"targetDir"` // 归类结果:目标目录
|
||||
FileSize int64 `xorm:"bigint 'file_size'" json:"fileSize"` // 文件大小(校验内容是否变化)
|
||||
HeadMd5 string `xorm:"char(32) 'head_md5'" json:"headMd5"` // 文件头 256KB MD5(大文件指纹)
|
||||
TailMd5 string `xorm:"char(32) 'tail_md5'" json:"tailMd5"` // 文件尾 256KB MD5(大文件指纹)
|
||||
UpdatedAt time.Time `xorm:"updated 'updated_at'" json:"updatedAt"` // 最近归类时间
|
||||
}
|
||||
|
||||
// Open 打开(或第一次时创建)数据库,并保证表结构齐全。
|
||||
func Open(dbPath string) (*Store, error) {
|
||||
e, err := xorm.NewEngine("sqlite3", dbPath)
|
||||
@@ -82,10 +96,18 @@ func Open(dbPath string) (*Store, error) {
|
||||
return nil, err
|
||||
}
|
||||
e.ShowSQL(false)
|
||||
// SQLite 同一时刻只允许一个写连接,这里限制成 1 个,避免并发写冲突
|
||||
e.SetMaxOpenConns(1)
|
||||
// 自动创建/补全三张表(媒体记录表、配置表、色情识别记录表)
|
||||
if err := e.Sync(new(MediaInfo), new(AppConfig), new(NsfwRecord)); err != nil {
|
||||
// SQLite 同一时刻只允许一个写连接。这里把读连接放开几路:
|
||||
// 扫描探测阶段是并发读(查指纹),被 1 个连接串行化会拖慢并行;
|
||||
// 写仍走单事务(批量入库),由 SQLite 自己的锁机制串行,不会冲突。
|
||||
e.SetMaxOpenConns(8)
|
||||
// 给数据库加个忙等待,万一遇到并发写也不至于立刻报 SQLITE_BUSY
|
||||
e.SetMaxIdleConns(8)
|
||||
if _, err := e.Exec("PRAGMA busy_timeout = 10000"); err != nil {
|
||||
_ = e.Close()
|
||||
return nil, err
|
||||
}
|
||||
// 自动创建/补全四张表(媒体记录表、配置表、色情识别记录表、续传表)
|
||||
if err := e.Sync(new(MediaInfo), new(AppConfig), new(NsfwRecord), new(ScanLog)); err != nil {
|
||||
_ = e.Close()
|
||||
return nil, err
|
||||
}
|
||||
@@ -138,12 +160,29 @@ func (s *Store) GetConfig(key string) string {
|
||||
// ============================ 色情识别记录 ============================
|
||||
|
||||
// SaveNsfwBatch 批量保存色情识别记录(一次扫描攒一批,最后统一写入)。
|
||||
// 一次扫描可能命中几千张,一条 INSERT 塞不下(SQLite 单条语句最多 999 个变量),
|
||||
// 所以按 chunkSize 一组一组插,外面套一个事务保证要么全成功要么全回滚。
|
||||
func (s *Store) SaveNsfwBatch(records []*NsfwRecord) error {
|
||||
if s == nil || s.engine == nil || len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := s.engine.Insert(records)
|
||||
return err
|
||||
const chunkSize = 100
|
||||
sess := s.engine.NewSession()
|
||||
defer sess.Close()
|
||||
if err := sess.Begin(); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := 0; i < len(records); i += chunkSize {
|
||||
end := i + chunkSize
|
||||
if end > len(records) {
|
||||
end = len(records)
|
||||
}
|
||||
if _, err := sess.Insert(records[i:end]); err != nil {
|
||||
_ = sess.Rollback()
|
||||
return err
|
||||
}
|
||||
}
|
||||
return sess.Commit()
|
||||
}
|
||||
|
||||
// ClearNsfw 清空色情识别记录。每次开启色情识别的扫描开始时调用,
|
||||
@@ -156,6 +195,48 @@ func (s *Store) ClearNsfw() error {
|
||||
return err
|
||||
}
|
||||
|
||||
// ============================ 续传表 ============================
|
||||
|
||||
// LoadScanLog 把全部续传记录读出来,按源路径建索引,供扫描时快速判断"这个文件上次归到哪了"。
|
||||
func (s *Store) LoadScanLog() (map[string]ScanLog, error) {
|
||||
m := make(map[string]ScanLog)
|
||||
if s == nil || s.engine == nil {
|
||||
return m, nil
|
||||
}
|
||||
list := make([]ScanLog, 0)
|
||||
if err := s.engine.Find(&list); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, e := range list {
|
||||
m[e.SourcePath] = e
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// SaveScanLogs 批量保存续传记录(按源路径 upsert)。攒一批一起写,少开事务。
|
||||
// 单行删除 + 单行插入,没有 999 变量的限制,一批全在一个事务里。
|
||||
func (s *Store) SaveScanLogs(entries []ScanLog) error {
|
||||
if s == nil || s.engine == nil || len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
sess := s.engine.NewSession()
|
||||
defer sess.Close()
|
||||
if err := sess.Begin(); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, e := range entries {
|
||||
if _, err := sess.Where("source_path = ?", e.SourcePath).Delete(new(ScanLog)); err != nil {
|
||||
_ = sess.Rollback()
|
||||
return err
|
||||
}
|
||||
if _, err := sess.Insert(&e); err != nil {
|
||||
_ = sess.Rollback()
|
||||
return err
|
||||
}
|
||||
}
|
||||
return sess.Commit()
|
||||
}
|
||||
|
||||
// ListNsfw 查全部色情识别记录,按检测分从高到低排(方便先看最"实锤"的)。
|
||||
func (s *Store) ListNsfw() ([]NsfwRecord, error) {
|
||||
if s == nil || s.engine == nil {
|
||||
|
||||
Reference in New Issue
Block a user