diff --git a/apps/rmbg/service.py b/apps/rmbg/service.py index 8508586..1a5ec59 100644 --- a/apps/rmbg/service.py +++ b/apps/rmbg/service.py @@ -14,6 +14,8 @@ import asyncio import io import multiprocessing as mp from concurrent.futures import ProcessPoolExecutor +import uuid +from settings import settings # 关闭不必要的警告 warnings.filterwarnings("ignore", category=UserWarning) @@ -28,6 +30,10 @@ class RmbgService: self.model_path = model_path self.model = None self.device = None + self.save_dir = settings.save_dir + self.download_url = settings.download_url + # 确保保存目录存在 + os.makedirs(self.save_dir, exist_ok=True) self._load_model() def _load_model(self): @@ -77,6 +83,27 @@ class RmbgService: buffered = io.BytesIO() image.save(buffered, format="PNG") return base64.b64encode(buffered.getvalue()).decode('utf-8') + + def save_image_to_file(self, image): + """ + 保存图片到jfile/files目录并返回URL + + Args: + image: PIL Image对象 + + Returns: + 图片URL + """ + # 生成唯一文件名 + filename = f"rmbg_{uuid.uuid4().hex[:10]}.png" + file_path = os.path.join(self.save_dir, filename) + + # 保存图片 + image.save(file_path, format="PNG") + + # 构建URL + image_url = f"{self.download_url}/{filename}" + return image_url async def remove_background(self, image_path): """ @@ -107,12 +134,12 @@ class RmbgService: image = Image.open(image_path).convert("RGB") image_no_bg = self.process_image(image) - # 转换为base64 - image_content = self.image_to_base64(image_no_bg) + # 保存图片到文件并获取URL + image_url = self.save_image_to_file(image_no_bg) return { "status": "success", - "image_content": image_content + "image_url": image_url } finally: @@ -138,12 +165,12 @@ class RmbgService: image = Image.open(io.BytesIO(file_content)).convert("RGB") image_no_bg = self.process_image(image) - # 转换为base64 - image_content = self.image_to_base64(image_no_bg) + # 保存图片到文件并获取URL + image_url = self.save_image_to_file(image_no_bg) return { "status": "success", - "image_content": image_content + "image_url": image_url } except Exception as e: @@ -175,7 +202,7 @@ class RmbgService: "total": total, "original_url": url_str, "status": "success", - "image_content": result["image_content"], + "image_url": result["image_url"], "success_count": success_count, "error_count": error_count, "message": "处理成功" diff --git a/apps/rmbg/settings.py b/apps/rmbg/settings.py index 872f597..6a3b2bc 100644 --- a/apps/rmbg/settings.py +++ b/apps/rmbg/settings.py @@ -16,6 +16,11 @@ class Settings(BaseSettings): upload_url: str = "http://images.jingrow.com:8080/api/v1/image" + # 图片保存配置 + save_dir: str = "../jfile/files" + # Japi 静态资源下载URL + download_url: str = "https://api.jingrow.com/files" + # Jingrow Jcloud API 配置 jingrow_api_url: str = "https://cloud.jingrow.com" jingrow_api_key: Optional[str] = None