f6e09dbff2
- palette.py + rembg.py: implement from stubs (Pillow median-cut + rembg u2net) - vlm.py: rename Spark2→steev (Strix Halo / Ollama); bump max_tokens 1024→4096 (qwen3-vl:32b thinking mode consumes budget tokens — 4096 min for valid output) - settings.py: rename spark2_vlm_*/spark1_flux_* → vlm_*/flux_*; real defaults (steev 100.88.167.87:11434 Ollama, gx10 100.90.100.10:8188 ComfyUI) - tests/: conftest.py + test_palette.py + test_rembg.py + test_integration_e2e.py (28 unit + 10 integration; 38/38 passing — VLM raw/polished/ugc + FLUX render) - CLAUDE.md: rewrite to accurate phase status + infra + layout - requirements.txt + pyproject.toml: add Pillow, rembg, pytest-asyncio deps Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
"""Unit tests for POST /rembg/cutout."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
|
|
def _make_png_b64(color: tuple[int, int, int] = (180, 100, 50), size: int = 64) -> str:
|
|
img = Image.new("RGB", (size, size), color=color)
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
|
|
|
|
def test_rembg_missing_image_returns_400(client) -> None:
|
|
resp = client.post("/rembg/cutout", json={})
|
|
assert resp.status_code == 400
|
|
|
|
|
|
def test_rembg_bad_base64_returns_400(client) -> None:
|
|
resp = client.post("/rembg/cutout", json={"image_base64": "%%%bad%%%"})
|
|
assert resp.status_code == 400
|
|
|
|
|
|
def test_rembg_returns_png_with_alpha(client, red_png_b64) -> None:
|
|
"""rembg removes background → output is PNG with alpha channel."""
|
|
resp = client.post("/rembg/cutout", json={"image_base64": red_png_b64})
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["content_type"] == "image/png"
|
|
assert "image_base64" in body
|
|
|
|
raw = base64.b64decode(body["image_base64"])
|
|
img = Image.open(io.BytesIO(raw))
|
|
assert img.mode == "RGBA", f"expected RGBA, got {img.mode}"
|
|
|
|
|
|
def test_rembg_alpha_matting_flag_round_trips(client, red_png_b64) -> None:
|
|
resp = client.post("/rembg/cutout", json={"image_base64": red_png_b64, "alpha_matting": False})
|
|
assert resp.status_code == 200
|
|
assert resp.json()["alpha_matting"] is False
|