Add new functionality for image rectification and depth metadata handling
- Implemented a preview tool for rectifying images with interactive aspect controls in `preview.py`. - Added tests for depth metadata loading to ensure consistent schema usage in `test_depth_metadata.py`. - Developed comprehensive tests for cutout calibration processes in `test_cutout_calibration.py`, including stability checks and depth calibration. - Created tests for frame compositing operations in `test_frame_compositing.py`, ensuring correct handling of depth and transparency. - Introduced material slot inheritance functionality in `material_slots.py` to manage geometry material assignments. - Built shared shader nodes for material processing in `shader.py`, including roughness and bump mapping setups.
This commit is contained in:
@@ -1,9 +1,6 @@
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
import numpy as np
|
||||
from tests.anyimage.server.support import load_server_module, FakeJobContext
|
||||
@@ -25,17 +22,6 @@ class ServerJobsTest(unittest.TestCase):
|
||||
cls.upscale_job = load_server_module("jobs.upscale")
|
||||
|
||||
|
||||
def test_upscale_progress_omits_single_frame_count(self):
|
||||
self.assertEqual(
|
||||
self.upscale_job.frame_progress_message("Upscaling image", 0, 1),
|
||||
"Upscaling image",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.upscale_job.frame_progress_message("Upscaling image", 1, 3),
|
||||
"Upscaling image (2/3)",
|
||||
)
|
||||
|
||||
|
||||
def test_plane_job_requests_depth_and_matching_normal_space(self):
|
||||
for plane_type, normal_space in (("DEPTH", "OBJECT"), ("RELIEF", "TANGENT")):
|
||||
normal_key = f"{normal_space.lower()}_normal"
|
||||
@@ -51,8 +37,8 @@ class ServerJobsTest(unittest.TestCase):
|
||||
with (
|
||||
patch.object(
|
||||
self.depth_plane_job,
|
||||
"input_frames",
|
||||
return_value=nullcontext([input_path]),
|
||||
"validate_input_path",
|
||||
return_value=input_path,
|
||||
),
|
||||
patch.object(
|
||||
self.depth_plane_job,
|
||||
@@ -96,8 +82,10 @@ class ServerJobsTest(unittest.TestCase):
|
||||
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_BASE", "device": "cpu"})
|
||||
with Image.open(root / "result" / "image.0001.png") as image:
|
||||
result = self.background.run(context, {"input": str(input_path), "model": "BEN2_BASE", "device": "cpu"})
|
||||
self.assertEqual(result, {"foreground": "foreground.png"})
|
||||
self.assertEqual([p.name for p in context.directory.glob("*.png")], ["foreground.png"])
|
||||
with Image.open(root / "result" / "foreground.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})
|
||||
@@ -152,27 +140,6 @@ class ServerJobsTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
def test_ben2_background_removal_processes_each_frame(self):
|
||||
from PIL import Image
|
||||
onnx_ben2 = load_server_module("models.onnx_ben2")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
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")
|
||||
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_BASE", "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))
|
||||
|
||||
|
||||
def test_background_alpha_job_preserves_float_precision(self):
|
||||
from PIL import Image
|
||||
onnx_ben2 = load_server_module("models.onnx_ben2")
|
||||
@@ -192,15 +159,13 @@ class ServerJobsTest(unittest.TestCase):
|
||||
self.assertFalse(list(context.directory.glob("*.png")))
|
||||
|
||||
|
||||
def test_upscale_job_saves_importable_upscale_sequence(self):
|
||||
def test_upscale_job_returns_one_image_file(self):
|
||||
from PIL import Image
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
frames = root / "frames"
|
||||
frames.mkdir()
|
||||
(frames / "a.png").touch()
|
||||
(frames / "b.png").touch()
|
||||
input_path = root / "input.png"
|
||||
input_path.touch()
|
||||
context = FakeJobContext(root / "result", "upscale-image")
|
||||
with patch.object(
|
||||
self.upscale_job,
|
||||
@@ -210,23 +175,19 @@ class ServerJobsTest(unittest.TestCase):
|
||||
result = self.upscale_job.run(
|
||||
context,
|
||||
{
|
||||
"input": str(frames),
|
||||
"input": str(input_path),
|
||||
"model": "REALESRGAN_GENERAL_WDN_X4V3",
|
||||
"video_frames": 2,
|
||||
|
||||
"device": "cpu",
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(infer_one.call_count, 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
||||
[False, True],
|
||||
)
|
||||
self.assertEqual(result, {"upscale": "."})
|
||||
colors = sorted(context.directory.glob("upscale.*.png"))
|
||||
self.assertEqual(infer_one.call_count, 1)
|
||||
self.assertEqual(result, {"upscale": "upscale.png"})
|
||||
colors = sorted(context.directory.glob("*.png"))
|
||||
self.assertEqual(
|
||||
[path.name for path in colors],
|
||||
["upscale.0001.png", "upscale.0002.png"],
|
||||
["upscale.png"],
|
||||
)
|
||||
self.assertIn(
|
||||
(
|
||||
@@ -261,12 +222,12 @@ class ServerJobsTest(unittest.TestCase):
|
||||
{
|
||||
"input": str(source_path),
|
||||
"model": "REALESRGAN_GENERAL_WDN_X4V3",
|
||||
"video_frames": 1,
|
||||
|
||||
"device": "cpu",
|
||||
},
|
||||
)
|
||||
|
||||
with Image.open(context.directory / "upscale.0001.png") as result:
|
||||
with Image.open(context.directory / "upscale.png") as result:
|
||||
self.assertEqual(result.mode, "RGBA")
|
||||
self.assertEqual(result.size, (10, 6))
|
||||
alpha = np.asarray(result.getchannel("A"))
|
||||
|
||||
Reference in New Issue
Block a user