icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement BiRefNet and HAT Sharper models with background removal capabilities

Browse Source 
- 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:
马南
2026-09-10 20:50:44 +08:00
parent 8ae3c1d5b7
commit e8129fe495
70 changed files with 1890 additions and 1208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/tools/models/test_cli.py
View File

@@ -50,10 +50,10 @@ def test_valid_model_cache_skips_all_preparation():
@pytest.mark.parametrize("valid", [True, False])
def test_pisa_export_validates_all_outputs_before_replacing(tmp_path, valid):
def test_isolated_export_validates_all_outputs_before_replacing(tmp_path, valid):
module = load_sync_models()
digest = module.hashlib.sha256(b"new").hexdigest()
models = tuple(module.ModelFile("PISA_SR", "pisa-sr", name, 3, digest) for name in ("a.onnx", "b.onnx"))
models = tuple(module.ModelFile("HAT_GAN_X4_SHARPER", "hat-gan-x4-sharper", name, 3, digest) for name in ("a.onnx", "b.onnx"))
def export(command, **kwargs):
destination = Path(command[-1])
@@ -71,7 +71,7 @@ def test_pisa_export_validates_all_outputs_before_replacing(tmp_path, valid):
module.prepare_model(models[0])
assert not models[0].destination.exists()
assert models[1].destination.read_bytes() == b"old"
assert not list(tmp_path.glob(".pisa-sr-*"))
assert not list(tmp_path.glob(".hat-gan-x4-sharper-*"))
def test_local_export_failure_cleans_temporary_output(tmp_path):
@@ -120,7 +120,7 @@ class SyncModelsTests(unittest.TestCase):
builds = self.sync.MODEL_BUILDS
self.assertEqual(
set(builds),
{"REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4"},
{"REALESRGAN_GENERAL_WDN_X4V3"},
)
for build in builds.values():
self.assertTrue(build.source_url.startswith("https://github.com/"))

63
tests/tools/nodes/test_depth_relief_planes.py
View File

@@ -235,10 +235,67 @@ def test_surface_interface_and_node_structure(surface):
validate_node_structure(obj.modifiers[0].node_group)
@pytest.mark.parametrize("thickness", [0.0, 0.4])
@pytest.mark.parametrize("direction", [(0.2, -0.1, 1), (-0.2, 0.1, 1), (0.4, -0.2, 2), (0, 0, 0), (1, 0, 0)])
def test_relief_direction_adjusts_depth_and_keeps_base(relief, thickness, direction):
obj, set_value, _, _, image = relief
pixels = np.ones((8, 16, 4), dtype=np.float32)
pixels[..., 2] = 2
image.pixels.foreach_set(pixels.ravel())
image.update()
set_value("Thickness", thickness)
set_value("Reference Depth", 2.0)
layers = evaluated(obj)[0]
weights = np.clip(layers[:, 2] / thickness, 0, 1) if thickness else np.ones(len(layers))
set_value("Depth Scale", 0.5)
base = evaluated(obj)
set_value("Depth Direction", direction)
result = evaluated(obj)
front = base[0][:, 2] > 0
expected = base[0].copy()
if direction[2]:
expected[front, 2] += 0.5 * weights[front] * (expected[front, :2] @ np.array(direction[:2])) / direction[2]
np.testing.assert_allclose(result[0], expected, atol=1e-6)
assert result[1] == base[1]
np.testing.assert_array_equal(result[2], base[2])
@pytest.mark.parametrize("thickness", [0.0, 0.4])
def test_relief_smooth_preserves_topology_uv_and_materials(relief, thickness):
obj, set_value, _, _, image = relief
pixels = np.asarray(image.pixels[:], dtype=np.float32).reshape(8, 16, 4)
pixels[..., 2] += (np.arange(16) % 2) * 0.5
image.pixels.foreach_set(pixels.ravel())
image.update()
set_value("Thickness", thickness)
set_value("Reference Depth", 2.0)
set_value("Depth Scale", 1.0)
set_value("Depth Direction", (0.1, -0.1, 1))
base = evaluated(obj)
set_value("Smooth", 3)
result = evaluated(obj)
assert np.isfinite(result[0]).all()
assert np.max(abs(result[0] - base[0])) > 1e-4
assert result[1] == base[1]
np.testing.assert_array_equal(result[2], base[2])
assert result[3:] == base[3:]
set_value("Smooth Weight", 0.0)
np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)
set_value("Smooth Weight", 1.0)
set_value("Smooth", 0)
np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)
def test_relief_interface_and_node_structure(relief):
obj, _, _, inputs, _ = relief
assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale",
"Reference Depth", "Uniform Scale", "Depth Image"]
assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale", "Smooth",
"Reference Depth", "Smooth Weight", "Depth Direction", "Uniform Scale", "Depth Image"]
assert inputs["Smooth"].default_value == 0
assert (inputs["Smooth"].min_value, inputs["Smooth"].max_value) == (0, 20)
assert inputs["Smooth Weight"].default_value == 1
assert inputs["Smooth Weight"].subtype == "FACTOR"
assert tuple(inputs["Depth Direction"].default_value) == (0, 0, 1)
assert inputs["Depth Direction"].subtype == "DIRECTION"
assert inputs["Thickness"].default_value == 0.5
assert inputs["Subdivide"].default_value == 6
assert inputs["Depth Scale"].default_value == 0.5
@@ -356,6 +413,8 @@ def test_plane_interface_panels(source, kind):
controls = ["Reference Depth"]
if kind == "DEPTH":
controls += ["Normal Smooth", "Smooth Weight", "Validity Threshold"]
else:
controls += ["Smooth Weight", "Depth Direction"]
assert [item.name for item in panels[0].interface_items] == controls
assert [item.name for item in panels[1].interface_items] == ["Uniform Scale", "Depth Image"]
assert all(panel.default_closed for panel in panels)

2
tests/tools/nodes/test_shadeless.py
View File

@@ -1,6 +1,7 @@
"""Render the bundled Shadeless group's opacity and HDR color endpoints."""
import bpy
import addon_utils
import numpy as np
from anyimage.common.node import load_node_group
@@ -8,6 +9,7 @@ from anyimage.common.node import load_node_group
def test_shadeless_alpha_is_mixed_inside_the_group(tmp_path):
bpy.ops.wm.read_factory_settings(use_empty=True)
addon_utils.enable("cycles", default_set=False)
group = load_node_group("O Shadeless")
inputs = {s.name: s for s in group.interface.items_tree if s.item_type == "SOCKET" and s.in_out == "INPUT"}
assert list(inputs) == ["Color", "Alpha"] and inputs["Alpha"].default_value == 1

22
tests/tools/test_pack.py
View File

@@ -47,7 +47,7 @@ class ExtensionPackagingTest(unittest.TestCase):
records = {}
with tempfile.TemporaryDirectory() as directory:
directory_path = Path(directory)
wheel_path = directory_path / "blendjob-0.1.16-py3-none-any.whl"
wheel_path = directory_path / "blendjob-0.1.17-py3-none-any.whl"
wheel_path.write_bytes(b"wheel")
original_dist = pack.DIST_DIR
pack.DIST_DIR = directory_path / "dist"
@@ -141,7 +141,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertNotIn("server/selection/polygon_mask.py", names)
self.assertNotIn("server/jobs/prepare_geometry_data.py", names)
self.assertNotIn("server/geometry/cutout.py", names)
self.assertIn("server/models/ben2.py", names)
self.assertIn("server/models/background.py", names)
self.assertNotIn("server/models/da3.py", names)
self.assertIn("server/models/moge.py", names)
self.assertNotIn("server/jobs/debug_predictions.py", names)
@@ -150,6 +150,8 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertIn("server/model_catalog.py", names)
self.assertIn("server/model_download.py", names)
self.assertIn("server/models/onnx_ben2.py", names)
self.assertIn("server/models/onnx_birefnet.py", names)
self.assertNotIn("server/models/onnx_pisa_sr.py", names)
self.assertNotIn("server/models/onnx_da3.py", names)
self.assertIn("server/models/onnx_moge2.py", names)
self.assertIn("server/models/onnx_moge3.py", names)
@@ -159,7 +161,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertIn("operators/cutout_tool/polygon.py", names)
self.assertIn("licenses/scikit-image.txt", names)
self.assertIn(
"wheels/blendjob-0.1.16-py3-none-any.whl",
"wheels/blendjob-0.1.17-py3-none-any.whl",
names,
)
self.assertFalse(any(name.startswith("blendjob/") for name in names))
@@ -177,7 +179,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertEqual(
pack.project_dependencies(),
[
"blendjob==0.1.16",
"blendjob==0.1.17",
"numpy>=1.26,<2.0",
"scipy==1.15.3",
],
@@ -188,7 +190,7 @@ class ExtensionPackagingTest(unittest.TestCase):
wheels = [
Path(name)
for name in (
"blendjob-0.1.16-py3-none-any.whl",
"blendjob-0.1.17-py3-none-any.whl",
"numpy-1.26.4-cp311-cp311-win_amd64.whl",
"scipy-1.15.3-cp311-cp311-win_amd64.whl",
)
@@ -217,12 +219,12 @@ class ExtensionPackagingTest(unittest.TestCase):
if target.name == "3.13":
constraints = Path(command[command.index("--constraint") + 1])
self.assertEqual(
constraints.read_text(), "blendjob==0.1.16"
constraints.read_text(), "blendjob==0.1.17"
)
else:
self.assertNotIn("--constraint", command)
for name in (
"blendjob-0.1.16-py3-none-any.whl",
"blendjob-0.1.17-py3-none-any.whl",
f"scipy-1.15.3-{tag}-{tag}-{config.pip_platform}.whl",
f"numpy-2.1.0-{tag}-{tag}-{config.pip_platform}.whl",
):
@@ -238,7 +240,7 @@ class ExtensionPackagingTest(unittest.TestCase):
("numpy>=1.26,<2.0", "numpy>=2.1,<3.0"),
):
command = call.args[0]
self.assertIn("blendjob==0.1.16", command)
self.assertIn("blendjob==0.1.17", command)
self.assertIn(str(pack.SOURCE_DIR / "wheels"), command)
self.assertIn(f"--python-version={version}", command)
self.assertIn("--implementation=cp", command)
@@ -272,7 +274,7 @@ class ExtensionPackagingTest(unittest.TestCase):
def fake_run(command, check):
target = Path(command[command.index("--dest") + 1])
if target.name == "3.11":
(target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel")
(target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel")
with patch.object(pack.subprocess, "run", side_effect=fake_run):
with self.assertRaisesRegex(RuntimeError, "Python 3.13"):
@@ -341,7 +343,7 @@ class ExtensionPackagingTest(unittest.TestCase):
target = Path(command[command.index("--dest") + 1])
if target.name == "3.13":
raise RuntimeError("download failed")
(target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel")
(target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel")
with (
patch.object(pack, "SOURCE_DIR", source),