feat: add outfit generation core
This commit is contained in:
@@ -0,0 +1,131 @@
|
|||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from PIL import Image, ImageOps
|
||||||
|
|
||||||
|
from core.models import OutfitResult, OutfitTask
|
||||||
|
from services.ai_image_service import ImageApiClient
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
QUALITY_SMALL = 75
|
||||||
|
QUALITY_BALANCED = 85
|
||||||
|
QUALITY_HIGH = 92
|
||||||
|
QUALITY_PRESETS = {
|
||||||
|
"small": QUALITY_SMALL,
|
||||||
|
"balanced": QUALITY_BALANCED,
|
||||||
|
"high": QUALITY_HIGH,
|
||||||
|
"小文件": QUALITY_SMALL,
|
||||||
|
"均衡": QUALITY_BALANCED,
|
||||||
|
"高清": QUALITY_HIGH,
|
||||||
|
}
|
||||||
|
|
||||||
|
MAX_JPG_BYTES = 2 * 1024 * 1024
|
||||||
|
_INVALID_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
|
||||||
|
|
||||||
|
|
||||||
|
def render_prompt(template, task):
|
||||||
|
"""Render an outfit prompt for one task."""
|
||||||
|
return str(template).replace("{title}", task.title).replace("{product_id}", task.product_id)
|
||||||
|
|
||||||
|
|
||||||
|
def safe_product_filename(product_id):
|
||||||
|
"""Return a filesystem-safe base filename without changing Excel data."""
|
||||||
|
name = _INVALID_FILENAME_CHARS.sub("_", str(product_id))
|
||||||
|
name = name.strip().strip(".")
|
||||||
|
return name or "outfit"
|
||||||
|
|
||||||
|
|
||||||
|
def make_outfit_output_path(output_dir, product_id):
|
||||||
|
"""Return output_dir/<product_id>.jpg without overwriting existing files."""
|
||||||
|
target_dir = Path(output_dir)
|
||||||
|
stem = safe_product_filename(product_id)
|
||||||
|
candidate = target_dir / (stem + ".jpg")
|
||||||
|
counter = 1
|
||||||
|
while candidate.exists():
|
||||||
|
candidate = target_dir / ("{}_{}.jpg".format(stem, counter))
|
||||||
|
counter += 1
|
||||||
|
return candidate
|
||||||
|
|
||||||
|
|
||||||
|
def save_jpg_under_limit(image_bytes, output_path, quality=QUALITY_BALANCED, max_bytes=MAX_JPG_BYTES):
|
||||||
|
"""Save image bytes as 1:1 JPG, reducing quality/size until under limit."""
|
||||||
|
output_path = Path(output_path)
|
||||||
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with Image.open(BytesIO(image_bytes)) as opened:
|
||||||
|
img = ImageOps.exif_transpose(opened).convert("RGB")
|
||||||
|
img = _crop_square(img)
|
||||||
|
|
||||||
|
quality = _coerce_quality(quality)
|
||||||
|
min_quality = 55
|
||||||
|
while True:
|
||||||
|
_save_jpeg(img, output_path, quality)
|
||||||
|
if output_path.stat().st_size <= max_bytes:
|
||||||
|
return output_path
|
||||||
|
if quality > min_quality:
|
||||||
|
quality = max(min_quality, quality - 7)
|
||||||
|
continue
|
||||||
|
if img.width <= 512 or img.height <= 512:
|
||||||
|
return output_path
|
||||||
|
new_size = max(512, int(img.width * 0.85))
|
||||||
|
img = img.resize((new_size, new_size), Image.LANCZOS)
|
||||||
|
quality = _coerce_quality(quality)
|
||||||
|
|
||||||
|
|
||||||
|
def generate_outfit_image(
|
||||||
|
task,
|
||||||
|
prompt_template,
|
||||||
|
output_dir,
|
||||||
|
model_config,
|
||||||
|
quality=QUALITY_BALANCED,
|
||||||
|
resolution="1K",
|
||||||
|
api_client=None,
|
||||||
|
):
|
||||||
|
"""Generate one outfit image and return OutfitResult. Never raises."""
|
||||||
|
if not isinstance(task, OutfitTask):
|
||||||
|
raise TypeError("task must be OutfitTask")
|
||||||
|
|
||||||
|
try:
|
||||||
|
prompt = render_prompt(prompt_template, task)
|
||||||
|
client = api_client or ImageApiClient(model_config)
|
||||||
|
image_bytes = client.generate(prompt, task.garment_path, resolution=resolution)
|
||||||
|
output_path = make_outfit_output_path(output_dir, task.product_id)
|
||||||
|
save_jpg_under_limit(image_bytes, output_path, quality=quality)
|
||||||
|
logger.info("Generated outfit row %s -> %s", task.row_index, output_path)
|
||||||
|
return OutfitResult(
|
||||||
|
task=task,
|
||||||
|
success=True,
|
||||||
|
output_path=str(output_path),
|
||||||
|
attempts=1,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Outfit generation failed for row %s", task.row_index)
|
||||||
|
return OutfitResult(
|
||||||
|
task=task,
|
||||||
|
success=False,
|
||||||
|
error=str(exc),
|
||||||
|
attempts=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _coerce_quality(quality):
|
||||||
|
if isinstance(quality, str):
|
||||||
|
return QUALITY_PRESETS.get(quality, QUALITY_BALANCED)
|
||||||
|
try:
|
||||||
|
return max(1, min(95, int(quality)))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return QUALITY_BALANCED
|
||||||
|
|
||||||
|
|
||||||
|
def _crop_square(img):
|
||||||
|
width, height = img.size
|
||||||
|
side = min(width, height)
|
||||||
|
left = (width - side) // 2
|
||||||
|
top = (height - side) // 2
|
||||||
|
return img.crop((left, top, left + side, top + side))
|
||||||
|
|
||||||
|
|
||||||
|
def _save_jpeg(img, output_path, quality):
|
||||||
|
img.save(str(output_path), format="JPEG", quality=quality, optimize=True)
|
||||||
@@ -1049,7 +1049,7 @@
|
|||||||
- [x] `core/models.py` 新增 `OutfitTask` / `OutfitResult`(纯 dataclass,Python 3.7 兼容,不依赖 PySide6)
|
- [x] `core/models.py` 新增 `OutfitTask` / `OutfitResult`(纯 dataclass,Python 3.7 兼容,不依赖 PySide6)
|
||||||
- [x] `services/excel_service.py`:读行 → `List[OutfitTask]`、写回 D/E/F、占用检测、跳过「完成」/按设置重试「失败」/空字段安全跳过、每行即存
|
- [x] `services/excel_service.py`:读行 → `List[OutfitTask]`、写回 D/E/F、占用检测、跳过「完成」/按设置重试「失败」/空字段安全跳过、每行即存
|
||||||
- [x] `services/ai_image_service.py`:移植旧项目 `ImageApiClient`(多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验);注意 PEP585 类型注解改 Python 3.7 写法
|
- [x] `services/ai_image_service.py`:移植旧项目 `ImageApiClient`(多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验);注意 PEP585 类型注解改 Python 3.7 写法
|
||||||
- [ ] `core/ai_outfit.py`:单行生成纯逻辑编排(提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`)
|
- [x] `core/ai_outfit.py`:单行生成纯逻辑编排(提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`)
|
||||||
- [ ] 单测:Excel 读写、提示词渲染、取图、命名去重(API 用 mock);可在 Python 3.7 运行、不依赖 GUI
|
- [ ] 单测:Excel 读写、提示词渲染、取图、命名去重(API 用 mock);可在 Python 3.7 运行、不依赖 GUI
|
||||||
|
|
||||||
### 19.2 批量编排 — docs/11 §14 阶段 2
|
### 19.2 批量编排 — docs/11 §14 阶段 2
|
||||||
|
|||||||
@@ -0,0 +1,123 @@
|
|||||||
|
"""Tests for single-row AI outfit generation core."""
|
||||||
|
import shutil
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestAiOutfitCore(unittest.TestCase):
|
||||||
|
def setUp(self):
|
||||||
|
self.tmp = Path(tempfile.mkdtemp())
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
shutil.rmtree(str(self.tmp), ignore_errors=True)
|
||||||
|
|
||||||
|
def _task(self, product_id="TY001"):
|
||||||
|
from core.models import OutfitTask
|
||||||
|
|
||||||
|
return OutfitTask(
|
||||||
|
row_index=2,
|
||||||
|
title="纯棉短袖",
|
||||||
|
product_id=product_id,
|
||||||
|
garment_path=str(self.tmp / "garment.png"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _image_bytes(self, size=(640, 480), color=(200, 120, 80)):
|
||||||
|
img = Image.new("RGB", size, color)
|
||||||
|
buffer = BytesIO()
|
||||||
|
img.save(buffer, format="PNG")
|
||||||
|
return buffer.getvalue()
|
||||||
|
|
||||||
|
def test_render_prompt_replaces_placeholders(self):
|
||||||
|
from core.ai_outfit import render_prompt
|
||||||
|
|
||||||
|
prompt = render_prompt("商品 {title} / {product_id}", self._task())
|
||||||
|
|
||||||
|
self.assertEqual(prompt, "商品 纯棉短袖 / TY001")
|
||||||
|
|
||||||
|
def test_safe_product_filename_replaces_invalid_chars(self):
|
||||||
|
from core.ai_outfit import safe_product_filename
|
||||||
|
|
||||||
|
self.assertEqual(safe_product_filename('TY:00/1*?"'), "TY_00_1___")
|
||||||
|
|
||||||
|
def test_make_output_path_avoids_overwrite(self):
|
||||||
|
from core.ai_outfit import make_outfit_output_path
|
||||||
|
|
||||||
|
first = make_outfit_output_path(self.tmp, "TY001")
|
||||||
|
first.write_text("exists")
|
||||||
|
|
||||||
|
second = make_outfit_output_path(self.tmp, "TY001")
|
||||||
|
|
||||||
|
self.assertEqual(second.name, "TY001_1.jpg")
|
||||||
|
|
||||||
|
def test_save_jpg_under_limit_outputs_square_jpg(self):
|
||||||
|
from core.ai_outfit import save_jpg_under_limit
|
||||||
|
|
||||||
|
out = self.tmp / "out.jpg"
|
||||||
|
save_jpg_under_limit(self._image_bytes(size=(640, 480)), out, quality=85, max_bytes=200000)
|
||||||
|
|
||||||
|
self.assertTrue(out.exists())
|
||||||
|
self.assertLessEqual(out.stat().st_size, 200000)
|
||||||
|
with Image.open(str(out)) as img:
|
||||||
|
self.assertEqual(img.format, "JPEG")
|
||||||
|
self.assertEqual(img.size[0], img.size[1])
|
||||||
|
|
||||||
|
def test_generate_outfit_image_success(self):
|
||||||
|
from core.ai_outfit import generate_outfit_image
|
||||||
|
|
||||||
|
client = _FakeClient(self._image_bytes())
|
||||||
|
|
||||||
|
result = generate_outfit_image(
|
||||||
|
self._task(),
|
||||||
|
"为 {title} 生成 {product_id}",
|
||||||
|
self.tmp,
|
||||||
|
model_config={"url": "https://api", "model": "m", "api_key": "k"},
|
||||||
|
api_client=client,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(result.success, result.error)
|
||||||
|
self.assertTrue(Path(result.output_path).exists())
|
||||||
|
self.assertEqual(client.prompt, "为 纯棉短袖 生成 TY001")
|
||||||
|
self.assertTrue(client.image_path.endswith("garment.png"))
|
||||||
|
|
||||||
|
def test_generate_outfit_image_failure(self):
|
||||||
|
from core.ai_outfit import generate_outfit_image
|
||||||
|
|
||||||
|
with self.assertLogs("core.ai_outfit", level="ERROR"):
|
||||||
|
result = generate_outfit_image(
|
||||||
|
self._task(),
|
||||||
|
"prompt",
|
||||||
|
self.tmp,
|
||||||
|
model_config={},
|
||||||
|
api_client=_FailingClient(),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(result.success)
|
||||||
|
self.assertIn("boom", result.error)
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeClient:
|
||||||
|
def __init__(self, image_bytes):
|
||||||
|
self._image_bytes = image_bytes
|
||||||
|
self.prompt = None
|
||||||
|
self.image_path = None
|
||||||
|
|
||||||
|
def generate(self, prompt, image_path, resolution="1K"):
|
||||||
|
self.prompt = prompt
|
||||||
|
self.image_path = str(image_path)
|
||||||
|
return self._image_bytes
|
||||||
|
|
||||||
|
|
||||||
|
class _FailingClient:
|
||||||
|
def generate(self, prompt, image_path, resolution="1K"):
|
||||||
|
raise RuntimeError("boom")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user