130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
"""
|
|
lineart-extractor 单元测试
|
|
==========================
|
|
运行: pytest tests/ -v
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from lineartization import (
|
|
LineArtConfig,
|
|
extract_lineart,
|
|
extract_lineart_file,
|
|
load_image,
|
|
save_image,
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 测试用图: 合成"白底 + 黑字 + 彩色块"
|
|
# --------------------------------------------------------------------------- #
|
|
@pytest.fixture
|
|
def sample_image():
|
|
"""构造一张 400x600 的合成图: 白底 + 黑色矩形(模拟文字) + 彩色块。"""
|
|
img = np.full((400, 600, 3), 255, np.uint8)
|
|
|
|
# 中央"文字区": 密集小黑块
|
|
rng = np.random.default_rng(42)
|
|
for _ in range(120):
|
|
x = rng.integers(180, 420)
|
|
y = rng.integers(150, 250)
|
|
img[y:y + 4, x:x + 4] = 0
|
|
|
|
# 左侧彩色块(模拟山体)
|
|
img[60:160, 20:180] = (60, 160, 80) # 绿
|
|
img[160:220, 20:180] = (80, 120, 200) # 偏蓝
|
|
|
|
# 右侧一个红色圆(模拟灯笼)
|
|
import cv2
|
|
cv2.circle(img, (500, 120), 40, (40, 40, 200), 3)
|
|
|
|
return img
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 测试
|
|
# --------------------------------------------------------------------------- #
|
|
def test_load_save_roundtrip(tmp_path, sample_image):
|
|
"""读写往返一致。"""
|
|
p = tmp_path / "in.png"
|
|
save_image(str(p), sample_image)
|
|
loaded = load_image(str(p))
|
|
assert loaded.shape == sample_image.shape
|
|
assert np.allclose(loaded, sample_image, atol=2)
|
|
|
|
|
|
def test_extract_returns_binary(sample_image):
|
|
"""输出必须是二值(0/255)白底黑线。"""
|
|
out = extract_lineart(sample_image, LineArtConfig(enable_green_smoothing=False))
|
|
assert out.dtype == np.uint8
|
|
assert out.ndim == 2
|
|
uniq = np.unique(out)
|
|
assert set(uniq.tolist()).issubset({0, 255})
|
|
assert out.shape == sample_image.shape[:2]
|
|
|
|
|
|
def test_extract_has_content(sample_image):
|
|
"""输出不能空白、也不能全黑。"""
|
|
out = extract_lineart(sample_image, LineArtConfig(enable_green_smoothing=False))
|
|
black_ratio = (out < 128).mean() * 100
|
|
assert 0.1 < black_ratio < 90.0
|
|
|
|
|
|
def test_line_width_effect(sample_image):
|
|
"""线宽参数应影响黑占比(越粗越多)。"""
|
|
cfg1 = LineArtConfig(line_width=1, enable_green_smoothing=False)
|
|
cfg3 = LineArtConfig(line_width=3, enable_green_smoothing=False)
|
|
r1 = (extract_lineart(sample_image, cfg1) < 128).mean()
|
|
r3 = (extract_lineart(sample_image, cfg3) < 128).mean()
|
|
assert r3 > r1
|
|
|
|
|
|
def test_green_smoothing_toggle(sample_image):
|
|
"""绿块抹平开关都应能正常出图。"""
|
|
for flag in (True, False):
|
|
cfg = LineArtConfig(enable_green_smoothing=flag)
|
|
out = extract_lineart(sample_image, cfg)
|
|
assert (out < 128).mean() > 0
|
|
|
|
|
|
def test_protect_areas(sample_image):
|
|
"""保护区域内的线条不应被删。"""
|
|
cfg = LineArtConfig(
|
|
enable_green_smoothing=False,
|
|
protect_areas=[(0, 200, 0, 400)],
|
|
)
|
|
out = extract_lineart(sample_image, cfg)
|
|
assert (out < 128).sum() > 0
|
|
|
|
|
|
def test_file_interface(tmp_path, sample_image):
|
|
"""extract_lineart_file 接口正常。"""
|
|
src = tmp_path / "src.png"
|
|
dst = tmp_path / "dst.png"
|
|
save_image(str(src), sample_image)
|
|
result = extract_lineart_file(str(src), str(dst),
|
|
LineArtConfig(enable_green_smoothing=False))
|
|
assert os.path.exists(result)
|
|
assert result == str(dst)
|
|
|
|
|
|
def test_load_missing_file():
|
|
"""读取不存在的文件应抛异常。"""
|
|
with pytest.raises((FileNotFoundError, Exception)):
|
|
load_image("___no_such_file___.png")
|
|
|
|
|
|
def test_chinese_path(tmp_path, sample_image):
|
|
"""中文路径应正常工作。"""
|
|
src = tmp_path / "中文图片.png"
|
|
dst = tmp_path / "输出_láthair.png"
|
|
save_image(str(src), sample_image)
|
|
out = extract_lineart_file(str(src), str(dst),
|
|
LineArtConfig(enable_green_smoothing=False))
|
|
assert os.path.exists(out)
|