feat(rmbg): 添加图片URL返回功能,移除base64字段

- 添加保存图片到jfile/files目录的功能
- 返回结果从image_content改为image_url
- 添加save_dir和download_url配置项
- 大幅减少返回数据大小,提升传输性能
- file和batch接口均支持返回图片URL
This commit is contained in:
jingrow 2025-11-21 01:34:31 +08:00
parent b3818a9ab1
commit ba4a57849d
2 changed files with 39 additions and 7 deletions

View File

@ -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": "处理成功"

View File

@ -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