diff --git a/apps/jfile/files/25_5uvvahf6kj_d9c433c5_1.png b/apps/jfile/files/25_5uvvahf6kj_d9c433c5_1.png deleted file mode 100644 index 9d6fe0e..0000000 Binary files a/apps/jfile/files/25_5uvvahf6kj_d9c433c5_1.png and /dev/null differ diff --git a/apps/jfile/files/split_2acec28387_1.png b/apps/jfile/files/split_2acec28387_1.png deleted file mode 100644 index 6256135..0000000 Binary files a/apps/jfile/files/split_2acec28387_1.png and /dev/null differ diff --git a/apps/jfile/files/split_2acec28387_2.png b/apps/jfile/files/split_2acec28387_2.png deleted file mode 100644 index b1c2363..0000000 Binary files a/apps/jfile/files/split_2acec28387_2.png and /dev/null differ diff --git a/apps/jfile/files/split_2acec28387_3.png b/apps/jfile/files/split_2acec28387_3.png deleted file mode 100644 index 40a3a1f..0000000 Binary files a/apps/jfile/files/split_2acec28387_3.png and /dev/null differ diff --git a/apps/jfile/files/split_2acec28387_4.png b/apps/jfile/files/split_2acec28387_4.png deleted file mode 100644 index 1b82fcb..0000000 Binary files a/apps/jfile/files/split_2acec28387_4.png and /dev/null differ diff --git a/apps/jfile/files/split_c4191175e4_1.png b/apps/jfile/files/split_c4191175e4_1.png deleted file mode 100644 index 6bfabf3..0000000 Binary files a/apps/jfile/files/split_c4191175e4_1.png and /dev/null differ diff --git a/apps/jfile/files/split_c4191175e4_2.png b/apps/jfile/files/split_c4191175e4_2.png deleted file mode 100644 index 525bcf5..0000000 Binary files a/apps/jfile/files/split_c4191175e4_2.png and /dev/null differ diff --git a/apps/jfile/files/split_c4191175e4_3.png b/apps/jfile/files/split_c4191175e4_3.png deleted file mode 100644 index 7421285..0000000 Binary files a/apps/jfile/files/split_c4191175e4_3.png and /dev/null differ diff --git a/apps/jfile/files/split_c4191175e4_4.png b/apps/jfile/files/split_c4191175e4_4.png deleted file mode 100644 index efa8c20..0000000 Binary files a/apps/jfile/files/split_c4191175e4_4.png and /dev/null differ diff --git a/apps/rmbg/service.py b/apps/rmbg/service.py index 1a5ec59..9127e4e 100644 --- a/apps/rmbg/service.py +++ b/apps/rmbg/service.py @@ -1,5 +1,4 @@ import os -import requests import tempfile from urllib.parse import urlparse from PIL import Image @@ -12,9 +11,8 @@ import gc import base64 import asyncio import io -import multiprocessing as mp -from concurrent.futures import ProcessPoolExecutor import uuid +import httpx from settings import settings # 关闭不必要的警告 @@ -34,6 +32,8 @@ class RmbgService: self.download_url = settings.download_url # 确保保存目录存在 os.makedirs(self.save_dir, exist_ok=True) + # 创建异步HTTP客户端（复用连接，提高性能） + self.http_client = httpx.AsyncClient(timeout=30.0, limits=httpx.Limits(max_keepalive_connections=20)) self._load_model() def _load_model(self): @@ -45,11 +45,10 @@ class RmbgService: self.model = self.model.to(self.device) self.model.eval() # 设置为评估模式 - def process_image(self, image): - """处理图像，移除背景""" + def _process_image_sync(self, image): + """同步处理图像，移除背景（内部方法，在线程池中执行）""" image_size = image.size # 转换图像 - t0 = time.time() transform_image = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), @@ -58,12 +57,10 @@ class RmbgService: input_images = transform_image(image).unsqueeze(0).to(self.device) # 推理 - t0 = time.time() with torch.no_grad(): preds = self.model(input_images)[-1].sigmoid().cpu() # 处理预测结果 - t0 = time.time() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize(image_size) @@ -77,6 +74,12 @@ class RmbgService: gc.collect() return image + + async def process_image(self, image): + """异步处理图像，移除背景（在线程池中执行同步操作）""" + # 将同步的GPU操作放到线程池中执行，避免阻塞事件循环 + loop = asyncio.get_event_loop() + return await loop.run_in_executor(None, self._process_image_sync, image) def image_to_base64(self, image): """将PIL Image对象转换为base64字符串""" @@ -120,8 +123,8 @@ class RmbgService: # 检查是否是URL if self.is_valid_url(image_path): try: - # 下载图片到临时文件 - temp_file = self.download_image(image_path) + # 异步下载图片到临时文件 + temp_file = await self.download_image(image_path) image_path = temp_file except Exception as e: raise Exception(f"下载图片失败: {e}") @@ -130,12 +133,22 @@ class RmbgService: if not os.path.exists(image_path): raise FileNotFoundError(f"输入图像不存在: {image_path}") - # 加载并处理图像 - image = Image.open(image_path).convert("RGB") - image_no_bg = self.process_image(image) + # 加载图像（IO操作，在线程池中执行） + loop = asyncio.get_event_loop() + image = await loop.run_in_executor( + None, + lambda: Image.open(image_path).convert("RGB") + ) - # 保存图片到文件并获取URL - image_url = self.save_image_to_file(image_no_bg) + # 异步处理图像 + image_no_bg = await self.process_image(image) + + # 保存图片到文件并获取URL（IO操作，在线程池中执行） + image_url = await loop.run_in_executor( + None, + self.save_image_to_file, + image_no_bg + ) return { "status": "success", @@ -161,12 +174,22 @@ class RmbgService: 处理后的图像内容 """ try: - # 从文件内容创建PIL Image对象 - image = Image.open(io.BytesIO(file_content)).convert("RGB") - image_no_bg = self.process_image(image) + # 从文件内容创建PIL Image对象（IO操作，在线程池中执行） + loop = asyncio.get_event_loop() + image = await loop.run_in_executor( + None, + lambda: Image.open(io.BytesIO(file_content)).convert("RGB") + ) - # 保存图片到文件并获取URL - image_url = self.save_image_to_file(image_no_bg) + # 异步处理图像 + image_no_bg = await self.process_image(image) + + # 保存图片到文件并获取URL（IO操作，在线程池中执行） + image_url = await loop.run_in_executor( + None, + self.save_image_to_file, + image_no_bg + ) return { "status": "success", @@ -178,51 +201,61 @@ class RmbgService: async def process_batch(self, urls): """ - 批量处理多个URL图像，流式返回结果 + 批量处理多个URL图像，并发处理并流式返回结果 Args: urls: 图片URL列表 Yields: - 每个图片的处理结果 + 每个图片的处理结果（按完成顺序返回） """ total = len(urls) success_count = 0 error_count = 0 - for i, url in enumerate(urls, 1): + # 创建并发任务 + async def process_single_url(index, url): + """处理单个URL的包装函数""" try: url_str = str(url) result = await self.remove_background(url_str) - success_count += 1 - - # 确保返回正确的数据格式 - yield { - "index": i, + return { + "index": index, "total": total, "original_url": url_str, "status": "success", "image_url": result["image_url"], - "success_count": success_count, - "error_count": error_count, "message": "处理成功" } - except Exception as e: - error_count += 1 - yield { - "index": i, + return { + "index": index, "total": total, "original_url": str(url), "status": "error", "error": str(e), - "success_count": success_count, - "error_count": error_count, "message": f"处理失败: {str(e)}" } + + # 创建所有任务 + tasks = [ + process_single_url(i, url) + for i, url in enumerate(urls, 1) + ] + + # 并发执行所有任务，使用as_completed按完成顺序返回 + for coro in asyncio.as_completed(tasks): + result = await coro + if result["status"] == "success": + success_count += 1 + else: + error_count += 1 - # 让出控制权，避免阻塞 - await asyncio.sleep(0) + # 更新统计信息 + result["success_count"] = success_count + result["error_count"] = error_count + + yield result def is_valid_url(self, url): """验证URL是否有效""" @@ -232,20 +265,34 @@ class RmbgService: except: return False - def download_image(self, url): - """从URL下载图片到临时文件""" - response = requests.get(url, stream=True) - response.raise_for_status() - - # 创建临时文件 - temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png') - with open(temp_file.name, 'wb') as f: - for chunk in response.iter_content(chunk_size=8192): - f.write(chunk) - return temp_file.name + async def download_image(self, url): + """异步从URL下载图片到临时文件""" + try: + response = await self.http_client.get(url) + response.raise_for_status() + + # 创建临时文件并写入内容（IO操作，在线程池中执行） + def write_temp_file(content): + temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png') + temp_file.write(content) + temp_file.close() + return temp_file.name + + loop = asyncio.get_event_loop() + temp_file_path = await loop.run_in_executor( + None, + write_temp_file, + response.content + ) + + return temp_file_path + except Exception as e: + raise Exception(f"下载图片失败: {e}") - def cleanup(self): + async def cleanup(self): """清理资源""" + # 关闭HTTP客户端 + await self.http_client.aclose() if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect()