Add new node asset scripts for geometry and shader operations
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters. - Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections. - Added `validate.py` to validate the integrity of the newly created node groups and their inputs. - Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies. - Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
@@ -41,8 +41,8 @@ class ServerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.server_module = load_server_module()
|
||||
importlib.import_module("server.jobs.cutout_depth")
|
||||
importlib.import_module("server.jobs.mesh")
|
||||
importlib.import_module("server.selection.polygon_mask")
|
||||
importlib.import_module("server.jobs.prepare_geometry_data")
|
||||
from blendjob import JobServer
|
||||
from server.model_manager import ModelManager, ModelSessionCache
|
||||
|
||||
@@ -82,7 +82,6 @@ class ServerTest(unittest.TestCase):
|
||||
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.ModelManager(Path(directory))
|
||||
snapshot = manager.snapshot()
|
||||
@@ -262,8 +261,7 @@ class ServerTest(unittest.TestCase):
|
||||
"refine-image-mask",
|
||||
"remove-background",
|
||||
"upscale-image",
|
||||
"extract-frames",
|
||||
"generate-mesh",
|
||||
"prepare-geometry-data",
|
||||
"debug-da3",
|
||||
"debug-moge2",
|
||||
"debug-ben2",
|
||||
@@ -387,36 +385,8 @@ class ServerTest(unittest.TestCase):
|
||||
)
|
||||
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.app_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 / "image.0001.jpg").is_file())
|
||||
|
||||
self.assertEqual(status["state"], "succeeded")
|
||||
|
||||
def test_polygon_selection_mask_uses_even_odd_fill_for_self_intersection(self):
|
||||
image_mask = sys.modules["server.jobs.image_mask"]
|
||||
image_mask = sys.modules["server.selection.polygon_mask"]
|
||||
polygon = (
|
||||
(1, 1),
|
||||
(9, 1),
|
||||
@@ -438,7 +408,7 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertEqual(int(mask.sum()), 48)
|
||||
|
||||
def test_polygon_selection_mask_nonzero_keeps_same_direction_overlap(self):
|
||||
image_mask = sys.modules["server.jobs.image_mask"]
|
||||
image_mask = sys.modules["server.selection.polygon_mask"]
|
||||
polygon = (
|
||||
(1, 1),
|
||||
(7, 1),
|
||||
@@ -467,7 +437,7 @@ class ServerTest(unittest.TestCase):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
image_mask = sys.modules["server.jobs.image_mask"]
|
||||
image_mask = sys.modules["server.selection.polygon_mask"]
|
||||
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)
|
||||
@@ -487,7 +457,7 @@ class ServerTest(unittest.TestCase):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
image_mask = sys.modules["server.jobs.image_mask"]
|
||||
image_mask = sys.modules["server.selection.polygon_mask"]
|
||||
pixels = np.zeros((4, 5, 4), dtype=np.uint8)
|
||||
pixels[1:3, 2:4] = (20, 40, 60, 255)
|
||||
result, bounds = image_mask.crop_to_alpha(
|
||||
@@ -501,10 +471,11 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertTrue(np.all(result[:, :, 3] == 255))
|
||||
|
||||
def test_mesh_refine_selection_keeps_source_size_and_selection_alpha(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "source.png"
|
||||
pixels = np.full((6, 8, 4), (20, 40, 60, 255), dtype=np.uint8)
|
||||
Image.fromarray(pixels, mode="RGBA").save(source_path)
|
||||
@@ -518,12 +489,16 @@ class ServerTest(unittest.TestCase):
|
||||
check_cancelled=lambda: None,
|
||||
)
|
||||
|
||||
def passthrough(path, *_args):
|
||||
def passthrough(path, *_args, **_kwargs):
|
||||
return Image.open(path).convert("RGBA"), 0.0, 0.0
|
||||
|
||||
output_path = job_directory / "image.0001.png"
|
||||
with patch.object(mesh_job, "remove_image_background", side_effect=passthrough):
|
||||
mesh_job._write_refined_color(
|
||||
with patch.object(
|
||||
cutout_job,
|
||||
"infer_foreground",
|
||||
side_effect=passthrough,
|
||||
):
|
||||
cutout_job.refine_selection(
|
||||
context,
|
||||
{"refine_model": "BEN2_BASE", "device": "cpu"},
|
||||
source_path,
|
||||
@@ -537,17 +512,20 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertEqual(tuple(result[5, 7]), (0, 0, 0, 0))
|
||||
|
||||
def test_mesh_cutout_images_crop_color_and_preserve_source_alpha(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "analysis-source.png"
|
||||
selection_path = self.storage_root / "selection-source.png"
|
||||
output_directory = self.storage_root / "cutout-images"
|
||||
output_directory.mkdir()
|
||||
source = np.full((6, 8, 4), (10, 80, 160, 255), dtype=np.uint8)
|
||||
source[2, 4, 3] = 64
|
||||
source[2, 3, 3] = 25
|
||||
source[2, 4, 3] = 26
|
||||
source[3, 5, 3] = 0
|
||||
Image.fromarray(source, mode="RGBA").save(source_path)
|
||||
selection = np.zeros((6, 8, 4), dtype=np.uint8)
|
||||
selection[1:5, 2:7] = (200, 20, 40, 255)
|
||||
@@ -555,7 +533,7 @@ class ServerTest(unittest.TestCase):
|
||||
Image.fromarray(selection, mode="RGBA").save(selection_path)
|
||||
|
||||
color_path, mask_path, analysis_path, bounds_path = (
|
||||
mesh_job._write_cutout_images(
|
||||
cutout_job.write_images(
|
||||
source_path,
|
||||
selection_path,
|
||||
output_directory,
|
||||
@@ -567,18 +545,26 @@ class ServerTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(json.loads(bounds_path.read_text()), [2, 1, 7, 5])
|
||||
self.assertEqual(color.shape, (4, 5, 4))
|
||||
self.assertEqual(tuple(color[1, 2]), (10, 80, 160, 64))
|
||||
self.assertEqual(tuple(color[1, 1]), (10, 80, 160, 25))
|
||||
self.assertEqual(int(mask[1, 1, 3]), 0)
|
||||
self.assertEqual(tuple(analysis[1, 1]), (0, 0, 0, 0))
|
||||
self.assertEqual(tuple(color[1, 2]), (10, 80, 160, 26))
|
||||
self.assertGreater(int(mask[1, 2, 3]), 25)
|
||||
self.assertEqual(int(mask[0, 0, 3]), 0)
|
||||
self.assertEqual(tuple(color[0, 0]), (10, 80, 160, 255))
|
||||
self.assertEqual(tuple(analysis[0, 0]), (0, 0, 0, 0))
|
||||
self.assertEqual(tuple(analysis[1, 2]), (10, 80, 160, 64))
|
||||
self.assertEqual(tuple(analysis[1, 2]), (10, 80, 160, 26))
|
||||
self.assertEqual(tuple(color[2, 3]), (10, 80, 160, 0))
|
||||
self.assertEqual(int(mask[2, 3, 3]), 0)
|
||||
self.assertEqual(tuple(analysis[2, 3]), (0, 0, 0, 0))
|
||||
|
||||
def test_mesh_unrefined_cutout_keeps_selection_outside_the_source(self):
|
||||
def test_mesh_unrefined_cutout_discards_selection_outside_the_source(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules['server.jobs.mesh']
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / 'outside-source.png'
|
||||
full_selection_path = self.storage_root / 'outside-selection-full.png'
|
||||
output_directory = self.storage_root / 'outside-cutout'
|
||||
@@ -587,14 +573,14 @@ class ServerTest(unittest.TestCase):
|
||||
polygon = ((-2, 1), (10, 1), (10, 5), (-2, 5))
|
||||
|
||||
_selection_path, selection_bounds = (
|
||||
mesh_job._write_unbounded_selection(
|
||||
cutout_job.write_selection(
|
||||
source_path,
|
||||
polygon,
|
||||
full_selection_path,
|
||||
)
|
||||
)
|
||||
color_path, mask_path, analysis_path, bounds_path = (
|
||||
mesh_job._write_cutout_images(
|
||||
cutout_job.write_images(
|
||||
source_path,
|
||||
full_selection_path,
|
||||
output_directory,
|
||||
@@ -606,19 +592,19 @@ class ServerTest(unittest.TestCase):
|
||||
analysis = np.asarray(Image.open(analysis_path).convert('RGBA'))
|
||||
|
||||
self.assertEqual(selection_bounds, (-2, 1, 10, 5))
|
||||
self.assertEqual(json.loads(bounds_path.read_text()), [-2, 1, 10, 5])
|
||||
self.assertEqual(color.shape, (4, 12, 4))
|
||||
self.assertEqual(tuple(color[1, 0]), (0, 0, 0, 0))
|
||||
self.assertEqual(tuple(color[1, 2]), (10, 80, 160, 255))
|
||||
self.assertEqual(json.loads(bounds_path.read_text()), [0, 1, 8, 5])
|
||||
self.assertEqual(color.shape, (4, 8, 4))
|
||||
self.assertEqual(tuple(color[1, 0]), (10, 80, 160, 255))
|
||||
self.assertEqual(tuple(color[1, 7]), (10, 80, 160, 255))
|
||||
self.assertGreater(int(mask[1, 0, 3]), 127)
|
||||
self.assertEqual(tuple(analysis[1, 0]), (0, 0, 0, 0))
|
||||
self.assertEqual(tuple(analysis[1, 0]), (10, 80, 160, 255))
|
||||
|
||||
def test_mesh_job_keeps_unrefined_outside_selection_with_normal(self):
|
||||
def test_mesh_job_discards_unrefined_outside_selection_with_normal(self):
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "mesh-source.png"
|
||||
pixels = np.full((6, 8, 4), (20, 40, 60, 255), dtype=np.uint8)
|
||||
Image.fromarray(pixels, mode="RGBA").save(source_path)
|
||||
@@ -650,8 +636,8 @@ class ServerTest(unittest.TestCase):
|
||||
"mesh_detail": 64,
|
||||
"invert": False,
|
||||
"device": "cpu",
|
||||
"normal_model": "MOGE2_VITS_NORMAL",
|
||||
"normal_resolution_level": 5,
|
||||
"model": "MOGE2_VITS_NORMAL",
|
||||
"resolution_level": 5,
|
||||
"normal_analysis": True,
|
||||
},
|
||||
)
|
||||
@@ -678,24 +664,74 @@ class ServerTest(unittest.TestCase):
|
||||
)
|
||||
self.assertEqual(
|
||||
json.loads((job_directory / "bounds.json").read_text()),
|
||||
[-2, 1, 10, 5],
|
||||
[0, 1, 8, 5],
|
||||
)
|
||||
self.assertEqual(tuple(analysis_pixels[2, 2, :3]), (20, 40, 60))
|
||||
self.assertGreater(int(analysis_pixels[2, 2, 3]), 240)
|
||||
self.assertEqual(tuple(analysis_pixels[2, 0]), (0, 0, 0, 0))
|
||||
self.assertEqual(tuple(color_pixels[2, 2, :3]), (20, 40, 60))
|
||||
self.assertGreater(int(color_pixels[2, 2, 3]), 240)
|
||||
self.assertEqual(tuple(color_pixels[2, 0]), (0, 0, 0, 0))
|
||||
self.assertEqual(color_pixels.shape[:2], (4, 12))
|
||||
self.assertEqual(analysis_pixels.shape[:2], (4, 12))
|
||||
self.assertEqual(normal_pixels.shape[:2], (4, 12))
|
||||
self.assertEqual(tuple(normal_pixels[0, 0]), (128, 128, 255))
|
||||
self.assertEqual(tuple(normal_pixels[2, 2]), (255, 127, 127))
|
||||
self.assertEqual(tuple(analysis_pixels[2, 0, :3]), (20, 40, 60))
|
||||
self.assertEqual(tuple(color_pixels[2, 0, :3]), (20, 40, 60))
|
||||
self.assertGreater(int(color_pixels[2, 0, 3]), 240)
|
||||
self.assertEqual(tuple(color_pixels[2, 7, :3]), (20, 40, 60))
|
||||
self.assertEqual(color_pixels.shape[:2], (4, 8))
|
||||
self.assertEqual(analysis_pixels.shape[:2], (4, 8))
|
||||
self.assertEqual(normal_pixels.shape[:2], (4, 8))
|
||||
self.assertEqual(tuple(normal_pixels[0, 0]), (255, 127, 127))
|
||||
self.assertEqual(tuple(normal_pixels[2, 0]), (255, 127, 127))
|
||||
|
||||
def test_cutout_limits_the_cropped_model_input_but_keeps_color_resolution(self):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "limited-cutout-source.png"
|
||||
Image.new("RGBA", (20, 10), (20, 40, 60, 255)).save(source_path)
|
||||
job_directory = self.storage_root / "limited-cutout-job"
|
||||
context = SimpleNamespace(
|
||||
directory=job_directory,
|
||||
progress=lambda *_args: None,
|
||||
check_cancelled=lambda: None,
|
||||
resource=lambda _name: object(),
|
||||
)
|
||||
prediction = SimpleNamespace(
|
||||
normal=np.zeros((1, 2, 8, 3), dtype=np.float32)
|
||||
)
|
||||
inferred_sizes = []
|
||||
|
||||
def infer(_manager, _parameters, images, **_kwargs):
|
||||
with Image.open(images[0]) as model_input:
|
||||
inferred_sizes.append(model_input.size)
|
||||
return prediction
|
||||
|
||||
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
|
||||
result = mesh_job.run(
|
||||
context,
|
||||
{
|
||||
"input": str(source_path),
|
||||
"shape": "BALLOON",
|
||||
"polygon": ((2, 3), (18, 3), (18, 7), (2, 7)),
|
||||
"mesh_detail": 64,
|
||||
"max_input_size": 8,
|
||||
"invert": False,
|
||||
"device": "cpu",
|
||||
"model": "MOGE2_VITS_NORMAL",
|
||||
"resolution_level": 5,
|
||||
"normal_analysis": True,
|
||||
},
|
||||
)
|
||||
|
||||
with Image.open(job_directory / "image.0001.png") as color:
|
||||
color_size = color.size
|
||||
with Image.open(job_directory / result["normal"]) as normal:
|
||||
normal_size = normal.size
|
||||
self.assertEqual(inferred_sizes, [(8, 2)])
|
||||
self.assertEqual(color_size, (16, 4))
|
||||
self.assertEqual(normal_size, (16, 4))
|
||||
self.assertFalse((job_directory / "model-input.png").exists())
|
||||
|
||||
def test_mesh_job_skips_moge_normal_when_option_is_disabled(self):
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "mesh-no-normal-source.png"
|
||||
Image.new("RGBA", (8, 6), (20, 40, 60, 255)).save(source_path)
|
||||
job_directory = self.storage_root / "mesh-no-normal-job"
|
||||
@@ -727,57 +763,165 @@ class ServerTest(unittest.TestCase):
|
||||
write_normal.assert_not_called()
|
||||
self.assertFalse((job_directory / "normal.png").exists())
|
||||
|
||||
def test_mesh_depth_balloon_uses_cropped_analysis_alpha(self):
|
||||
import numpy as np
|
||||
def test_mesh_depth_balloon_uses_metric_point_field(self):
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "mesh-depth-balloon-source.png"
|
||||
pixels = np.full((6, 6, 4), (20, 40, 60, 0), dtype=np.uint8)
|
||||
pixels[2:5, 2:5, 3] = 255
|
||||
Image.fromarray(pixels, mode="RGBA").save(source_path)
|
||||
Image.new("RGBA", (6, 6), (20, 40, 60, 255)).save(source_path)
|
||||
job_directory = self.storage_root / "mesh-depth-balloon-job"
|
||||
job_directory.mkdir()
|
||||
context = SimpleNamespace(
|
||||
directory=job_directory,
|
||||
progress=lambda *_args: None,
|
||||
check_cancelled=lambda: None,
|
||||
resource=lambda _name: object(),
|
||||
)
|
||||
received = []
|
||||
field_path = job_directory / "field.npz"
|
||||
|
||||
def infer(_manager, _parameters, images, **_kwargs):
|
||||
received.extend(images)
|
||||
depth = np.linspace(1.0, 2.0, 36, dtype=np.float32).reshape(6, 6)
|
||||
normal = np.zeros((1, 6, 6, 3), dtype=np.float32)
|
||||
normal[..., 0] = 1.0
|
||||
return SimpleNamespace(
|
||||
depth=depth[None],
|
||||
mask=np.ones((1, 6, 6), dtype=bool),
|
||||
normal=normal,
|
||||
with patch.object(
|
||||
mesh_job,
|
||||
"_write_depth_field",
|
||||
return_value=(field_path, None),
|
||||
) as write_surface:
|
||||
result = mesh_job.run(
|
||||
context,
|
||||
{
|
||||
"input": str(source_path),
|
||||
"shape": "DEPTH_BALLOON",
|
||||
"polygon": ((1, 1), (5, 1), (5, 5), (1, 5)),
|
||||
"normal_analysis": False,
|
||||
"video_frames": 1,
|
||||
"invert": False,
|
||||
},
|
||||
)
|
||||
|
||||
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
|
||||
depth_path, normal_path = mesh_job._write_depth_balloon(
|
||||
self.assertEqual(
|
||||
result,
|
||||
{
|
||||
"selection": "selection.png",
|
||||
"bounds": "bounds.json",
|
||||
"field": "field.npz",
|
||||
},
|
||||
)
|
||||
self.assertTrue(write_surface.called)
|
||||
self.assertEqual(
|
||||
write_surface.call_args.kwargs["selection_path"],
|
||||
job_directory / "selection.png",
|
||||
)
|
||||
|
||||
def test_moge_depth_and_normal_share_one_prediction(self):
|
||||
import numpy as np
|
||||
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
job_directory = self.storage_root / "shared-moge-prediction"
|
||||
job_directory.mkdir()
|
||||
color_path = job_directory / "analysis.png"
|
||||
normal_path = job_directory / "normal.png"
|
||||
prediction = object()
|
||||
valid = np.ones((2, 3), dtype=bool)
|
||||
frame = SimpleNamespace(
|
||||
depth=np.ones((2, 3), dtype=np.float32),
|
||||
points=np.ones((2, 3, 3), dtype=np.float32),
|
||||
intrinsics=np.eye(3, dtype=np.float32),
|
||||
valid_mask=valid,
|
||||
model_mask=valid,
|
||||
)
|
||||
context = SimpleNamespace(
|
||||
directory=job_directory,
|
||||
resource=lambda _name: object(),
|
||||
check_cancelled=lambda: None,
|
||||
)
|
||||
parameters = {
|
||||
"model": "MOGE2_VITS_NORMAL",
|
||||
"device": "cpu",
|
||||
"resolution_level": 5,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
mesh_job.moge2, "infer", return_value=prediction
|
||||
) as infer,
|
||||
patch.object(
|
||||
mesh_job.moge2, "geometry_frame", return_value=frame
|
||||
),
|
||||
patch.object(mesh_job, "_write_field") as write_field,
|
||||
patch.object(
|
||||
mesh_job,
|
||||
"_write_camera_normal_from_prediction",
|
||||
return_value=normal_path,
|
||||
) as write_normal,
|
||||
):
|
||||
field_path, written_normal_path = mesh_job._write_depth_field(
|
||||
context,
|
||||
{"device": "cpu"},
|
||||
source_path,
|
||||
parameters,
|
||||
color_path,
|
||||
include_normal=True,
|
||||
)
|
||||
|
||||
self.assertEqual(received, [source_path])
|
||||
with np.load(depth_path, allow_pickle=False) as data:
|
||||
self.assertEqual(tuple(data["origin"]), (0, 0))
|
||||
self.assertEqual(tuple(data["crop_size"]), (6, 6))
|
||||
self.assertFalse(bool(data["valid"][0, 0]))
|
||||
self.assertTrue(bool(data["valid"][2, 2]))
|
||||
normal = np.asarray(Image.open(normal_path).convert("RGB"))
|
||||
self.assertEqual(tuple(normal[0, 0]), (128, 128, 255))
|
||||
self.assertEqual(tuple(normal[2, 2]), (255, 127, 127))
|
||||
infer.assert_called_once()
|
||||
write_field.assert_called_once()
|
||||
write_normal.assert_called_once_with(
|
||||
prediction, color_path, normal_path
|
||||
)
|
||||
self.assertEqual(field_path, job_directory / "field.npz")
|
||||
self.assertEqual(written_normal_path, normal_path)
|
||||
|
||||
def test_depth_calibration_mask_is_separate_from_model_validity(self):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
selection_path = self.storage_root / "depth-selection.png"
|
||||
selection = np.zeros((2, 3, 4), dtype=np.uint8)
|
||||
selection[0, 1, 3] = 255
|
||||
selection[1, 2, 3] = 255
|
||||
Image.fromarray(selection, mode="RGBA").save(selection_path)
|
||||
frame = SimpleNamespace(
|
||||
valid_mask=np.asarray(
|
||||
((True, True, True), (True, True, False)),
|
||||
dtype=bool,
|
||||
)
|
||||
)
|
||||
|
||||
original_valid = frame.valid_mask.copy()
|
||||
result = mesh_job._selection_calibration_mask(frame, selection_path)
|
||||
|
||||
np.testing.assert_array_equal(
|
||||
result,
|
||||
((False, True, False), (False, False, False)),
|
||||
)
|
||||
np.testing.assert_array_equal(frame.valid_mask, original_valid)
|
||||
|
||||
def test_depth_field_writes_original_and_calibration_masks_separately(self):
|
||||
import numpy as np
|
||||
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
output_path = self.storage_root / "field.npz"
|
||||
valid = np.asarray(((True, True), (True, False)), dtype=bool)
|
||||
calibration = np.asarray(((False, True), (False, False)), dtype=bool)
|
||||
depth = np.asarray(((2.0, 2.1), (2.2, 0.0)), dtype=np.float32)
|
||||
points = np.zeros((2, 2, 3), dtype=np.float32)
|
||||
points[..., 2] = depth
|
||||
frame = SimpleNamespace(
|
||||
depth=depth,
|
||||
points=points,
|
||||
intrinsics=np.eye(3, dtype=np.float32),
|
||||
valid_mask=valid,
|
||||
model_mask=valid,
|
||||
)
|
||||
|
||||
mesh_job._write_field(frame, output_path, calibration)
|
||||
|
||||
with np.load(output_path, allow_pickle=False) as archive:
|
||||
np.testing.assert_array_equal(archive["valid"], valid)
|
||||
np.testing.assert_array_equal(
|
||||
archive["calibration_valid"],
|
||||
calibration,
|
||||
)
|
||||
|
||||
def test_mesh_job_uses_upscale_source_only_for_analysis(self):
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
display_path = self.storage_root / "upscaled.png"
|
||||
analysis_path = self.storage_root / "original.png"
|
||||
Image.new("RGBA", (12, 8), (200, 20, 40, 255)).save(display_path)
|
||||
@@ -792,7 +936,7 @@ class ServerTest(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
mesh_job,
|
||||
"_write_surface_depth",
|
||||
"_write_depth_field",
|
||||
return_value=(field_path, None),
|
||||
) as write_surface:
|
||||
result = mesh_job.run(
|
||||
@@ -812,10 +956,11 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertEqual(color.size, (12, 8))
|
||||
|
||||
def test_mesh_masked_color_keeps_same_direction_overlap(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "opaque-overlap.png"
|
||||
output_path = self.storage_root / "cutout-overlap.png"
|
||||
Image.new("RGBA", (12, 9), (20, 40, 60, 255)).save(source_path)
|
||||
@@ -833,17 +978,18 @@ class ServerTest(unittest.TestCase):
|
||||
(1, 1),
|
||||
)
|
||||
|
||||
mesh_job._write_masked_color(source_path, polygon, output_path)
|
||||
cutout_job._write_masked_color(source_path, polygon, output_path)
|
||||
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
|
||||
|
||||
self.assertGreater(int(alpha[3, 2]), 200)
|
||||
self.assertGreater(int(alpha[3, 5]), 200)
|
||||
self.assertGreater(int(alpha[3, 8]), 200)
|
||||
def test_mesh_masked_color_preserves_nonzero_winding_hole(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "opaque-hole.png"
|
||||
output_path = self.storage_root / "cutout-hole.png"
|
||||
Image.new("RGBA", (10, 10), (20, 40, 60, 255)).save(source_path)
|
||||
@@ -861,131 +1007,30 @@ class ServerTest(unittest.TestCase):
|
||||
(1, 1),
|
||||
)
|
||||
|
||||
mesh_job._write_masked_color(source_path, polygon, output_path)
|
||||
cutout_job._write_masked_color(source_path, polygon, output_path)
|
||||
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
|
||||
|
||||
self.assertGreater(int(alpha[2, 2]), 200)
|
||||
self.assertLess(int(alpha[4, 4]), 32)
|
||||
|
||||
def test_mesh_masked_color_antialiases_the_cutout_edge(self):
|
||||
cutout_job = importlib.import_module("server.selection.cutout_images")
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
mesh_job = sys.modules["server.jobs.mesh"]
|
||||
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
|
||||
source_path = self.storage_root / "opaque.png"
|
||||
output_path = self.storage_root / "cutout.png"
|
||||
Image.new("RGBA", (16, 16), (20, 40, 60, 255)).save(source_path)
|
||||
|
||||
mesh_job._write_masked_color(
|
||||
cutout_job._write_masked_color(
|
||||
source_path,
|
||||
((1, 1), (14, 1), (1, 14)),
|
||||
output_path,
|
||||
)
|
||||
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
|
||||
|
||||
def test_balloon_depth_analysis_normalizes_moge_distance(self):
|
||||
import numpy as np
|
||||
|
||||
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
|
||||
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_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.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(
|
||||
(
|
||||
(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)
|
||||
self.assertTrue(np.any((alpha > 0) & (alpha < 255)))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user