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

Add cutout profile and teddy balloon height calculations

Browse Source 
- Implemented cotangent Poisson field and Poisson balloon profile in cutout_profile.py.
- Added teddy balloon heights calculation and related functions in cutout_teddy.py.
- Introduced interaction utilities for image editing in interaction.py.
- Developed depth normalization and analysis functions in cutout_depth.py.
- Created a validation script for cutout assets in validate_cutout_asset.py.
This commit is contained in:
马南
2026-08-27 22:46:57 +08:00
parent 6c8358f92f
commit 7d39fe8c5e
40 changed files with 3058 additions and 3896 deletions
Download Patch File Download Diff File Expand all files Collapse all files

206
tests/test_server.py
View File

@@ -41,7 +41,7 @@ class ServerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.server_module = load_server_module()
importlib.import_module("server.jobs.balloon")
importlib.import_module("server.jobs.cutout_depth")
importlib.import_module("server.jobs.mesh")
from blendjob import JobServer
from server.model_manager import ModelManager, ModelSessionCache
@@ -415,32 +415,6 @@ class ServerTest(unittest.TestCase):
self.assertEqual(status["state"], "succeeded")
def test_balloon_lasso_input_crops_and_clears_outside_rgba(self):
import numpy as np
from PIL import Image
balloon = sys.modules["server.jobs.balloon"]
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
source_path = root / "source.png"
output_path = root / "lasso.png"
pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8)
pixels[1, 1, 3] = 64
Image.fromarray(pixels, mode="RGBA").save(source_path)
_path, _selection, origin, source_size = balloon.create_lasso_input(
source_path,
((1, 1), (5, 1), (1, 5)),
output_path,
)
result = np.asarray(Image.open(output_path).convert("RGBA"))
self.assertEqual(origin, (1, 1))
self.assertEqual(source_size, (6, 6))
self.assertEqual(result.shape, (4, 4, 4))
self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64))
self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0))
def test_polygon_selection_mask_uses_even_odd_fill_for_self_intersection(self):
image_mask = sys.modules["server.jobs.image_mask"]
polygon = (
@@ -783,7 +757,7 @@ class ServerTest(unittest.TestCase):
)
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
depth_path, normal_path = mesh_job._write_balloon_depth(
depth_path, normal_path = mesh_job._write_depth_balloon(
context,
{"device": "cpu"},
source_path,
@@ -909,53 +883,10 @@ class ServerTest(unittest.TestCase):
)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
def test_balloon_job_sends_the_masked_crop_to_ben2(self):
import numpy as np
from PIL import Image
balloon = sys.modules["server.jobs.balloon"]
source_path = self.storage_root / "source.png"
pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8)
Image.fromarray(pixels, mode="RGBA").save(source_path)
job_directory = self.storage_root / "job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
resource=lambda _name: SimpleNamespace(
directory=lambda _model: self.storage_root / "model"
),
)
received = []
def infer(input_path, *_args, **_kwargs):
image = Image.open(input_path).convert("RGBA")
received.append(np.asarray(image))
return image, 0.0, 1.0
with patch.object(balloon, "infer_foreground", side_effect=infer):
result = balloon.run(
context,
{
"input": str(source_path),
"polygon": ((1, 1), (5, 1), (1, 5)),
"contour_analysis": True,
"depth_analysis": False,
"contour_model": "BEN2_BASE",
"device": "cpu",
},
)
self.assertEqual(result, {"contour": "contour.json"})
self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0))
self.assertTrue((job_directory / "lasso-input.png").is_file())
self.assertTrue((job_directory / "foreground.png").is_file())
self.assertTrue((job_directory / "contour.json").is_file())
def test_balloon_depth_analysis_normalizes_moge_distance(self):
import numpy as np
balloon = sys.modules["server.jobs.balloon"]
balloon = sys.modules["server.jobs.cutout_depth"]
y, x = np.indices((7, 9), dtype=np.float32)
depth = 2.0 + x * 0.1 + y * 0.2
depth[3, 4] -= 0.5
@@ -996,10 +927,44 @@ class ServerTest(unittest.TestCase):
places=5,
)
def test_balloon_depth_plane_is_stable_for_a_slender_diagonal_mask(self):
import numpy as np
balloon = sys.modules["server.jobs.cutout_depth"]
y, x = np.indices((101, 101), dtype=np.float64)
normalized_x = x / 50.0 - 1.0
normalized_y = y / 50.0 - 1.0
distance = y - x
selection = np.abs(distance) <= 1.0
along = (normalized_x + normalized_y) * 0.5
depth = 4.0 + 0.4 * along
depth += np.where(distance > 0.0, 0.05, 0.0)
depth -= np.where(distance < 0.0, 0.05, 0.0)
_signal, valid, _reference, baseline = balloon.normalized_moge_depth(
depth,
np.ones_like(selection),
selection,
)
fitted = (
baseline[0] * normalized_x
+ baseline[1] * normalized_y
+ baseline[2]
)
self.assertLessEqual(
float(np.ptp(fitted)),
float(np.ptp(fitted[valid])) * 1.25,
)
self.assertGreater(
abs(float(np.corrcoef(fitted[valid], along[valid])[0, 1])),
0.99,
)
def test_balloon_relative_depth_ignores_absolute_distance_and_scale(self):
import numpy as np
balloon = sys.modules["server.jobs.balloon"]
balloon = sys.modules["server.jobs.cutout_depth"]
y, x = np.indices((8, 10), dtype=np.float32)
depth = 4.0 + x * 0.04 + y * 0.03
depth[2:6, 3:7] -= np.asarray(
@@ -1022,100 +987,5 @@ class ServerTest(unittest.TestCase):
np.testing.assert_allclose(rescaled[0], original[0], atol=2e-4)
np.testing.assert_allclose(rescaled[3], original[3], atol=2e-4)
def test_balloon_depth_job_sends_masked_crop_only_to_moge(self):
import numpy as np
from PIL import Image
balloon = sys.modules["server.jobs.balloon"]
source_path = self.storage_root / "depth-source.png"
pixels = np.full((6, 6, 4), (80, 100, 120, 255), dtype=np.uint8)
Image.fromarray(pixels, mode="RGBA").save(source_path)
job_directory = self.storage_root / "depth-job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
resource=lambda _name: SimpleNamespace(),
)
received = []
def infer(_manager, _parameters, images, **_kwargs):
image = Image.open(images[0]).convert("RGBA")
received.append(np.asarray(image))
height, width = image.size[1], image.size[0]
depth = np.linspace(
1.0,
2.0,
height * width,
dtype=np.float32,
).reshape(height, width)
return SimpleNamespace(
depth=depth[None],
mask=np.ones((1, height, width), dtype=bool),
normal=np.broadcast_to(
np.asarray((0.0, 0.0, -1.0), dtype=np.float32),
(1, height, width, 3),
).copy(),
)
with (
patch.object(balloon.moge2, "infer", side_effect=infer),
patch.object(
balloon,
"infer_foreground",
side_effect=AssertionError("BEN2 should not run"),
),
):
result = balloon.run(
context,
{
"input": str(source_path),
"polygon": ((1, 1), (5, 1), (1, 5)),
"contour_analysis": False,
"depth_analysis": True,
"depth_model": "MOGE2_VITS_NORMAL",
"resolution_level": 5,
"device": "cpu",
},
)
self.assertEqual(
result,
{"depth": "depth.npz", "normal": "normal.png"},
)
self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0))
self.assertFalse((job_directory / "foreground.png").exists())
with np.load(job_directory / "depth.npz", allow_pickle=False) as data:
self.assertEqual(data["depth_signal"].shape, (4, 4))
self.assertEqual(data["baseline_vector"].shape, (3,))
self.assertEqual(tuple(data["origin"]), (1, 1))
self.assertEqual(tuple(data["source_size"]), (6, 6))
self.assertTrue(data["valid"].any())
normal = np.asarray(Image.open(job_directory / "normal.png").convert("RGB"))
self.assertEqual(normal.shape, (6, 6, 3))
self.assertEqual(tuple(normal[0, 0]), (128, 128, 255))
def test_balloon_contour_uses_alpha_above_half_and_largest_component(self):
import numpy as np
balloon = sys.modules["server.jobs.balloon"]
alpha = np.zeros((6, 7), dtype=np.uint8)
alpha[0, 0] = 255
alpha[1:4, 2:5] = 128
alpha[1, 1] = 127
selection_mask = np.ones_like(alpha, dtype=bool)
contour = balloon.contour_from_alpha(
alpha,
selection_mask,
origin=(10, 20),
)
self.assertEqual(
set(map(tuple, contour)),
{(12, 21), (15, 21), (15, 24), (12, 24)},
)
if __name__ == "__main__":
unittest.main()