Implement BiRefNet and HAT Sharper models with background removal capabilities
- Added validation documentation for AI model tiers. - Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction. - Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference. - Developed tests for background models, ensuring proper dispatch and session reuse. - Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values. - Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing. - Created `deform_sampling.py` for deformable convolution operations used in BiRefNet. - Added utility functions for compressing ONNX weights without altering graph arithmetic. - Introduced a reference plane node for depth field adjustments in geometry processing.
This commit is contained in:
36
tests/anyimage/server/test_model_birefnet_onnx.py
Normal file
36
tests/anyimage/server/test_model_birefnet_onnx.py
Normal file
@@ -0,0 +1,36 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from server.models import onnx_birefnet
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1024, 2048])
|
||||
def test_probability_alpha_keeps_transparency_and_source_dimensions(size):
|
||||
session = SimpleNamespace(
|
||||
get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, size, size])],
|
||||
get_outputs=lambda: [SimpleNamespace(name="alpha")],
|
||||
)
|
||||
mask = np.full((1, 1, size, size), 0.37, dtype=np.float32)
|
||||
with patch.object(onnx_birefnet, "run_session", return_value=[mask]) as run:
|
||||
result = onnx_birefnet.infer_alpha(session, Image.new("RGB", (31, 17), (64, 128, 192)), release_memory=False)
|
||||
assert result.shape == (17, 31)
|
||||
assert result.dtype == np.float32
|
||||
np.testing.assert_allclose(result, 0.37, atol=1e-6)
|
||||
assert run.call_args.kwargs["release_memory"] is False
|
||||
feed = run.call_args.args[2]["image"]
|
||||
np.testing.assert_allclose(feed[0, :, 0, 0], (np.array([64, 128, 192]) / 255 - [.485, .456, .406]) / [.229, .224, .225], atol=1e-6)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [np.nan, np.inf, -0.1, 1.1])
|
||||
def test_invalid_alpha_is_rejected(value):
|
||||
session = SimpleNamespace(
|
||||
get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, 1024, 1024])],
|
||||
get_outputs=lambda: [SimpleNamespace(name="alpha")],
|
||||
)
|
||||
with patch.object(onnx_birefnet, "run_session", return_value=[np.full((1, 1, 1024, 1024), value)]):
|
||||
with pytest.raises(RuntimeError, match="invalid alpha"):
|
||||
onnx_birefnet.infer_alpha(session, Image.new("RGB", (2, 3)))
|
||||
Reference in New Issue
Block a user