# main.py from fastapi import FastAPI, File, UploadFile, Depends, HTTPException, Security from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from fastapi.security import APIKeyHeader from kafka import KafkaProducer from redis import Redis import os import json import uuid from datetime import datetime, timedelta, timezone import string from decord import VideoReader from PIL import Image from fastapi.responses import FileResponse import logging from config import * app = FastAPI() v1_app = FastAPI() app.mount("/v1", v1_app) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 配置 KAFKA_BROKER = KAFKA_BROKER REDIS_HOST = REDIS_HOST REDIS_PORT = REDIS_PORT REDIS_PASSWORD = REDIS_PASSWORD REDIS_DB = MAIN_REDIS_DB REDIS_API_DB = REDIS_API_DB REDIS_API_USAGE_DB = REDIS_API_USAGE_DB UPLOAD_DIR = UPLOAD_DIR RESULT_DIR = RESULT_DIR MAX_FILE_AGE = timedelta(hours=1) # 确保目录存在 os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(RESULT_DIR, exist_ok=True) # 定义支持的任务类型 KAFKA_TOPICS = { 'pose': 'pose', 'mediapipe': 'mediapipe', 'qwenvl': 'qwenvl', 'yolo': 'yolo', 'fall': 'fall', 'face': 'face', 'cpm': 'cpm', 'compare': 'compare', 'qwenvl_analyze': 'qwenvl_analyze', 'cpm_analyze': 'cpm_analyze' } TASK_TYPES = list(KAFKA_TOPICS.keys()) # 初始化 Kafka Producer producer = KafkaProducer( bootstrap_servers=[KAFKA_BROKER], value_serializer=lambda v: json.dumps(v).encode('utf-8') ) # 初始化 Redis redis_client = Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB ) redis_api_client = Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_API_DB ) redis_api_usage_client = Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_API_USAGE_DB ) redis_pose_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['pose']['redis_db']) redis_cpm_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['cpm']['redis_db']) redis_yolo_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['yolo']['redis_db']) redis_face_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['face']['redis_db']) redis_fall_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['fall']['redis_db']) redis_mediapipe_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['mediapipe']['redis_db']) redis_qwenvl_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['qwenvl']['redis_db']) redis_compare_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['compare']['redis_db']) redis_qwenvl_analyze_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['qwenvl_analyze']['redis_db']) redis_cpm_analyze_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=WORKER_CONFIGS['cpm_analyze']['redis_db']) @v1_app.get('/favicon.ico', include_in_schema=False) async def favicon(): file_name = "favicon.ico" file_path = os.path.join(app.root_path, "static", file_name) if os.path.isfile(file_path): return FileResponse(file_path) else: return {"message": "Favicon not found"}, 404 # 定义API密钥头部 api_key_header = APIKeyHeader(name="Authorization", auto_error=False) # 定义 base62 字符集 BASE62 = string.digits + string.ascii_letters # 验证API密钥的函数 async def get_api_key(api_key: str = Security(api_key_header)): if api_key and api_key.startswith("Bearer "): key = api_key.split(" ")[1] if key.startswith("obs-"): return key raise HTTPException( status_code=401, detail="无效的API密钥", headers={"WWW-Authenticate": "Bearer"}, ) async def verify_api_key(api_key: str = Depends(get_api_key)): logging.info(f"验证API密钥: {api_key}") redis_key = f"api_key:{api_key}" api_key_info = redis_api_client.hgetall(redis_key) if not api_key_info: logging.warning(f"API密钥不存在: {api_key}") raise HTTPException(status_code=401, detail="无效的API密钥") api_key_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in api_key_info.items()} if api_key_info.get('is_active') != '1': logging.warning(f"API密钥已停用: {api_key}") raise HTTPException(status_code=401, detail="API密钥已停用") expires_at = datetime.fromisoformat(api_key_info.get('expires_at')) if datetime.now(timezone.utc) > expires_at: logging.warning(f"API密钥已过期: {api_key}") raise HTTPException(status_code=401, detail="API密钥已过期") usage_info = redis_api_usage_client.hgetall(redis_key) usage_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in usage_info.items()} logging.info(f"API密钥验证成功: {api_key}") return { "api_key": api_key, **api_key_info, **usage_info } async def update_token_usage(api_key: str, new_tokens_used: int, model_name: str): redis_key = f"api_key:{api_key}" current_time = datetime.now(timezone.utc).isoformat() pipe = redis_api_usage_client.pipeline() # 更新总的token使用量 pipe.hincrby(redis_key, "tokens_used", new_tokens_used) pipe.hset(redis_key, "last_used_at", current_time) # 更新特定模型的token使用量 model_tokens_field = f"{model_name}_tokens_used" model_last_used_field = f"{model_name}_last_used_at" pipe.hincrby(redis_key, model_tokens_field, new_tokens_used) pipe.hset(redis_key, model_last_used_field, current_time) pipe.execute() def calculate_tokens(file_path: str, file_type: str) -> int: base_tokens = 0 try: file_size = os.path.getsize(file_path) # 获取文件大小（字节） # 基础token：每MB文件大小消耗10个token base_tokens = int((file_size / (1024 * 1024)) * 0.1) if file_type == "image": img = Image.open(file_path) width, height = img.size pixel_count = width * height # 图片token：每100个像素额外消耗5个token image_tokens = int((pixel_count / 100000000) * 0.1) base_tokens += image_tokens elif file_type == "video": vr = VideoReader(file_path) fps = vr.get_avg_fps() frame_count = len(vr) width, height = vr[0].shape[1], vr[0].shape[0] pixel_count = width * height * frame_count duration = frame_count / fps # 视频时长（秒） # 视频token：每100万像素每秒额外消耗1个token video_tokens = int((pixel_count / 100000000) * (duration / 60)) base_tokens += video_tokens return max(1, base_tokens) # 确保至少返回1个token except Exception as e: print(f"计算token时出错: {str(e)}") return 1 # 出错时返回默认值1 async def upload_file(file: UploadFile, task_type: str, api_key_info: dict): if task_type not in KAFKA_TOPICS: raise HTTPException(status_code=400, detail="不支持的任务类型") content = await file.read() file_extension = os.path.splitext(file.filename)[1].lower() new_filename = f"{uuid.uuid4()}{file_extension}" file_path = os.path.join(UPLOAD_DIR, new_filename) with open(file_path, "wb") as f: f.write(content) # 计算 token file_type = "image" if file_extension in ['.jpg', '.jpeg', '.png'] else "video" tokens_required = calculate_tokens(file_path, file_type) # 检查并更新 token 使用量 api_key = api_key_info['api_key'] usage_key = f"api_key:{api_key}" total_tokens = int(redis_api_usage_client.hget(usage_key, "total_tokens") or 0) tokens_used = int(redis_api_usage_client.hget(usage_key, "tokens_used") or 0) if tokens_used + tokens_required > total_tokens: raise HTTPException(status_code=403, detail="Token 余额不足") # 更新 token 使用量 await update_token_usage(api_key, tokens_required, task_type) # 创建任务记录 task_id = str(uuid.uuid4()) task_data = { "task_id": task_id, "filename": new_filename, "file_type": file_type, "task_type": task_type, "status": "queued", "created_at": datetime.now(timezone.utc).isoformat(), } # 存储任务信息到Redis redis_client.set(f"task:{task_id}", json.dumps(task_data)) logging.info(f"任务信息已存储到Redis: {task_id}") # 发送任务到对应的Kafka主题 kafka_topic = KAFKA_TOPICS[task_type] producer.send(kafka_topic, task_data) logging.info(f"任务已发送到Kafka主题: {kafka_topic}") # 获取更新后的 token 使用情况 updated_api_key_info = await verify_api_key(api_key) new_tokens_used = int(updated_api_key_info.get("tokens_used", 0)) model_tokens_used = int(updated_api_key_info.get(f"{task_type}_tokens_used", 0)) response_data = { "message": "文件已上传并排队等待处理", "task_id": task_id, "tokens_used": tokens_required, "total_tokens_used": new_tokens_used, f"{task_type}_tokens_used": model_tokens_used, "tokens_remaining": total_tokens - new_tokens_used } logging.info(f"上传文件完成: {task_id}") return JSONResponse(content=response_data) # 为每个任务类型创建单独的端点 @v1_app.post("/pose") async def upload_pose(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): logging.info(f"收到 /pose端点的请求") return await upload_file(file, task_type="pose", api_key_info=api_key_info) @v1_app.post("/cpm") async def upload_cpm(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="cpm", api_key_info=api_key_info) @v1_app.post("/qwenvl") async def upload_qwenvl(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="qwenvl", api_key_info=api_key_info) @v1_app.post("/qwenvl_analyze") async def upload_qwenvl(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="qwenvl_analyze", api_key_info=api_key_info) @v1_app.post("/cpm_analyze") async def upload_qwenvl(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="cpm_analyze", api_key_info=api_key_info) @v1_app.post("/yolo") async def upload_yolo(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="yolo", api_key_info=api_key_info) @v1_app.post("/fall") async def upload_fall(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="fall", api_key_info=api_key_info) @v1_app.post("/face") async def upload_face(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="face", api_key_info=api_key_info) @v1_app.post("/mediapipe") async def upload_mediapipe(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="mediapipe", api_key_info=api_key_info) @v1_app.post("/compare") async def upload_compare(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)): return await upload_file(file, task_type="compare", api_key_info=api_key_info) @v1_app.get("/result/{task_id}") async def get_result(task_id: str, api_key: str = Depends(get_api_key)): api_key_info = await verify_api_key(api_key) if not api_key_info: raise HTTPException(status_code=403, detail="无效的API密钥") # 从 REDIS_DB (15) 获取任务状态 task_info = redis_client.hgetall(f"task:{task_id}") if not task_info: raise HTTPException(status_code=404, detail="Task not found") task_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in task_info.items()} if task_info['status'] != 'completed': return {"status": task_info['status'], "message": "Task is not completed yet"} result_type = task_info['result_type'] result_key = task_info['result_key'] # 根据任务类型选择相应的 Redis 客户端 redis_client_map = { 'pose': redis_pose_client, 'cpm': redis_cpm_client, 'yolo': redis_yolo_client, 'face': redis_face_client, 'fall': redis_fall_client, 'mediapipe': redis_mediapipe_client, 'qwenvl': redis_qwenvl_client, 'compare': redis_compare_client, 'qwenvl_analyze': redis_qwenvl_analyze_client, 'cpm_analyze': redis_cpm_analyze_client } result_redis = redis_client_map.get(result_type) if not result_redis: raise HTTPException(status_code=400, detail="Unsupported result type") result = result_redis.hgetall(result_key) if not result: raise HTTPException(status_code=404, detail=f"{result_type.upper()} result not found") result = {k.decode('utf-8'): v.decode('utf-8') for k, v in result.items()} # 将 result 字段解析为 JSON（如果存在） if 'result' in result: result['result'] = json.loads(result['result']) return { "status": "completed", "result_type": result_type, "result": result } @v1_app.get("/annotated/{task_id}") async def get_annotated_image(task_id: str, api_key: str = Depends(get_api_key)): api_key_info = await verify_api_key(api_key) if not api_key_info: raise HTTPException(status_code=403, detail="无效的API密钥") # 从 REDIS_DB (15) 获取任务信息 task_info = redis_client.hgetall(f"task:{task_id}") if not task_info: raise HTTPException(status_code=404, detail="Task not found") task_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in task_info.items()} if task_info['status'] != 'completed': raise HTTPException(status_code=400, detail="Task is not completed yet") result_type = task_info.get('result_type') result_key = task_info.get('result_key') if not result_key: raise HTTPException(status_code=404, detail="Result key not found") if result_type in ['cpm', 'qwenvl','cpm_analyze', 'qwenvl_analyze']: raise HTTPException(status_code=400, detail="Annotated image not available for this task type") # 根据任务类型选择相应的 Redis 客户端 redis_client_map = { 'pose': redis_pose_client, 'yolo': redis_yolo_client, 'face': redis_face_client, 'fall': redis_fall_client, 'mediapipe': redis_mediapipe_client, 'compare': redis_compare_client } result_redis = redis_client_map.get(result_type) if not result_redis: raise HTTPException(status_code=400, detail="Unsupported result type") result = result_redis.hgetall(result_key) if not result: raise HTTPException(status_code=404, detail=f"{result_type.upper()} result not found") result = {k.decode('utf-8'): v.decode('utf-8') for k, v in result.items()} result_file = result.get('result_file') if not result_file: raise HTTPException(status_code=404, detail="Result file not found") file_path = os.path.join(RESULT_DIR, result_file) if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Result image file not found") return FileResponse(file_path, media_type="image/png") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8005)