Add new tests and functionality for image processing and geometry handling
- Introduced color_image.py for preparing standalone color images for tests. - Added hdr_image.py to create HDR textures for testing. - Implemented test_image_depth_panorama.py to evaluate panorama topology, UVs, and depth discontinuities. - Created test_shadeless.py to validate the opacity and HDR color endpoints of the Shadeless group. - Developed moge3_export.py for exporting dynamic MoGe-3 backbone and ONNX graphs. - Added material.py to inherit material slots onto generated geometry. - Built image_depth_panorama.py to construct a quad panorama with depth separation and sampling.
This commit is contained in:
@@ -54,10 +54,9 @@ class ServerJobsTest(unittest.TestCase):
|
||||
"input_frames",
|
||||
return_value=nullcontext([input_path]),
|
||||
),
|
||||
patch.object(self.depth_plane_job, "copy_color_images"),
|
||||
patch.object(
|
||||
self.depth_plane_job,
|
||||
"generate_moge2_artifacts",
|
||||
"generate_moge_artifacts",
|
||||
return_value={
|
||||
"depth": depth_texture,
|
||||
"depth_metadata": metadata,
|
||||
@@ -85,100 +84,23 @@ class ServerJobsTest(unittest.TestCase):
|
||||
|
||||
|
||||
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
||||
inference_options = []
|
||||
save_options = []
|
||||
|
||||
class FakeImage:
|
||||
size = (1, 1)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_args):
|
||||
return None
|
||||
|
||||
def convert(self, mode):
|
||||
self.mode = mode
|
||||
return self
|
||||
|
||||
def getchannel(self, _name):
|
||||
return self
|
||||
|
||||
def putalpha(self, _alpha):
|
||||
return None
|
||||
|
||||
def paste(self, *_args):
|
||||
return None
|
||||
|
||||
def save(self, path, **options):
|
||||
save_options.append(options)
|
||||
Path(path).touch()
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
self.assertEqual(model_dir, "model")
|
||||
return cls()
|
||||
|
||||
def to(model_self, device):
|
||||
self.assertEqual(device.type, "cpu")
|
||||
return model_self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
def inference(self, image, **options):
|
||||
inference_options.append((image, options))
|
||||
return FakeImage()
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.cuda = SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
synchronize=lambda: None,
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("server.models.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image, **options: (
|
||||
inference_options.append((image, options)) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(
|
||||
open=lambda _path: FakeImage(),
|
||||
new=lambda *_args, **_kwargs: FakeImage(),
|
||||
)
|
||||
fake_pil.ImageChops = SimpleNamespace(
|
||||
multiply=lambda _source, _foreground: FakeImage()
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"server.models.onnx_ben2": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
from PIL import Image
|
||||
onnx_ben2 = load_server_module("models.onnx_ben2")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
input_path = root / "input.png"
|
||||
input_path.touch()
|
||||
Image.new("RGBA", (2, 1), (255, 64, 32, 128)).save(input_path)
|
||||
context = FakeJobContext(root / "result", "remove-background")
|
||||
self.background.run(
|
||||
context,
|
||||
{"input": str(input_path), "model": "BEN2", "device": "cpu"},
|
||||
)
|
||||
|
||||
self.assertTrue((root / "result" / "image.0001.png").is_file())
|
||||
|
||||
self.assertEqual(
|
||||
inference_options[0][1],
|
||||
{"release_memory": True},
|
||||
)
|
||||
self.assertEqual(
|
||||
save_options[0],
|
||||
{"format": "PNG", "compress_level": 1},
|
||||
)
|
||||
with (
|
||||
patch.object(onnx_ben2, "create_session", return_value=object()),
|
||||
patch.object(onnx_ben2, "infer_alpha", return_value=np.full((1, 2), 128 / 255, np.float32)) as infer,
|
||||
patch.object(Image.Image, "save", autospec=True, side_effect=Image.Image.save) as save,
|
||||
):
|
||||
self.background.run(context, {"input": str(input_path), "model": "BEN2", "device": "cpu"})
|
||||
with Image.open(root / "result" / "image.0001.png") as image:
|
||||
self.assertEqual(image.getpixel((0, 0)), (255, 64, 32, 64))
|
||||
self.assertEqual(infer.call_args.kwargs, {"release_memory": True})
|
||||
self.assertEqual(save.call_args.kwargs, {"format": "PNG", "compress_level": 1})
|
||||
self.assertIn("Background removed in", context.status["message"])
|
||||
|
||||
|
||||
@@ -190,7 +112,7 @@ class ServerJobsTest(unittest.TestCase):
|
||||
(((255, 0, 0, 0), (0, 255, 0, 128)),),
|
||||
dtype=np.uint8,
|
||||
)
|
||||
model_result = Image.new("RGBA", (2, 1), (10, 20, 30, 128))
|
||||
model_result = np.full((1, 2), 128 / 255, np.float32)
|
||||
received = []
|
||||
|
||||
def infer(_model, image, **_options):
|
||||
@@ -202,7 +124,7 @@ class ServerJobsTest(unittest.TestCase):
|
||||
Image.fromarray(source_pixels, mode="RGBA").save(input_path)
|
||||
with (
|
||||
patch.object(onnx_ben2, "create_session", return_value=object()),
|
||||
patch.object(onnx_ben2, "infer", side_effect=infer),
|
||||
patch.object(onnx_ben2, "infer_alpha", side_effect=infer),
|
||||
):
|
||||
result, _load_ms, _inference_ms = self.ben2.infer_foreground(
|
||||
input_path,
|
||||
@@ -226,111 +148,48 @@ class ServerJobsTest(unittest.TestCase):
|
||||
self.assertEqual(tuple(result_pixels[0, 1]), (0, 255, 0, 64))
|
||||
self.assertEqual(
|
||||
tuple(np.asarray(raw_result.convert("RGBA"))[0, 0]),
|
||||
(10, 20, 30, 128),
|
||||
(0, 0, 0, 128),
|
||||
)
|
||||
|
||||
|
||||
def test_ben2_background_removal_processes_each_frame(self):
|
||||
inference_options = []
|
||||
|
||||
class FakeImage:
|
||||
size = (1, 1)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_args):
|
||||
return None
|
||||
|
||||
def convert(self, mode):
|
||||
self.mode = mode
|
||||
return self
|
||||
|
||||
def getchannel(self, _name):
|
||||
return self
|
||||
|
||||
def putalpha(self, _alpha):
|
||||
return None
|
||||
|
||||
def paste(self, *_args):
|
||||
return None
|
||||
|
||||
def save(self, path, **options):
|
||||
Path(path).touch()
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
return cls()
|
||||
|
||||
def to(model_self, device):
|
||||
return model_self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
def inference(self, image, **options):
|
||||
inference_options.append((image, options))
|
||||
return FakeImage()
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.cuda = SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
synchronize=lambda: None,
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("server.models.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image, **options: (
|
||||
inference_options.append((image, options)) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(
|
||||
open=lambda _path: FakeImage(),
|
||||
new=lambda *_args, **_kwargs: FakeImage(),
|
||||
)
|
||||
fake_pil.ImageChops = SimpleNamespace(
|
||||
multiply=lambda _source, _foreground: FakeImage()
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"server.models.onnx_ben2": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
from PIL import Image
|
||||
onnx_ben2 = load_server_module("models.onnx_ben2")
|
||||
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()
|
||||
frames = root / "frames"
|
||||
frames.mkdir()
|
||||
for index in (1, 2):
|
||||
Image.new("RGB", (2, 1), (index * 50, 0, 0)).save(frames / f"frame.{index:04d}.png")
|
||||
context = FakeJobContext(root / "result", "remove-background")
|
||||
self.background.run(
|
||||
context,
|
||||
{
|
||||
"input": str(input_directory),
|
||||
"model": "BEN2",
|
||||
"device": "cpu",
|
||||
"video_frames": 8,
|
||||
},
|
||||
)
|
||||
with (
|
||||
patch.object(onnx_ben2, "create_session", return_value=object()),
|
||||
patch.object(onnx_ben2, "infer_alpha", return_value=np.ones((1, 2), np.float32)) as infer,
|
||||
):
|
||||
self.background.run(context, {"input": str(frames), "model": "BEN2", "device": "cpu", "video_frames": 8})
|
||||
self.assertEqual([call.kwargs["release_memory"] for call in infer.call_args_list], [False, True])
|
||||
for index in (1, 2):
|
||||
with Image.open(context.directory / f"image.{index:04d}.png") as image:
|
||||
self.assertEqual(image.getpixel((0, 0)), (index * 50, 0, 0, 255))
|
||||
|
||||
result_directory = context.directory
|
||||
self.assertTrue(
|
||||
(result_directory / "image.0001.png").is_file()
|
||||
)
|
||||
self.assertTrue(
|
||||
(result_directory / "image.0002.png").is_file()
|
||||
)
|
||||
|
||||
self.assertEqual(len(inference_options), 2)
|
||||
self.assertEqual(
|
||||
[options["release_memory"] for _image, options in inference_options],
|
||||
[False, True],
|
||||
)
|
||||
def test_background_alpha_job_preserves_float_precision(self):
|
||||
from PIL import Image
|
||||
onnx_ben2 = load_server_module("models.onnx_ben2")
|
||||
alpha = np.array([[0.123456, 0.987654]], dtype=np.float32)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
path = root / "input.png"
|
||||
Image.new("RGB", (2, 1)).save(path)
|
||||
context = FakeJobContext(root / "result", "remove-background")
|
||||
with (
|
||||
patch.object(onnx_ben2, "create_session", return_value=object()),
|
||||
patch.object(onnx_ben2, "infer_alpha", return_value=alpha),
|
||||
):
|
||||
result = self.background.run(context, {"input": str(path), "model": "BEN2", "device": "cpu", "output_kind": "alpha"})
|
||||
self.assertEqual(result, {"alpha": "alpha.npy"})
|
||||
np.testing.assert_array_equal(np.load(context.directory / result["alpha"], allow_pickle=False), alpha)
|
||||
self.assertFalse(list(context.directory.glob("*.png")))
|
||||
|
||||
|
||||
def test_upscale_job_saves_importable_upscale_sequence(self):
|
||||
|
||||
Reference in New Issue
Block a user