From 7391f6a0c604084dc0f564375ea1faaab89339b2 Mon Sep 17 00:00:00 2001 From: fanpt <320622572@qq.com> Date: Fri, 21 Jun 2024 10:16:56 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E4=BB=A3=E7=A0=81=E6=B3=A8?= =?UTF-8?q?=E9=87=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/main.py b/main.py index 2b2d069..c0bc2f5 100644 --- a/main.py +++ b/main.py @@ -10,18 +10,21 @@ from hashlib import md5 from typing import Dict app = FastAPI() -tasks: Dict[str, dict] = {} # To store the status and result of each task -md5_to_uid: Dict[str, str] = {} # To map md5 to uid +tasks: Dict[str, dict] = {} # 存储每个任务的状态和结果 +md5_to_uid: Dict[str, str] = {} # 将md5映射到uid -ALLOWED_IMAGE_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'} +ALLOWED_IMAGE_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'} # 允许的图片扩展名 +# 保存上传的文件 def save_upload_file(upload_file: UploadFile, filename: str): with open(filename, "wb") as buffer: buffer.write(upload_file.file.read()) +# 检查文件是否是允许的类型 def is_allowed_file(filename: str): return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_IMAGE_EXTENSIONS +# 生成视频处理命令 def generate_video_command(result_dir: str, img_path: str, audio_path: str, video_path: str): return [ "python", "script.py", @@ -31,10 +34,12 @@ def generate_video_command(result_dir: str, img_path: str, audio_path: str, vide "--ref_eyeblink", video_path, ] +# 获取最新的子目录 def get_latest_sub_dir(result_dir: str): sub_dirs = [os.path.join(result_dir, d) for d in os.listdir(result_dir) if os.path.isdir(os.path.join(result_dir, d))] return max(sub_dirs, key=os.path.getmtime, default=None) +# 获取视频时长 def get_video_duration(video_path: str): result = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", video_path], @@ -42,15 +47,18 @@ def get_video_duration(video_path: str): ) return float(result.stdout.strip()) +# 运行FFmpeg命令 def run_ffmpeg_command(command: list): subprocess.run(command, check=True) +# 将文件保存到最终目的地 def save_to_final_destination(source_path: str, destination_dir: str): os.makedirs(destination_dir, exist_ok=True) destination_path = os.path.join(destination_dir, os.path.basename(source_path)) os.rename(source_path, destination_path) return destination_path +# 计算文件的MD5 def get_file_md5(file_path: str): hash_md5 = md5() with open(file_path, "rb") as f: @@ -58,10 +66,12 @@ def get_file_md5(file_path: str): hash_md5.update(chunk) return hash_md5.hexdigest() +# 记录视频映射关系 def record_video_mapping(image_md5: str, video_path: str, record_file: str): with open(record_file, "a") as f: f.write(f"{image_md5} {video_path}\n") +# 检查是否存在已处理的视频 def check_existing_video(image_md5: str, record_file: str): if not os.path.exists(record_file): return None @@ -72,6 +82,7 @@ def check_existing_video(image_md5: str, record_file: str): return video_path return None +# 处理视频生成 def process_video(uid: str, image_md5: str, img_path: str, audio_type: str): record_file = f"{audio_type}_video_record.txt" audio_path = f"./examples/driven_audio/{audio_type}_audio.wav" @@ -106,13 +117,15 @@ def process_video(uid: str, image_md5: str, img_path: str, audio_type: str): else: tasks[uid]['status'] = 'failed' +# 将视频编码为base64 def encode_video_to_base64(video_path: str): with open(video_path, "rb") as video_file: return base64.b64encode(video_file.read()).decode('utf-8') +# 处理视频请求 async def handle_video_request(background_tasks: BackgroundTasks, image: UploadFile, audio_type: str): if not is_allowed_file(image.filename): - return JSONResponse(status_code=400, content={"code": 400, "message": "Invalid file type. Only images (png, jpg, jpeg, gif) are allowed.", "uid": "", "video": ""}) + return JSONResponse(status_code=400, content={"code": 500, "message": "文件类型无效。只允许图像文件（png, jpg, jpeg, gif）。", "uid": "", "video": ""}) img_path = os.path.join(audio_type, image.filename) save_upload_file(image, img_path) @@ -122,38 +135,41 @@ async def handle_video_request(background_tasks: BackgroundTasks, image: UploadF existing_video = check_existing_video(image_md5, record_file) if existing_video: video_base64 = encode_video_to_base64(existing_video) - return JSONResponse(content={"code": 200, "message": "Video retrieved successfully.", "uid": "", "video": video_base64}) + return JSONResponse(content={"code": 500, "message": "视频检索成功。", "uid": "", "video": video_base64}) if image_md5 in md5_to_uid: uid = md5_to_uid[image_md5] - return JSONResponse(content={"code": 200, "message": "Video is being generated, please check back later.", "uid": uid, "video": ""}) + return JSONResponse(content={"code": 200, "message": "视频正在生成，请稍后再试。", "uid": uid, "video": ""}) uid = str(uuid.uuid4()) tasks[uid] = {'status': 'processing'} md5_to_uid[image_md5] = uid background_tasks.add_task(process_video, uid, image_md5, img_path, audio_type) - return JSONResponse(content={"code": 200, "message": "Video is being generated, please check back later.", "uid": uid, "video": ""}) + return JSONResponse(content={"code": 200, "message": "视频正在生成，请稍后再试。", "uid": uid, "video": ""}) +# 动态视频生成端点 @app.post("/dynamic-video") async def generate_dynamic_video(background_tasks: BackgroundTasks, image: UploadFile = File(...)): return await handle_video_request(background_tasks, image, "dynamic") +# 静态视频生成端点 @app.post("/silent-video") async def generate_silent_video(background_tasks: BackgroundTasks, image: UploadFile = File(...)): return await handle_video_request(background_tasks, image, "silent") +# 获取任务状态端点 @app.get("/status/{uid}") async def get_status(uid: str): task = tasks.get(uid) if not task: - return JSONResponse(status_code=404, content={"code": 500, "status": "not found", "message": "Task not found", "video": ""}) + return JSONResponse(status_code=404, content={"code": 500, "status": "not found", "message": "任务未找到", "video": ""}) if task['status'] == 'completed': video_base64 = encode_video_to_base64(task['video_path']) - return JSONResponse(content={"code": 200, "status": task['status'], "message": "Video generation completed.", "video": video_base64}) + return JSONResponse(content={"code": 200, "status": task['status'], "message": "视频生成完成。", "video": video_base64}) elif task['status'] == 'failed': - return JSONResponse(status_code=500, content={"code": 500, "status": task['status'], "message": "Video generation failed", "video": ""}) + return JSONResponse(status_code=500, content={"code": 500, "status": task['status'], "message": "视频生成失败", "video": ""}) else: - return JSONResponse(content={"code": 200, "status": task['status'], "message": "Video is being generated.", "video": ""}) + return JSONResponse(content={"code": 200, "status": task['status'], "message": "视频正在生成。", "video": ""}) if __name__ == "__main__": import uvicorn