nan
f42cba1932
- Implemented ExtractImageRegion operator for extracting image regions using box or lasso selection. - Added RemoveImageRegionBackground job for background removal from selected image regions. - Created utility functions for polygon bounds, selection canvas creation, and applying selection masks. - Integrated image processing using PIL and numpy for efficient pixel manipulation. - Enhanced user experience with modal drawing and status updates during extraction.
653 lines
24 KiB
Python
653 lines
24 KiB
Python
import importlib.util
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_server_module():
|
|
server_directory = PROJECT_ROOT / "src" / "anyimage" / "server"
|
|
specification = importlib.util.spec_from_file_location(
|
|
"server",
|
|
server_directory / "__init__.py",
|
|
submodule_search_locations=[str(server_directory)],
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
sys.modules["server"] = module
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FakeCuda:
|
|
class OutOfMemoryError(RuntimeError):
|
|
pass
|
|
|
|
@staticmethod
|
|
def is_available():
|
|
return False
|
|
|
|
@staticmethod
|
|
def empty_cache():
|
|
return None
|
|
|
|
|
|
class ServerTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.server_module = load_server_module()
|
|
|
|
def setUp(self):
|
|
self.jobs_temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.jobs_temp.cleanup)
|
|
self.storage_root = Path(self.jobs_temp.name)
|
|
|
|
def test_required_model_download_skips_ready_model(self):
|
|
manager = self.server_module.ModelManager(Path("."))
|
|
downloads = []
|
|
context = SimpleNamespace(progress=lambda *_args: None)
|
|
with (
|
|
patch.object(
|
|
manager,
|
|
"ready",
|
|
side_effect=lambda key: key == "BEN2_BASE",
|
|
),
|
|
patch.object(
|
|
manager,
|
|
"_download",
|
|
side_effect=lambda _context, key, _progress: downloads.append(key),
|
|
),
|
|
):
|
|
result = manager.download_missing(
|
|
context,
|
|
("MOGE2_VITS_NORMAL", "BEN2_BASE"),
|
|
)
|
|
|
|
self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"])
|
|
self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]})
|
|
|
|
def test_models_are_owned_by_the_server_resource(self):
|
|
self.assertFalse((PROJECT_ROOT / "src" / "anyimage" / "models.py").exists())
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
manager = self.server_module.ModelManager(Path(directory))
|
|
snapshot = manager.snapshot()
|
|
self.assertEqual(
|
|
snapshot,
|
|
{"loaded": {"ben2": "", "da3": "", "moge2": ""}},
|
|
)
|
|
self.assertIn("server.model_catalog", sys.modules)
|
|
self.assertNotIn("model_catalog", sys.modules)
|
|
|
|
def create_cache(self):
|
|
cache = self.server_module.ModelSessionCache.__new__(
|
|
self.server_module.ModelSessionCache
|
|
)
|
|
cache.ben2_model = None
|
|
cache.ben2_key = None
|
|
cache.da3_model = None
|
|
cache.da3_key = None
|
|
cache.moge_model = None
|
|
cache.moge_key = None
|
|
return cache
|
|
|
|
def create_server(self):
|
|
source = self.server_module.server
|
|
server = self.server_module.JobServer(
|
|
"AnyImage Job Server", storage_root=self.storage_root
|
|
)
|
|
server.handlers.update(source.handlers)
|
|
server.resource_factories.update(source.resource_factories)
|
|
server.bind(self.storage_root)
|
|
return server
|
|
|
|
def wait_for_job(self, context, timeout=2.0):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if context._snapshot()["state"] in {"succeeded", "failed", "cancelled"}:
|
|
return context._snapshot()
|
|
time.sleep(0.01)
|
|
self.fail(f"Job did not finish: {context._snapshot()}")
|
|
|
|
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
ben2 = ModuleType("server.models.onnx_ben2")
|
|
ben2.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(sys.modules, {"server.models.onnx_ben2": ben2}):
|
|
first, _load_ms = cache.get_ben2("model", device)
|
|
second, cached_load_ms = cache.get_ben2("model", device)
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cuda")])
|
|
self.assertEqual(cache.snapshot()["ben2"], "model")
|
|
self.assertEqual(cached_load_ms, 0.0)
|
|
|
|
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
da3 = ModuleType("server.models.onnx_da3")
|
|
da3.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"server.models.onnx_da3": da3,
|
|
},
|
|
):
|
|
first = cache.get_da3("model", device)
|
|
second = cache.get_da3("model", device)
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cuda")])
|
|
self.assertEqual(cache.snapshot()["da3"], "model")
|
|
|
|
def test_depth_model_cache_keeps_only_da3_or_moge(self):
|
|
cache = self.create_cache()
|
|
cache.ben2_model = object()
|
|
cache.ben2_key = ("ben2", "cuda")
|
|
cache.moge_model = object()
|
|
cache.moge_key = ("moge", "cuda")
|
|
device = SimpleNamespace(type="cuda")
|
|
da3 = object()
|
|
|
|
def load_da3(_model_dir, _device, key):
|
|
cache.da3_model = da3
|
|
cache.da3_key = key
|
|
return da3
|
|
|
|
with patch.object(cache, "_load_da3", side_effect=load_da3):
|
|
self.assertIs(cache.get_da3("da3", device), da3)
|
|
|
|
self.assertIsNone(cache.moge_model)
|
|
self.assertIsNotNone(cache.ben2_model)
|
|
|
|
moge = object()
|
|
|
|
fake_moge = ModuleType("server.models.onnx_moge2")
|
|
fake_moge.create_session = lambda _model_dir, _device: moge
|
|
with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}):
|
|
self.assertIs(cache.get_moge2("moge", "cuda"), moge)
|
|
|
|
self.assertIsNone(cache.da3_model)
|
|
self.assertIsNotNone(cache.ben2_model)
|
|
|
|
def test_dedicated_server_runs_only_one_job_at_a_time(self):
|
|
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
|
first_started = threading.Event()
|
|
release_first = threading.Event()
|
|
order = []
|
|
|
|
@server.job("example")
|
|
def example(_context, parameters):
|
|
order.append(parameters["name"])
|
|
if parameters["name"] == "first":
|
|
first_started.set()
|
|
release_first.wait(1.0)
|
|
|
|
first = server.submit("example", {"name": "first"}, job_id="job")
|
|
self.assertTrue(first_started.wait(1.0))
|
|
second = server.submit("example", {"name": "second"}, job_id="other")
|
|
|
|
self.assertIsNotNone(first)
|
|
self.assertIsNotNone(second)
|
|
self.assertTrue(server.snapshot("instance")["busy"])
|
|
self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
|
|
self.assertNotIn("clients", server.snapshot("instance"))
|
|
release_first.set()
|
|
self.wait_for_job(first)
|
|
self.wait_for_job(second)
|
|
self.assertEqual(order, ["first", "second"])
|
|
server.close()
|
|
|
|
def test_server_registers_tasks_instead_of_model_names(self):
|
|
server = self.create_server()
|
|
self.assertEqual(
|
|
set(server.handlers),
|
|
{
|
|
"download-model",
|
|
"download-required-models",
|
|
"analyze-balloon",
|
|
"extract-image-background",
|
|
"remove-background",
|
|
"extract-frames",
|
|
"generate-depth",
|
|
"generate-surface",
|
|
"debug-da3",
|
|
"debug-moge2",
|
|
"debug-ben2",
|
|
},
|
|
)
|
|
|
|
def test_queued_job_can_be_cancelled_before_it_runs(self):
|
|
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
|
first_started = threading.Event()
|
|
release_first = threading.Event()
|
|
executed = []
|
|
|
|
@server.job("example")
|
|
def example(_context, parameters):
|
|
executed.append(parameters["name"])
|
|
if parameters["name"] == "first":
|
|
first_started.set()
|
|
release_first.wait(1.0)
|
|
|
|
first = server.submit("example", {"name": "first"})
|
|
self.assertTrue(first_started.wait(1.0))
|
|
second = server.submit("example", {"name": "second"})
|
|
|
|
self.assertIs(server.cancel(second.job_id), second)
|
|
release_first.set()
|
|
self.wait_for_job(first)
|
|
second_status = self.wait_for_job(second)
|
|
|
|
self.assertEqual(second_status["state"], "cancelled")
|
|
self.assertEqual(executed, ["first"])
|
|
server.close()
|
|
|
|
def test_cancelling_job_keeps_server_and_model_cache_alive(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": {"ben2": "BEN2_BASE", "da3": ""}},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.server_module, "ModelManager", return_value=fake_models):
|
|
server = self.create_server()
|
|
started = threading.Event()
|
|
|
|
@server.job("cancellable")
|
|
def cancellable(context, _parameters):
|
|
started.set()
|
|
while not context._is_cancelled():
|
|
time.sleep(0.001)
|
|
context.check_cancelled()
|
|
|
|
context = server.submit("cancellable", {}, job_id="job")
|
|
self.assertTrue(started.wait(1.0))
|
|
self.assertIsNotNone(server.cancel("job"))
|
|
self.assertTrue(context._is_cancelled())
|
|
status = self.wait_for_job(context)
|
|
|
|
self.assertEqual(status["state"], "cancelled")
|
|
self.assertFalse(server.shutdown_event.is_set())
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"],
|
|
"BEN2_BASE",
|
|
)
|
|
|
|
def test_job_progress_is_kept_in_server_memory(self):
|
|
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
context.progress(0.5, "Half way")
|
|
|
|
context = server.submit("example", {})
|
|
self.wait_for_job(context)
|
|
|
|
self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
|
|
self.assertEqual(context.directory.name, context.job_id)
|
|
self.assertEqual(context._snapshot()["state"], "succeeded")
|
|
self.assertEqual(context._snapshot()["progress"], 1.0)
|
|
|
|
def test_runtime_storage_root_owns_the_jobs_directory(self):
|
|
storage_root = self.storage_root
|
|
server = self.server_module.JobServer("Test")
|
|
server.bind(storage_root)
|
|
|
|
@server.job("example")
|
|
def example(_context, _parameters):
|
|
return {"value": 42}
|
|
|
|
context = server.submit("example", {})
|
|
self.wait_for_job(context)
|
|
|
|
self.assertEqual(server.jobs_directory, storage_root / "jobs")
|
|
self.assertEqual(context.storage_root, storage_root.resolve())
|
|
self.assertEqual(context._snapshot()["result"], {"value": 42})
|
|
|
|
def test_job_server_owns_and_clears_shared_resources(self):
|
|
events = []
|
|
resource = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": "model"},
|
|
clear=lambda: events.append("clear"),
|
|
close=lambda: events.append("close"),
|
|
)
|
|
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
|
server.add_resource("model_manager", resource)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
self.assertIs(context.resource("model_manager"), resource)
|
|
release.wait(1.0)
|
|
|
|
release = threading.Event()
|
|
context = server.submit("example", {}, job_id="job")
|
|
self.assertFalse(server.clear_resource("model_manager"))
|
|
release.set()
|
|
self.wait_for_job(context)
|
|
self.assertTrue(server.clear_resource("model_manager"))
|
|
server.close()
|
|
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["model_manager"]["loaded"],
|
|
"model",
|
|
)
|
|
self.assertEqual(events, ["clear", "close"])
|
|
|
|
def test_server_dispatches_frame_extraction_without_models(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.server_module, "ModelManager", return_value=fake_models):
|
|
server = self.create_server()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
(input_directory / "frame-0001.jpg").touch()
|
|
(input_directory / "frame-0002.jpg").touch()
|
|
context = server.submit(
|
|
"extract-frames",
|
|
{
|
|
"input": str(input_directory),
|
|
"video_frames": 16,
|
|
"start_frame": 1,
|
|
},
|
|
job_id="job",
|
|
)
|
|
status = self.wait_for_job(context)
|
|
self.assertTrue((context.directory / "color.0001.jpg").is_file())
|
|
|
|
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_image_extract_canvas_clears_rgba_outside_the_selection(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
image_extract = sys.modules["server.jobs.image_extract"]
|
|
source_path = self.storage_root / "source.png"
|
|
pixels = np.full((2, 2, 4), (20, 40, 60, 255), dtype=np.uint8)
|
|
Image.fromarray(pixels, mode="RGBA").save(source_path)
|
|
|
|
canvas, _selection, bounds = image_extract.create_selection_canvas(
|
|
source_path,
|
|
((-1, -1), (2, -1), (-1, 2)),
|
|
)
|
|
result = np.asarray(canvas)
|
|
|
|
self.assertEqual(bounds, (-1, -1, 2, 2))
|
|
self.assertEqual(result.shape, (3, 3, 4))
|
|
self.assertEqual(tuple(result[2, 2]), (0, 0, 0, 0))
|
|
self.assertEqual(tuple(result[1, 1]), (20, 40, 60, 255))
|
|
|
|
def test_image_extract_crops_to_alpha_with_pixel_padding(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
image_extract = sys.modules["server.jobs.image_extract"]
|
|
pixels = np.zeros((4, 5, 4), dtype=np.uint8)
|
|
pixels[1:3, 2:4] = (20, 40, 60, 255)
|
|
result, bounds = image_extract.crop_to_alpha(
|
|
Image.fromarray(pixels, mode="RGBA"),
|
|
(10, 20, 15, 24),
|
|
padding=1,
|
|
)
|
|
result = np.asarray(result)
|
|
|
|
self.assertEqual(bounds, (11, 20, 15, 24))
|
|
self.assertEqual(result.shape, (4, 4, 4))
|
|
self.assertTrue(np.all(result[[0, -1], :, 3] == 0))
|
|
self.assertTrue(np.all(result[:, [0, -1], 3] == 0))
|
|
self.assertTrue(np.all(result[1:3, 1:3, 3] == 255))
|
|
|
|
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"]
|
|
y, x = np.indices((7, 9), dtype=np.float32)
|
|
depth = 2.0 + x * 0.1 + y * 0.2
|
|
depth[3, 4] -= 0.5
|
|
depth[1, 1] += 3.0
|
|
model_mask = np.ones_like(depth, dtype=bool)
|
|
selection_mask = np.ones_like(depth, dtype=bool)
|
|
selection_mask[0, 0] = False
|
|
|
|
depth_signal, valid, reference, baseline_vector = balloon.normalized_moge_depth(
|
|
depth,
|
|
model_mask,
|
|
selection_mask,
|
|
)
|
|
normalized_x = x / (depth.shape[1] - 1) * 2.0 - 1.0
|
|
normalized_y = y / (depth.shape[0] - 1) * 2.0 - 1.0
|
|
fitted_plane = (
|
|
baseline_vector[0] * normalized_x
|
|
+ baseline_vector[1] * normalized_y
|
|
+ baseline_vector[2]
|
|
)
|
|
displacement = depth_signal - fitted_plane
|
|
|
|
self.assertFalse(valid[0, 0])
|
|
self.assertEqual(depth_signal[0, 0], 0.0)
|
|
self.assertGreater(float(displacement[3, 4]), 1.0)
|
|
self.assertLess(float(displacement[3, 4]), 1.01)
|
|
self.assertLess(displacement[1, 1], -1.0)
|
|
self.assertGreater(reference, 2.0)
|
|
self.assertEqual(baseline_vector.shape, (3,))
|
|
self.assertLess(float(displacement.min()), -1.0)
|
|
self.assertLess(float(displacement.max()), 1.01)
|
|
self.assertLess(float(np.percentile(displacement, 90.0)), 0.01)
|
|
|
|
positive = displacement[valid & (displacement > 0.0)]
|
|
self.assertAlmostEqual(
|
|
float(np.sum(np.square(positive)) / np.sum(positive)),
|
|
1.0,
|
|
places=5,
|
|
)
|
|
|
|
def test_balloon_relative_depth_ignores_absolute_distance_and_scale(self):
|
|
import numpy as np
|
|
|
|
balloon = sys.modules["server.jobs.balloon"]
|
|
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(
|
|
(
|
|
(0.1, 0.2, 0.3, 0.1),
|
|
(0.2, 0.4, 0.8, 0.2),
|
|
(0.15, 0.3, 0.5, 0.15),
|
|
(0.05, 0.1, 0.2, 0.05),
|
|
),
|
|
dtype=np.float32,
|
|
)
|
|
mask = np.ones_like(depth, dtype=bool)
|
|
|
|
original = balloon.normalized_moge_depth(depth, mask, mask)
|
|
farther = balloon.normalized_moge_depth(depth + 100.0, mask, mask)
|
|
rescaled = balloon.normalized_moge_depth(depth * 7.0, mask, mask)
|
|
|
|
np.testing.assert_allclose(farther[0], original[0], atol=2e-4)
|
|
np.testing.assert_allclose(farther[3], original[3], atol=2e-4)
|
|
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()
|