73 lines
2.0 KiB
Python
73 lines
2.0 KiB
Python
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from trellis2.pipelines.trellis2_image_to_3d import Trellis2ImageTo3DPipeline
|
|
|
|
|
|
class _FakeRembg:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def __call__(self, image):
|
|
self.calls += 1
|
|
rgba = image.convert("RGBA")
|
|
alpha = np.zeros((image.height, image.width), dtype=np.uint8)
|
|
alpha[2:6, 2:6] = 255
|
|
rgba.putalpha(Image.fromarray(alpha, mode="L"))
|
|
return rgba
|
|
|
|
|
|
def _pipeline(rembg):
|
|
pipeline = object.__new__(Trellis2ImageTo3DPipeline)
|
|
pipeline.low_vram = False
|
|
pipeline.rembg_model = rembg
|
|
pipeline._device = "cpu"
|
|
return pipeline
|
|
|
|
|
|
def test_rgba_with_transparency_bypasses_rmbg():
|
|
rembg = _FakeRembg()
|
|
rgba = np.full((8, 8, 4), 255, dtype=np.uint8)
|
|
rgba[:, :, :3] = [80, 120, 160]
|
|
rgba[:2, :, 3] = 0
|
|
|
|
output = _pipeline(rembg).preprocess_image(Image.fromarray(rgba, mode="RGBA"))
|
|
|
|
assert rembg.calls == 0
|
|
assert output.mode == "RGB"
|
|
assert output.size[0] > 0 and output.size[1] > 0
|
|
|
|
|
|
def test_opaque_rgba_uses_official_background_remover():
|
|
rembg = _FakeRembg()
|
|
opaque = Image.new("RGBA", (8, 8), (80, 120, 160, 255))
|
|
|
|
output = _pipeline(rembg).preprocess_image(opaque)
|
|
|
|
assert rembg.calls == 1
|
|
assert output.mode == "RGB"
|
|
assert output.size[0] > 0 and output.size[1] > 0
|
|
|
|
|
|
def test_soft_authored_alpha_is_preserved_without_rmbg():
|
|
rembg = _FakeRembg()
|
|
rgba = np.zeros((8, 8, 4), dtype=np.uint8)
|
|
rgba[2:6, 2:6, :3] = [80, 120, 160]
|
|
rgba[2:6, 2:6, 3] = 128
|
|
|
|
output = _pipeline(rembg).preprocess_image(Image.fromarray(rgba, mode="RGBA"))
|
|
|
|
assert rembg.calls == 0
|
|
assert output.size[0] > 0 and output.size[1] > 0
|
|
|
|
|
|
def test_empty_authored_alpha_has_actionable_error():
|
|
rembg = _FakeRembg()
|
|
empty = Image.new("RGBA", (8, 8), (80, 120, 160, 0))
|
|
|
|
with pytest.raises(RuntimeError, match="empty alpha mask"):
|
|
_pipeline(rembg).preprocess_image(empty)
|
|
|
|
assert rembg.calls == 0
|