diff --git a/src/core/ai_outfit.py b/src/core/ai_outfit.py new file mode 100644 index 0000000..2e04fe9 --- /dev/null +++ b/src/core/ai_outfit.py @@ -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/.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) diff --git a/tasks.md b/tasks.md index aa8b4ef..909d9a6 100644 --- a/tasks.md +++ b/tasks.md @@ -1049,7 +1049,7 @@ - [x] `core/models.py` 新增 `OutfitTask` / `OutfitResult`(纯 dataclass,Python 3.7 兼容,不依赖 PySide6) - [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 写法 -- [ ] `core/ai_outfit.py`:单行生成纯逻辑编排(提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`) +- [x] `core/ai_outfit.py`:单行生成纯逻辑编排(提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`) - [ ] 单测:Excel 读写、提示词渲染、取图、命名去重(API 用 mock);可在 Python 3.7 运行、不依赖 GUI ### 19.2 批量编排 — docs/11 §14 阶段 2 diff --git a/tests/test_ai_outfit.py b/tests/test_ai_outfit.py new file mode 100644 index 0000000..b664366 --- /dev/null +++ b/tests/test_ai_outfit.py @@ -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()