Tangledup-ai/sam3_local
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

admin update

Browse Source
This commit is contained in:
jeremygan2021
2026-02-17 12:17:23 +08:00
parent a4475e743d
commit f6c361dd13
3 changed files with 101 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

191
fastAPI_nocom.py
View File

@@ -1,191 +0,0 @@
import os
import uuid
import requests
from typing import Optional
from contextlib import asynccontextmanager
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends, status
from fastapi.security import APIKeyHeader
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from PIL import Image
from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor
from sam3.visualization_utils import plot_results
# ------------------- 配置与路径 -------------------
STATIC_DIR = "static"
RESULT_IMAGE_DIR = os.path.join(STATIC_DIR, "results")
os.makedirs(RESULT_IMAGE_DIR, exist_ok=True)
# ------------------- API Key 核心配置 (已加固) -------------------
VALID_API_KEY = "123quant-speed"
API_KEY_HEADER_NAME = "X-API-Key"
# 定义 Header 认证
api_key_header = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=False)
async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
"""
强制验证 API Key
"""
# 1. 检查是否有 Key
if not api_key:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Missing API Key. Please provide it in the header."
)
# 2. 检查 Key 是否正确
if api_key != VALID_API_KEY:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Invalid API Key."
)
# 3. 验证通过
return True
# ------------------- 生命周期管理 -------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
print("="*40)
print("✅ API Key 保护已激活")
print(f"✅ 有效 Key: {VALID_API_KEY}")
print("="*40)
print("正在加载 SAM3 模型到 GPU...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if not torch.cuda.is_available():
print("警告: 未检测到 GPU，将使用 CPU，速度会较慢。")
model = build_sam3_image_model()
model = model.to(device)
model.eval()
processor = Sam3Processor(model)
app.state.model = model
app.state.processor = processor
app.state.device = device
print(f"模型加载完成，设备: {device}")
yield
print("正在清理资源...")
# ------------------- FastAPI 初始化 -------------------
app = FastAPI(
lifespan=lifespan,
title="SAM3 Segmentation API",
description="## 🔒 受 API Key 保护\n请点击右上角 **Authorize** 并输入: `123quant-speed`",
)
# 手动添加 OpenAPI 安全配置，让 Docs 里的锁头生效
app.openapi_schema = None
def custom_openapi():
if app.openapi_schema:
return app.openapi_schema
from fastapi.openapi.utils import get_openapi
openapi_schema = get_openapi(
title=app.title,
version=app.version,
description=app.description,
routes=app.routes,
)
# 定义安全方案
openapi_schema["components"]["securitySchemes"] = {
"APIKeyHeader": {
"type": "apiKey",
"in": "header",
"name": API_KEY_HEADER_NAME,
}
}
# 为所有路径应用安全要求
for path in openapi_schema["paths"]:
for method in openapi_schema["paths"][path]:
openapi_schema["paths"][path][method]["security"] = [{"APIKeyHeader": []}]
app.openapi_schema = openapi_schema
return app.openapi_schema
app.openapi = custom_openapi
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
# ------------------- 辅助函数 -------------------
def load_image_from_url(url: str) -> Image.Image:
try:
headers = {'User-Agent': 'Mozilla/5.0'}
response = requests.get(url, headers=headers, stream=True, timeout=10)
response.raise_for_status()
image = Image.open(response.raw).convert("RGB")
return image
except Exception as e:
raise HTTPException(status_code=400, detail=f"无法下载图片: {str(e)}")
def generate_and_save_result(image: Image.Image, inference_state) -> str:
filename = f"seg_{uuid.uuid4().hex}.jpg"
save_path = os.path.join(RESULT_IMAGE_DIR, filename)
plot_results(image, inference_state)
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
return filename
# ------------------- API 接口 (强制依赖验证) -------------------
@app.post("/segment", dependencies=[Depends(verify_api_key)])
async def segment(
request: Request,
prompt: str = Form(...),
file: Optional[UploadFile] = File(None),
image_url: Optional[str] = Form(None)
):
if not file and not image_url:
raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)")
try:
if file:
image = Image.open(file.file).convert("RGB")
elif image_url:
image = load_image_from_url(image_url)
except Exception as e:
raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}")
processor = request.app.state.processor
try:
inference_state = processor.set_image(image)
output = processor.set_text_prompt(state=inference_state, prompt=prompt)
masks, boxes, scores = output["masks"], output["boxes"], output["scores"]
except Exception as e:
raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}")
try:
filename = generate_and_save_result(image, inference_state)
except Exception as e:
raise HTTPException(status_code=500, detail=f"绘图保存错误: {str(e)}")
file_url = request.url_for("static", path=f"results/{filename}")
return JSONResponse(content={
"status": "success",
"result_image_url": str(file_url),
"detected_count": len(masks),
"scores": scores.tolist() if torch.is_tensor(scores) else scores
})
if __name__ == "__main__":
import uvicorn
# 注意：如果你的文件名不是 fastAPI_nocom.py，请修改下面第一个参数
uvicorn.run(
"fastAPI_nocom:app",
host="127.0.0.1",
port=55600,
proxy_headers=True,
forwarded_allow_ips="*",
reload=False # 生产环境建议关闭 reload，确保代码完全重载
)

61
fastAPI_tarot.py
View File

@@ -73,6 +73,13 @@ dashscope.api_key = 'sk-ce2404f55f744a1987d5ece61c6bac58'
QWEN_MODEL = 'qwen-vl-max' # Default model
AVAILABLE_QWEN_MODELS = ["qwen-vl-max", "qwen-vl-plus"]
# 清理配置 (Cleanup Config)
CLEANUP_CONFIG = {
"enabled": os.getenv("AUTO_CLEANUP_ENABLED", "True").lower() == "true",
"lifetime": int(os.getenv("FILE_LIFETIME_SECONDS", "3600")),
"interval": int(os.getenv("CLEANUP_INTERVAL_SECONDS", "600"))
}
# API Tags (用于文档分类)
TAG_GENERAL = "General Segmentation (通用分割)"
TAG_TAROT = "Tarot Analysis (塔罗牌分析)"
@@ -108,14 +115,24 @@ async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
# 4. Lifespan Management (生命周期管理)
# ==========================================
async def cleanup_old_files(directory: str, lifetime_seconds: int, interval_seconds: int):
async def cleanup_old_files(directory: str):
"""
后台任务：定期清理过期的图片文件
- 动态读取 CLEANUP_CONFIG 配置
"""
print(f"🧹 自动清理任务已启动 | 目录: {directory} | 生命周期: {lifetime_seconds}s | 检查间隔: {interval_seconds}s")
print(f"🧹 自动清理任务已启动 | 目录: {directory}")
while True:
try:
await asyncio.sleep(interval_seconds)
# 动态读取配置
interval = CLEANUP_CONFIG["interval"]
lifetime = CLEANUP_CONFIG["lifetime"]
enabled = CLEANUP_CONFIG["enabled"]
await asyncio.sleep(interval)
if not enabled:
continue
current_time = time.time()
count = 0
# 遍历所有文件（包括子目录）
@@ -125,7 +142,7 @@ async def cleanup_old_files(directory: str, lifetime_seconds: int, interval_seco
# 获取文件修改时间
try:
file_mtime = os.path.getmtime(file_path)
if current_time - file_mtime > lifetime_seconds:
if current_time - file_mtime > lifetime:
os.remove(file_path)
count += 1
except OSError:
@@ -142,7 +159,7 @@ async def cleanup_old_files(directory: str, lifetime_seconds: int, interval_seco
pass
if count > 0:
print(f"🧹 已清理 {count} 个过期文件")
print(f"🧹 已清理 {count} 个过期文件 (Lifetime: {lifetime}s)")
except asyncio.CancelledError:
print("🛑 清理任务已停止")
@@ -185,13 +202,9 @@ async def lifespan(app: FastAPI):
# --- 启动后台清理任务 ---
cleanup_task_handle = None
# 优先读取环境变量，否则使用默认值
if os.getenv("AUTO_CLEANUP_ENABLED", "False").lower() == "true":
try:
lifetime = int(os.getenv("FILE_LIFETIME_SECONDS", "3600"))
interval = int(os.getenv("CLEANUP_INTERVAL_SECONDS", "600"))
cleanup_task_handle = asyncio.create_task(
cleanup_old_files(RESULT_IMAGE_DIR, lifetime, interval)
cleanup_old_files(RESULT_IMAGE_DIR)
)
except Exception as e:
print(f"启动清理任务失败: {e}")
@@ -1180,9 +1193,8 @@ async def trigger_cleanup():
Manually trigger cleanup
"""
try:
# Re-use logic from cleanup_old_files but force it for all files > 0 seconds if we want deep clean?
# Or just use the standard lifetime. Let's use standard lifetime but run it now.
lifetime = int(os.getenv("FILE_LIFETIME_SECONDS", "3600"))
# Use global config
lifetime = CLEANUP_CONFIG["lifetime"]
count = 0
current_time = time.time()
@@ -1207,7 +1219,7 @@ async def trigger_cleanup():
except:
pass
return {"status": "success", "message": f"Cleaned {count} files"}
return {"status": "success", "message": f"Cleaned {count} files (Lifetime: {lifetime}s)"}
except Exception as e:
return {"status": "error", "message": str(e)}
@@ -1222,13 +1234,28 @@ async def get_config(request: Request):
return {
"device": device,
"cleanup_enabled": os.getenv("AUTO_CLEANUP_ENABLED"),
"file_lifetime": os.getenv("FILE_LIFETIME_SECONDS"),
"cleanup_interval": os.getenv("CLEANUP_INTERVAL_SECONDS"),
"cleanup_config": CLEANUP_CONFIG,
"current_qwen_model": QWEN_MODEL,
"available_qwen_models": AVAILABLE_QWEN_MODELS
}
@app.post("/admin/api/config/cleanup", dependencies=[Depends(verify_admin)])
async def update_cleanup_config(
enabled: bool = Form(...),
lifetime: int = Form(...),
interval: int = Form(...)
):
"""
Update cleanup configuration
"""
global CLEANUP_CONFIG
CLEANUP_CONFIG["enabled"] = enabled
CLEANUP_CONFIG["lifetime"] = lifetime
CLEANUP_CONFIG["interval"] = interval
print(f"Updated Cleanup Config: {CLEANUP_CONFIG}")
return {"status": "success", "message": "Cleanup configuration updated", "config": CLEANUP_CONFIG}
@app.post("/admin/api/config/model", dependencies=[Depends(verify_admin)])
async def set_model(model: str = Form(...)):
"""

67
static/admin.html
View File

@@ -171,24 +171,38 @@
<div class="grid grid-cols-1 md:grid-cols-2 gap-6">
<div class="bg-white rounded-lg shadow p-6">
<h3 class="text-lg font-bold mb-4 text-gray-700">自动清理配置</h3>
<div class="space-y-3">
<div class="flex justify-between border-b pb-2">
<span class="text-gray-600">状态</span>
<span class="font-mono bg-green-100 text-green-800 px-2 rounded text-sm">运行中</span>
<div class="space-y-4">
<div class="flex justify-between items-center border-b pb-2">
<span class="text-gray-600">启用自动清理</span>
<label class="relative inline-flex items-center cursor-pointer">
<input type="checkbox" v-model="cleanupConfig.enabled" class="sr-only peer">
<div class="w-11 h-6 bg-gray-200 peer-focus:outline-none peer-focus:ring-4 peer-focus:ring-blue-300 rounded-full peer peer-checked:after:translate-x-full peer-checked:after:border-white after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-blue-600"></div>
</label>
</div>
<div class="flex justify-between border-b pb-2">
<span class="text-gray-600">文件保留时长</span>
<span class="font-mono">3600 秒 (1小时)</span>
<div class="border-b pb-2">
<div class="flex justify-between items-center mb-1">
<span class="text-gray-600">文件保留时长 (秒)</span>
<span class="text-xs text-gray-500">{{ (cleanupConfig.lifetime / 3600).toFixed(1) }} 小时</span>
</div>
<div class="flex justify-between border-b pb-2">
<span class="text-gray-600">检查间隔</span>
<span class="font-mono">600 秒 (10分钟)</span>
<input type="number" v-model.number="cleanupConfig.lifetime" class="w-full border rounded px-2 py-1 text-sm">
</div>
<div class="mt-4">
<button @click="triggerCleanup" :disabled="cleaning" class="w-full bg-yellow-500 hover:bg-yellow-600 text-white font-bold py-2 px-4 rounded transition">
<i class="fas fa-broom mr-2"></i> {{ cleaning ? '清理中...' : '立即执行清理' }}
<div class="border-b pb-2">
<div class="flex justify-between items-center mb-1">
<span class="text-gray-600">检查间隔 (秒)</span>
<span class="text-xs text-gray-500">{{ (cleanupConfig.interval / 60).toFixed(1) }} 分钟</span>
</div>
<input type="number" v-model.number="cleanupConfig.interval" class="w-full border rounded px-2 py-1 text-sm">
</div>
<div class="flex gap-2 mt-4">
<button @click="saveCleanupConfig" class="flex-1 bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded transition text-sm">
<i class="fas fa-save mr-2"></i> 保存配置
</button>
<button @click="triggerCleanup" :disabled="cleaning" class="flex-1 bg-yellow-500 hover:bg-yellow-600 text-white font-bold py-2 px-4 rounded transition text-sm">
<i class="fas fa-broom mr-2"></i> {{ cleaning ? '清理中...' : '立即清理' }}
</button>
<p class="text-xs text-gray-500 mt-2 text-center">将删除所有超过保留时长的文件</p>
</div>
</div>
</div>
@@ -247,6 +261,11 @@
const deviceInfo = ref('Loading...');
const currentModel = ref('');
const availableModels = ref([]);
const cleanupConfig = ref({
enabled: true,
lifetime: 3600,
interval: 600
});
// 检查登录状态
const checkLogin = () => {
@@ -306,11 +325,28 @@
deviceInfo.value = res.data.device;
currentModel.value = res.data.current_qwen_model;
availableModels.value = res.data.available_qwen_models;
if (res.data.cleanup_config) {
cleanupConfig.value = res.data.cleanup_config;
}
} catch (e) {
console.error(e);
}
};
const saveCleanupConfig = async () => {
try {
const formData = new FormData();
formData.append('enabled', cleanupConfig.value.enabled);
formData.append('lifetime', cleanupConfig.value.lifetime);
formData.append('interval', cleanupConfig.value.interval);
const res = await axios.post('/admin/api/config/cleanup', formData);
alert(res.data.message);
} catch (e) {
alert('保存配置失败: ' + (e.response?.data?.detail || e.message));
}
};
const updateModel = async () => {
try {
const formData = new FormData();
@@ -415,7 +451,8 @@
enterDir, navigateUp, deleteFile, triggerCleanup,
viewResult, previewImage, isImage, previewUrl,
formatDate, getTypeBadgeClass, cleaning, deviceInfo,
currentModel, availableModels, updateModel
currentModel, availableModels, updateModel,
cleanupConfig, saveCleanupConfig
};
}
}).mount('#app');