odboy/kenaito-media-picker
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

upd

Browse Source
This commit is contained in:
odboy
2025-09-15 21:17:57 +08:00
parent 524da0830a
commit 413ac308ac
9 changed files with 394 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/anime_nsfw/detector.go
+11 -5
View File
@@ -46,6 +46,7 @@ const requestTimeout = 60 * time.Second
type Detector struct {
url string
threshold float64
workers int
client *http.Client
cmd *exec.Cmd
scriptDir string
@@ -59,10 +60,13 @@ type Detector struct {
}
// Start 拉起本地 python 检测服务并等它就绪。
func Start(pythonPath string, threshold float64) (*Detector, error) {
func Start(pythonPath string, threshold float64, workers int) (*Detector, error) {
if threshold <= 0 || threshold >= 1 {
threshold = 0.5
}
if workers <= 0 {
workers = 2 // 默认 2 个 worker
}
py := resolvePython(pythonPath)
if py == "" {
return nil, fmt.Errorf("未找到 python，动漫敏感识别不可用")
@@ -85,6 +89,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
d := &Detector{
threshold: threshold,
workers: workers,
client: &http.Client{Timeout: requestTimeout},
python: py,
modelDir: modelDir,
@@ -108,6 +113,7 @@ func (d *Detector) launch() error {
ctx, cancel := context.WithCancel(context.Background())
cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
cmd.Env = append(os.Environ(), fmt.Sprintf("KM_ANIME_NSFW_WORKERS=%d", d.workers))
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stdout
if err := cmd.Start(); err != nil {
@@ -256,10 +262,10 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, ratings An
Nsfw bool `json:"nsfw"`
Score float64 `json:"score"`
Ratings struct {
General float64 `json:"general"`
Sensitive float64 `json:"sensitive"`
Questionable float64 `json:"questionable"`
Explicit float64 `json:"explicit"`
General float64 `json:"general"`
Sensitive float64 `json:"sensitive"`
Questionable float64 `json:"questionable"`
Explicit float64 `json:"explicit"`
} `json:"ratings"`
TopTags []map[string]interface{} `json:"top_tags"`
}
internal/api/config_controller.go
+85 -5
View File
@@ -8,15 +8,19 @@ import (
"github.com/gin-gonic/gin"
"kenaito-media-picker/internal/store"
"kenaito-media-picker/internal/sysinfo"
)
// 配置键名（存储在 AppConfig 表中）
const (
ConfigKeyNsfwType = "nsfw.type" // 检测类型："" / "human" / "anime"
ConfigKeyNsfwThreshold = "nsfw.threshold" // 人类检测阈值
ConfigKeyAnimeNsfwThreshold = "anime.threshold" // 二次元检测阈值
ConfigKeyScanDir = "scan.scanDir" // 扫描目录
ConfigKeyDestDir = "scan.destDir" // 存放目录
ConfigKeyNsfwType = "nsfw.type" // 检测类型："" / "human" / "anime"
ConfigKeyNsfwThreshold = "nsfw.threshold" // 人类检测阈值
ConfigKeyAnimeNsfwThreshold = "anime.threshold" // 二次元检测阈值
ConfigKeyNsfwWorkers = "nsfw.workers" // 人类检测 worker 数量
ConfigKeyAnimeNsfwWorkers = "anime.workers" // 二次元检测 worker 数量
ConfigKeyOrganizeWorkers = "organize.workers" // Go 端探测/MIME/移动 worker 数量
ConfigKeyScanDir = "scan.scanDir" // 扫描目录
ConfigKeyDestDir = "scan.destDir" // 存放目录
)
// ConfigController 配置相关接口。
@@ -34,6 +38,9 @@ type NsfwConfig struct {
Type string `json:"type"` // 检测类型："" / "human" / "anime"
HumanThreshold float64 `json:"humanThreshold"` // 人类检测阈值
AnimeThreshold float64 `json:"animeThreshold"` // 二次元检测阈值
HumanWorkers int `json:"humanWorkers"` // 人类检测 worker 数量
AnimeWorkers int `json:"animeWorkers"` // 二次元检测 worker 数量
GoWorkers int `json:"goWorkers"` // Go 端 worker 数量
}
// GetNsfwConfig 处理 GET /api/config/nsfw：获取 NSFW 配置。
@@ -54,10 +61,34 @@ func (c *ConfigController) GetNsfwConfig(ctx *gin.Context) {
}
}
humanWorkers := 4
if v := c.store.GetConfig(ConfigKeyNsfwWorkers); v != "" {
if w, err := strconv.Atoi(v); err == nil && w > 0 {
humanWorkers = w
}
}
animeWorkers := 2
if v := c.store.GetConfig(ConfigKeyAnimeNsfwWorkers); v != "" {
if w, err := strconv.Atoi(v); err == nil && w > 0 {
animeWorkers = w
}
}
goWorkers := 0 // 0 表示使用文件配置或自动计算
if v := c.store.GetConfig(ConfigKeyOrganizeWorkers); v != "" {
if w, err := strconv.Atoi(v); err == nil && w >= 0 {
goWorkers = w
}
}
Success(ctx, NsfwConfig{
Type: nsfwType,
HumanThreshold: humanThreshold,
AnimeThreshold: animeThreshold,
HumanWorkers: humanWorkers,
AnimeWorkers: animeWorkers,
GoWorkers: goWorkers,
})
}
@@ -66,6 +97,9 @@ type UpdateNsfwConfigReq struct {
Type string `json:"type"` // 检测类型："" / "human" / "anime"
HumanThreshold float64 `json:"humanThreshold"` // 人类检测阈值
AnimeThreshold float64 `json:"animeThreshold"` // 二次元检测阈值
HumanWorkers int `json:"humanWorkers"` // 人类检测 worker 数量
AnimeWorkers int `json:"animeWorkers"` // 二次元检测 worker 数量
GoWorkers int `json:"goWorkers"` // Go 端 worker 数量（0 表示使用文件配置）
}
// UpdateNsfwConfig 处理 PUT /api/config/nsfw：更新 NSFW 配置。
@@ -90,6 +124,17 @@ func (c *ConfigController) UpdateNsfwConfig(ctx *gin.Context) {
req.AnimeThreshold = 0.5
}
// 校验 worker 数量
if req.HumanWorkers <= 0 {
req.HumanWorkers = 4
}
if req.AnimeWorkers <= 0 {
req.AnimeWorkers = 2
}
if req.GoWorkers < 0 {
req.GoWorkers = 0
}
// 保存配置
if err := c.store.SaveConfig(ConfigKeyNsfwType, req.Type); err != nil {
Error(ctx, http.StatusInternalServerError, "保存配置失败："+err.Error())
@@ -103,10 +148,45 @@ func (c *ConfigController) UpdateNsfwConfig(ctx *gin.Context) {
Error(ctx, http.StatusInternalServerError, "保存配置失败："+err.Error())
return
}
if err := c.store.SaveConfig(ConfigKeyNsfwWorkers, strconv.Itoa(req.HumanWorkers)); err != nil {
Error(ctx, http.StatusInternalServerError, "保存配置失败："+err.Error())
return
}
if err := c.store.SaveConfig(ConfigKeyAnimeNsfwWorkers, strconv.Itoa(req.AnimeWorkers)); err != nil {
Error(ctx, http.StatusInternalServerError, "保存配置失败："+err.Error())
return
}
if err := c.store.SaveConfig(ConfigKeyOrganizeWorkers, strconv.Itoa(req.GoWorkers)); err != nil {
Error(ctx, http.StatusInternalServerError, "保存配置失败："+err.Error())
return
}
Success(ctx, nil)
}
// SystemInfo 系统信息响应。
type SystemInfo struct {
CPUCores int `json:"cpuCores"` // CPU 逻辑核心数
MemoryGB float64 `json:"memoryGB"` // 内存大小（GB）
RecommendedGo int `json:"recommendedGo"` // 推荐的 Go 端 worker 数
RecommendedNsfw int `json:"recommendedNsfw"` // 推荐的人类检测 worker 数
RecommendedAnime int `json:"recommendedAnime"` // 推荐的二次元检测 worker 数
}
// GetSystemInfo 处理 GET /api/config/system：获取系统信息和推荐配置。
func (c *ConfigController) GetSystemInfo(ctx *gin.Context) {
info := sysinfo.GetSystemInfo()
workers := sysinfo.CalculateWorkerConfig(info)
Success(ctx, SystemInfo{
CPUCores: info.CPUCores,
MemoryGB: info.MemoryGB,
RecommendedGo: workers.GoWorkers,
RecommendedNsfw: workers.NsfwWorkers,
RecommendedAnime: workers.AnimeNsfwWorkers,
})
}
// UpdateScanDirReq 更新扫描目录配置请求。
type UpdateScanDirReq struct {
ScanDir string `json:"scanDir"` // 扫描目录
internal/api/router.go
+1
View File
@@ -52,6 +52,7 @@ func NewRouter(st *store.Store) *gin.Engine {
cc := NewConfigController(st)
apiGroup.GET("/config/nsfw", cc.GetNsfwConfig)
apiGroup.PUT("/config/nsfw", cc.UpdateNsfwConfig)
apiGroup.GET("/config/system", cc.GetSystemInfo)
apiGroup.PUT("/config/scan-dir", cc.UpdateScanDir)
}
return r
internal/config/config.go
+8
View File
@@ -81,3 +81,11 @@ func WorkerCount() int {
}
return n
}
// SetWorkers 运行时更新 worker 数量（用于从数据库读取配置后覆盖文件配置）。
func SetWorkers(workers int) {
if workers < 0 {
workers = 0
}
appConfig.Workers = workers
}
internal/nsfw/detector.go
+9 -1
View File
@@ -60,6 +60,7 @@ const requestTimeout = 30 * time.Second
type Detector struct {
url string // 本地服务地址，如 http://127.0.0.1:54321
threshold float64 // 判定阈值 0~1
workers int // Python 推理 worker 数量
client *http.Client
cmd *exec.Cmd // python 进程句柄（Stop 时杀掉）
scriptDir string // 临时脚本所在目录（Stop 时清理）
@@ -75,12 +76,16 @@ type Detector struct {
// Start 拉起本地 python 检测服务并等它就绪。
// - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
// - threshold：判定阈值 0~1。
// - workers：Python 推理 worker 数量；<=0 表示使用默认值（4）。
//
// 找不到 python、装不了 nudenet、或服务起不来时返回错误（调用方跳过识别即可）。
func Start(pythonPath string, threshold float64) (*Detector, error) {
func Start(pythonPath string, threshold float64, workers int) (*Detector, error) {
if threshold <= 0 || threshold >= 1 {
threshold = 0.5 // 阈值传得不合理就用默认
}
if workers <= 0 {
workers = 4 // 默认 4 个 worker
}
// 找 python
py := resolvePython(pythonPath)
if py == "" {
@@ -107,6 +112,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
d := &Detector{
threshold: threshold,
workers: workers,
client: &http.Client{Timeout: requestTimeout},
python: py,
modelDir: modelDir,
@@ -135,6 +141,8 @@ func (d *Detector) launch() error {
ctx, cancel := context.WithCancel(context.Background())
cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
// 设置环境变量，传递 worker 数量给 Python 服务
cmd.Env = append(os.Environ(), fmt.Sprintf("KM_NSFW_WORKERS=%d", d.workers))
// 标准输出/错误都接到程序自身的输出：
// - 扫描期间 os.Stdout 会被接管进"运行日志"，这样 python 的报错（崩溃原因）也能在前端看到
// - 方便排查服务为什么挂掉
internal/service/scan.go
+27 -6
View File
@@ -74,8 +74,9 @@ type ScanStatus struct {
// 应用级配置键。目前只存了两个：上次使用的存放目录、色情检测服务是否真的起来了。
const (
ConfigKeyScanDir = "scan.scanDir"
ConfigKeyDestDir = "scan.destDir"
ConfigKeyScanDir = "scan.scanDir"
ConfigKeyDestDir = "scan.destDir"
ConfigKeyOrganizeWorkers = "organize.workers"
)
// nsfwDetectorReadyKey 记录最近一次扫描时色情检测服务是否真的起来了（python+nudenet 可用）。
@@ -625,6 +626,14 @@ func sanitizeLine(line string) string {
// 3. 准备解码器：把内嵌的 ffprobe 程序释放到当前目录（用它读视频/音频信息）
// 4. 整理：给每个文件探测信息 → 判断是否重复 → 分到分类目录 → 搬进去
func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) error {
// 从数据库读取 organize.workers 配置（支持运行时修改，立即生效）
if v := s.store.GetConfig(ConfigKeyOrganizeWorkers); v != "" {
if w, err := strconv.Atoi(v); err == nil && w >= 0 {
config.SetWorkers(w)
fmt.Printf("=== 使用数据库配置 organize.workers=%d\n", w)
}
}
// 确定整理后文件的存放目录：统一都放进「KMP已整理」这一层——
// 用户填了存放目录就放在它下面（存放目录/KMP已整理），没填就放在扫描目录下（扫描目录/KMP已整理）。
pickerDir := filepath.Join(rootDir, DefaultPickerDirName)
@@ -696,13 +705,19 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro
threshold = t
}
}
if d, err := nsfw.Start("", threshold); err != nil {
workers := 4
if v := s.store.GetConfig("nsfw.workers"); v != "" {
if w, err := strconv.Atoi(v); err == nil && w > 0 {
workers = w
}
}
if d, err := nsfw.Start("", threshold, workers); err != nil {
fmt.Printf("=== 色情识别未启用(图片照常处理): %v\n", err)
} else {
nsfwDet = d
defer nsfwDet.Stop()
s.setNsfwDetectorReady(true)
fmt.Println("=== 色情识别已启用(NudeNet)")
fmt.Printf("=== 色情识别已启用(NudeNet, workers=%d)\n", workers)
}
}
@@ -715,12 +730,18 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro
threshold = t
}
}
if d, err := anime_nsfw.Start("", threshold); err != nil {
workers := 2
if v := s.store.GetConfig("anime.workers"); v != "" {
if w, err := strconv.Atoi(v); err == nil && w > 0 {
workers = w
}
}
if d, err := anime_nsfw.Start("", threshold, workers); err != nil {
fmt.Printf("=== 动漫色情识别未启用(图片照常处理): %v\n", err)
} else {
animeNsfwDet = d
defer animeNsfwDet.Stop()
fmt.Println("=== 动漫色情识别已启用(WD Tagger v3)")
fmt.Printf("=== 动漫色情识别已启用(WD Tagger v3, workers=%d)\n", workers)
}
}
internal/sysinfo/sysinfo.go
+151
View File
@@ -0,0 +1,151 @@
// Package sysinfo 提供系统资源检测和 worker 数量自动计算功能。
package sysinfo
import (
"os/exec"
"runtime"
"strconv"
"strings"
)
// SystemInfo 系统资源信息
type SystemInfo struct {
CPUCores int // CPU 逻辑核心数
MemoryGB float64 // 内存大小（GB）
MemoryMB int // 内存大小（MB）
}
// GetSystemInfo 获取系统资源信息
func GetSystemInfo() SystemInfo {
info := SystemInfo{
CPUCores: runtime.NumCPU(),
}
// 获取内存信息
info.MemoryMB = getMemoryMB()
info.MemoryGB = float64(info.MemoryMB) / 1024.0
return info
}
// getMemoryMB 获取系统内存大小（MB）
func getMemoryMB() int {
switch runtime.GOOS {
case "windows":
return getMemoryWindows()
case "linux":
return getMemoryLinux()
case "darwin":
return getMemoryDarwin()
default:
return 0
}
}
// getMemoryWindows 使用 wmic 获取 Windows 内存
func getMemoryWindows() int {
cmd := exec.Command("wmic", "computersystem", "get", "totalphysicalmemory", "/value")
output, err := cmd.Output()
if err != nil {
return 0
}
// 解析输出：TotalPhysicalMemory=1234567890
for _, line := range strings.Split(string(output), "\n") {
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "TotalPhysicalMemory=") {
value := strings.TrimPrefix(line, "TotalPhysicalMemory=")
if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
return int(bytes / 1024 / 1024)
}
}
}
return 0
}
// getMemoryLinux 从 /proc/meminfo 获取 Linux 内存
func getMemoryLinux() int {
cmd := exec.Command("cat", "/proc/meminfo")
output, err := cmd.Output()
if err != nil {
return 0
}
// 解析 MemTotal: 12345678 kB
for _, line := range strings.Split(string(output), "\n") {
if strings.HasPrefix(line, "MemTotal:") {
fields := strings.Fields(line)
if len(fields) >= 2 {
if kb, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
return int(kb / 1024)
}
}
}
}
return 0
}
// getMemoryDarwin 使用 sysctl 获取 macOS 内存
func getMemoryDarwin() int {
cmd := exec.Command("sysctl", "-n", "hw.memsize")
output, err := cmd.Output()
if err != nil {
return 0
}
value := strings.TrimSpace(string(output))
if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
return int(bytes / 1024 / 1024)
}
return 0
}
// WorkerConfig worker 配置
type WorkerConfig struct {
GoWorkers int // Go 端 worker 数量（探测/MIME/移动）
NsfwWorkers int // NudeNet Python 推理 worker 数量
AnimeNsfwWorkers int // WD Tagger Python 推理 worker 数量
}
// CalculateWorkerConfig 根据系统资源自动计算 worker 配置
// 规则：
// - Go worker: CPU核心数 - 2（留2核给Python推理和系统）
// - NudeNet worker: 因为模型小（~300MB），可以开更多，上限为 CPU核心数
// - WD Tagger worker: 因为模型大（~1.5GB），受内存限制，每个worker约需2GB
func CalculateWorkerConfig(info SystemInfo) WorkerConfig {
config := WorkerConfig{}
// Go worker: CPU核心数 - 2，最少为1
goWorkers := info.CPUCores - 2
if goWorkers < 1 {
goWorkers = 1
}
config.GoWorkers = goWorkers
// NudeNet worker: 模型小，可以开到 CPU核心数，但留一些给系统
nsfwWorkers := info.CPUCores
if nsfwWorkers > 12 {
nsfwWorkers = 12 // 上限12，避免过度并发
}
config.NsfwWorkers = nsfwWorkers
// WD Tagger worker: 受内存限制，每个worker约需2GB
// 预留4GB给系统和其他进程
availableMemoryGB := info.MemoryGB - 4
if availableMemoryGB < 0 {
availableMemoryGB = 0
}
animeWorkers := int(availableMemoryGB / 2)
if animeWorkers < 1 {
animeWorkers = 1
}
if animeWorkers > info.CPUCores {
animeWorkers = info.CPUCores // 不超过CPU核心数
}
if animeWorkers > 12 {
animeWorkers = 12 // 上限12
}
config.AnimeNsfwWorkers = animeWorkers
return config
}