zxd

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",