92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
import base64
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
import app as worker
|
|
|
|
|
|
@pytest.mark.parametrize(("densities", "labels"), [
|
|
(["dense"], ["适中偏多"]),
|
|
(["sparse", "normal", "dense"], ["稀疏", "正常", "适中偏多"]),
|
|
([], []),
|
|
])
|
|
def test_interface22_selects_prompts_and_aligns_output_size(monkeypatch, densities, labels):
|
|
from hairline import service
|
|
|
|
prompts = []
|
|
rounded_png = cv2.imencode(".png", np.zeros((15, 13, 3), np.uint8))[1].tobytes()
|
|
|
|
def generate(image, styles, defaults, *, prompt, hairline_height, **kwargs):
|
|
assert styles == [1, 3]
|
|
assert hairline_height == -1
|
|
prompts.append(prompt)
|
|
return [
|
|
{"order": order, "hairline_type": key, "grown_png": rounded_png,
|
|
"overlay": np.zeros((16, 19, 4), np.uint8)}
|
|
for order, key in [(1, "ellipse"), (3, "heart")]
|
|
]
|
|
|
|
monkeypatch.setattr(service, "generate_grow_results_swap", generate)
|
|
image = cv2.imencode(".png", np.zeros((16, 19, 3), np.uint8))[1].tobytes()
|
|
form = {"gender": "female", "hair_style": "1,3", "hairline_height": "-1"}
|
|
if densities:
|
|
form["hair_density"] = densities
|
|
client = TestClient(worker.app, client=("127.0.0.1", 50000))
|
|
response = client.post("/api/v1/hair/grow-22", data=form,
|
|
files={"image_file": ("input.png", image, "image/png")})
|
|
if not densities:
|
|
assert response.status_code == 422
|
|
assert prompts == []
|
|
return
|
|
assert response.status_code == 200
|
|
assert prompts == [f"填充遮罩区域的头发，发量{label}，人物增加一点磨皮" for label in labels]
|
|
results = response.json()["data"]["results"]
|
|
assert [(item["hair_density"], item["order"]) for item in results] == [
|
|
(density, order) for density in densities for order in (1, 3)
|
|
]
|
|
for item in results:
|
|
jpg = np.frombuffer(base64.b64decode(item["grown_image_base64"]), np.uint8)
|
|
assert cv2.imdecode(jpg, cv2.IMREAD_COLOR).shape[:2] == (16, 19)
|
|
|
|
|
|
@pytest.mark.parametrize("dy", [-5, 5])
|
|
def test_height_moves_paste_and_redraw_band_after_swap(monkeypatch, dy):
|
|
from face_analysis import hairline_grow as grow
|
|
|
|
image = np.zeros((80, 80, 3), np.uint8)
|
|
source_mask = np.zeros((80, 80), dtype=bool)
|
|
source_mask[25:45, 20:60] = True
|
|
inner = np.array([[20, 42], [59, 42]])
|
|
outer = np.array([[20, 32], [59, 32]])
|
|
upper = np.ones_like(source_mask)
|
|
monkeypatch.setattr(grow.detector, "detect", lambda image: object())
|
|
monkeypatch.setattr(grow, "estimate_scale_factor", lambda *args: 50)
|
|
monkeypatch.setattr(grow, "compute_mask", lambda *args, **kwargs: (
|
|
source_mask, {"_inner_pts": inner, "_outer_pts": outer, "_upper_mask": upper,
|
|
"mask_pixels": int(source_mask.sum())}))
|
|
|
|
def swap(original, hairline_id, is_hr, ext_mask, denoising_strength, **kwargs):
|
|
assert original is image
|
|
assert ext_mask is source_mask
|
|
generated = np.zeros_like(image)
|
|
generated[30, :, :] = 180
|
|
return generated
|
|
|
|
def composite(original, generated, mask, *args, **kwargs):
|
|
assert np.flatnonzero(mask[:, 40]).tolist() == list(range(25 + dy, 45 + dy))
|
|
assert np.flatnonzero(generated[:, 40, 0]).tolist() == [30 + dy]
|
|
return np.where(mask[..., None], generated, original), mask.astype(np.float32)
|
|
|
|
def encode_red_mask(mask, h, w):
|
|
assert np.flatnonzero(mask[:, 40]).tolist() == list(range(27 + dy, 38 + dy))
|
|
return "verified"
|
|
|
|
monkeypatch.setattr(grow, "_call_swap", swap)
|
|
monkeypatch.setattr(grow, "_composite", composite)
|
|
monkeypatch.setattr(grow, "_red_mask_b64", encode_red_mask)
|
|
result = grow.generate_hairline_redraw(image, "chang_tuoyuan", paste_dy=dy)
|
|
assert result["steps"]["redraw_band_mask_base64"] == "verified"
|