71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
"""仅扩图的像素保护与单张 AI 编辑提示词回归测试。"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
SERVER_DIR = Path(__file__).resolve().parents[1]
|
|
if str(SERVER_DIR) not in sys.path:
|
|
sys.path.insert(0, str(SERVER_DIR))
|
|
|
|
from services.generator import prepare_outpaint_canvas
|
|
from services.prompts import build_prompt
|
|
|
|
|
|
def _data_url(image: Image.Image) -> str:
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
|
|
class OutpaintTests(unittest.TestCase):
|
|
def test_outpaint_restores_original_region_after_provider_response(self) -> None:
|
|
original = Image.new("RGB", (400, 400), (220, 20, 60))
|
|
canvas = asyncio.run(prepare_outpaint_canvas(_data_url(original), "1152x1536"))
|
|
|
|
self.assertEqual(canvas.size, (1152, 1536))
|
|
self.assertEqual(canvas.box, (0, 192, 1152, 1344))
|
|
|
|
# 模拟供应商把整张图都重绘成蓝色；服务端合成后中心原图必须仍为红色。
|
|
generated = Image.new("RGB", canvas.size, (0, 0, 255))
|
|
restored = Image.open(io.BytesIO(canvas.restore_original(_png_bytes(generated)))).convert("RGB")
|
|
self.assertEqual(restored.getpixel((576, 768)), (220, 20, 60))
|
|
self.assertEqual(restored.getpixel((576, 80)), (0, 0, 255))
|
|
|
|
def test_expanded_margin_is_transparent_in_mask(self) -> None:
|
|
original = Image.new("RGB", (400, 400), (1, 2, 3))
|
|
canvas = asyncio.run(prepare_outpaint_canvas(_data_url(original), "1152x1536"))
|
|
mask_bytes = base64.b64decode(canvas.mask_data_url.split(",", 1)[1])
|
|
mask = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
|
|
self.assertEqual(mask.getpixel((576, 80))[3], 0)
|
|
self.assertEqual(mask.getpixel((576, 768))[3], 255)
|
|
|
|
|
|
def _png_bytes(image: Image.Image) -> bytes:
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
|
|
class ExclusiveEditPromptTests(unittest.TestCase):
|
|
def test_user_edit_requirements_do_not_inherit_marketing_or_light_studio_prompt(self) -> None:
|
|
prompt = build_prompt(
|
|
"pcore", "gpt-image-2", "custom", {"title": "product"}, 5, "ru",
|
|
requirements="背景必须是纯黑色，其他不要改变",
|
|
requirements_exclusive=True,
|
|
)
|
|
self.assertIn("背景必须是纯黑色", prompt)
|
|
self.assertNotIn("e-commerce marketing image", prompt)
|
|
self.assertNotIn("plain light background", prompt)
|
|
self.assertNotIn("Render concise Russian marketing text", prompt)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|