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:
2025-06-19 21:07:35 +08:00
parent ea134afa27
commit 957dfe1134
18 changed files with 908 additions and 111 deletions
+13 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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)
}
+47
View File
@@ -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
View File
@@ -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
View File
@@ -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)
+74
View File
@@ -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("内容变了(大小相同)应不命中续传")
}
}
+17
View File
@@ -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++ {
+24 -7
View File
@@ -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)
}
}
+8
View File
@@ -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)
+23 -4
View File
@@ -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 {
+44
View File
@@ -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))
}
}
+40
View File
@@ -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
View File
@@ -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 {