diff --git a/fastAPI_tarot.py b/fastAPI_tarot.py index 6271c4b..8b8b050 100644 --- a/fastAPI_tarot.py +++ b/fastAPI_tarot.py @@ -19,6 +19,7 @@ import time import json import traceback import re +import asyncio from typing import Optional, List, Dict, Any from contextlib import asynccontextmanager @@ -100,6 +101,49 @@ 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): + """ + 后台任务：定期清理过期的图片文件 + """ + print(f"🧹 自动清理任务已启动 | 目录: {directory} | 生命周期: {lifetime_seconds}s | 检查间隔: {interval_seconds}s") + while True: + try: + await asyncio.sleep(interval_seconds) + current_time = time.time() + count = 0 + # 遍历所有文件（包括子目录） + for root, dirs, files in os.walk(directory): + for file in files: + file_path = os.path.join(root, file) + # 获取文件修改时间 + try: + file_mtime = os.path.getmtime(file_path) + if current_time - file_mtime > lifetime_seconds: + os.remove(file_path) + count += 1 + except OSError: + pass # 文件可能已被删除 + + # 尝试清理空目录 (可选，仅清理二级目录) + for root, dirs, files in os.walk(directory, topdown=False): + for dir in dirs: + dir_path = os.path.join(root, dir) + try: + if not os.listdir(dir_path): # 如果目录为空 + os.rmdir(dir_path) + except OSError: + pass + + if count > 0: + print(f"🧹 已清理 {count} 个过期文件") + + except asyncio.CancelledError: + print("🛑 清理任务已停止") + break + except Exception as e: + print(f"⚠️ 清理任务出错: {e}") + await asyncio.sleep(60) # 出错后等待一分钟再试 + @asynccontextmanager async def lifespan(app: FastAPI): """ @@ -132,9 +176,33 @@ async def lifespan(app: FastAPI): print(f"模型加载完成，设备: {device}") + # --- 启动后台清理任务 --- + 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) + ) + except Exception as e: + print(f"启动清理任务失败: {e}") + # ----------------------- + yield print("正在清理资源...") + + # --- 停止后台任务 --- + if cleanup_task_handle: + cleanup_task_handle.cancel() + try: + await cleanup_task_handle + except asyncio.CancelledError: + pass + # ------------------ + # 这里可以添加释放显存的逻辑，如果需要 # ========================================== @@ -537,14 +605,18 @@ async def segment( image_url: Optional[str] = Form(None, description="URL of the image"), save_segment_images: bool = Form(False, description="Whether to save and return individual segmented objects"), cutout: bool = Form(False, description="If True, returns transparent background PNGs; otherwise returns original crops"), + highlight: bool = Form(False, description="If True, darkens the background to highlight the subject (周边变黑放大)."), confidence: float = Form(0.7, description="Confidence threshold (0.0-1.0). Default is 0.7.") + ): """ **通用图像分割接口** - 支持上传图片或提供图片 URL - 支持自动将中文 Prompt 翻译为英文 + - 支持周边变黑放大效果 (Highlight Mode) - 支持手动设置置信度 (Confidence Threshold) + """ if not file and not image_url: raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)") @@ -601,7 +673,13 @@ async def segment( # 4. 结果可视化与保存 try: - filename = generate_and_save_result(image, inference_state) + if highlight: + filename = f"seg_highlight_{uuid.uuid4().hex}.jpg" + save_path = os.path.join(RESULT_IMAGE_DIR, filename) + # 使用 human_analysis_service 中的可视化函数 (周边变黑) + human_analysis_service.create_highlighted_visualization(image, masks, save_path) + else: + filename = generate_and_save_result(image, inference_state) except Exception as e: raise HTTPException(status_code=500, detail=f"绘图保存错误: {str(e)}") @@ -924,6 +1002,21 @@ async def segment_face( if __name__ == "__main__": import uvicorn + + # ========================================== + # 自动清理图片配置 (Auto Cleanup Config) + # ========================================== + # 设置是否开启自动清理 (True/False) + os.environ["AUTO_CLEANUP_ENABLED"] = "True" + + # 设置图片生命周期 (秒)，超过此时间的图片将被删除 + # 例如: 3600 = 1小时, 86400 = 1天 + os.environ["FILE_LIFETIME_SECONDS"] = "3600" + + # 设置检查间隔 (秒)，每隔多久检查一次 + os.environ["CLEANUP_INTERVAL_SECONDS"] = "600" + # ========================================== + # 启动服务器 uvicorn.run( "fastAPI_tarot:app", diff --git a/static/results/1771167659_ef9e57ff/face_0.jpg b/static/results/1771167659_ef9e57ff/face_0.jpg new file mode 100644 index 0000000..60c24f1 Binary files /dev/null and b/static/results/1771167659_ef9e57ff/face_0.jpg differ diff --git a/static/results/1771167659_ef9e57ff/seg_d6acc79b658b4534bedd7fcf9e23158e.jpg b/static/results/1771167659_ef9e57ff/seg_d6acc79b658b4534bedd7fcf9e23158e.jpg new file mode 100644 index 0000000..4135929 Binary files /dev/null and b/static/results/1771167659_ef9e57ff/seg_d6acc79b658b4534bedd7fcf9e23158e.jpg differ diff --git a/static/results/seg_9745a70b2b5c4cc6b0d1b69857704aed.jpg b/static/results/seg_9745a70b2b5c4cc6b0d1b69857704aed.jpg new file mode 100644 index 0000000..c4330df Binary files /dev/null and b/static/results/seg_9745a70b2b5c4cc6b0d1b69857704aed.jpg differ diff --git a/static/results/seg_highlight_7a8ddc7c47b547d1b8c574fbcf5b28cc.jpg b/static/results/seg_highlight_7a8ddc7c47b547d1b8c574fbcf5b28cc.jpg new file mode 100644 index 0000000..276f45f Binary files /dev/null and b/static/results/seg_highlight_7a8ddc7c47b547d1b8c574fbcf5b28cc.jpg differ