Add image perspective editing and related tools
- Implemented the Image Perspective tool for editing image corners in Blender. - Created a new operator `EditImagePerspective` to handle the perspective transformation. - Added a preview feature for the perspective adjustment. - Introduced the `ImagePolylineTool` for editing image masks with polylines. - Developed backend support for Real-ESRGAN model inference, including ONNX integration. - Added functionality for image masking and background removal. - Created unit tests for the Real-ESRGAN ONNX model to ensure correct preprocessing and inference.
This commit is contained in:
@@ -87,6 +87,7 @@ class AddonTest(unittest.TestCase):
|
||||
cls.depth_operator = importlib.import_module(
|
||||
"anyimage.operators.depth_plane.operator"
|
||||
)
|
||||
cls.debug = importlib.import_module("anyimage.operators.debug")
|
||||
cls.depth_build = importlib.import_module(
|
||||
"anyimage.operators.depth_plane.build"
|
||||
)
|
||||
@@ -95,20 +96,26 @@ class AddonTest(unittest.TestCase):
|
||||
)
|
||||
cls.image_background = importlib.import_module("anyimage.operators.background")
|
||||
cls.image_mesh = importlib.import_module("anyimage.operators.surface")
|
||||
cls.image_extract_common = importlib.import_module(
|
||||
"anyimage.operators.image_extract.common"
|
||||
cls.image_edit_common = importlib.import_module(
|
||||
"anyimage.operators.image_edit.common"
|
||||
)
|
||||
cls.image_region = importlib.import_module(
|
||||
"anyimage.operators.image_extract.region"
|
||||
cls.image_mask = importlib.import_module(
|
||||
"anyimage.operators.image_edit.mask"
|
||||
)
|
||||
cls.image_box = importlib.import_module(
|
||||
"anyimage.operators.image_extract.box"
|
||||
"anyimage.operators.image_edit.box"
|
||||
)
|
||||
cls.image_lasso = importlib.import_module(
|
||||
"anyimage.operators.image_extract.lasso"
|
||||
"anyimage.operators.image_edit.lasso"
|
||||
)
|
||||
cls.image_line = importlib.import_module(
|
||||
"anyimage.operators.image_edit.line"
|
||||
)
|
||||
cls.image_polyline = importlib.import_module(
|
||||
"anyimage.operators.image_edit.polyline"
|
||||
)
|
||||
cls.image_perspective = importlib.import_module(
|
||||
"anyimage.operators.image_extract.perspective"
|
||||
"anyimage.operators.image_edit.perspective"
|
||||
)
|
||||
cls.balloon = importlib.import_module("anyimage.operators.balloon.operator")
|
||||
cls.balloon_geometry = importlib.import_module(
|
||||
@@ -253,6 +260,57 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertIn(self.job_runtime.JobOperatorBase, operator.__mro__)
|
||||
self.assertIn(self.fake_bpy.types.Operator, operator.__mro__)
|
||||
|
||||
def test_realesrgan_debug_operator_is_a_direct_model_request(self):
|
||||
operator_type = self.debug.DebugRealesrganGeneralX4v3
|
||||
self.assertIn(operator_type, self.addon.CLASSES)
|
||||
self.assertEqual(
|
||||
operator_type.bl_idname,
|
||||
"anyimage.debug_realesrgan_general_x4v3",
|
||||
)
|
||||
self.assertEqual(
|
||||
operator_type.job_type,
|
||||
"debug-realesrgan-general-x4v3",
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image_path = Path(directory) / "input.png"
|
||||
image_path.touch()
|
||||
settings = SimpleNamespace(
|
||||
input_path=str(image_path),
|
||||
device="CPU",
|
||||
video_frames=3,
|
||||
)
|
||||
context = SimpleNamespace(
|
||||
scene=SimpleNamespace(anyimage_settings=settings)
|
||||
)
|
||||
with (
|
||||
patch.object(self.debug, "require_environment"),
|
||||
patch.object(self.debug, "require_model") as require_model,
|
||||
):
|
||||
request = operator_type().request(context)
|
||||
|
||||
require_model.assert_called_once_with(
|
||||
self.addon.properties.REALESRGAN_MODEL_KEY
|
||||
)
|
||||
self.assertEqual(
|
||||
request,
|
||||
{
|
||||
"input": str(image_path),
|
||||
"model": "REALESRGAN_GENERAL_X4V3",
|
||||
"device": "cpu",
|
||||
"video_frames": 3,
|
||||
"start_frame": 1,
|
||||
},
|
||||
)
|
||||
self.assertNotIn("tile", request)
|
||||
self.assertNotIn("strength", request)
|
||||
|
||||
panel = importlib.import_module("anyimage.panel")
|
||||
self.assertIn(
|
||||
panel.RealesrganGeneralX4v3DebugPanel,
|
||||
self.addon.CLASSES,
|
||||
)
|
||||
|
||||
def test_default_request_sends_every_operator_property(self):
|
||||
class Example(self.job_runtime.JobOperatorBase):
|
||||
value: object()
|
||||
@@ -614,6 +672,62 @@ class AddonTest(unittest.TestCase):
|
||||
panel = importlib.import_module("anyimage.panel")
|
||||
self.assertFalse(hasattr(panel, "ModelsPanel"))
|
||||
|
||||
def test_preferences_draws_upscale_model_download(self):
|
||||
preference_module = importlib.import_module("anyimage.preferences")
|
||||
downloads = []
|
||||
|
||||
class Layout:
|
||||
enabled = True
|
||||
|
||||
def prop(self, *_args, **_options):
|
||||
return None
|
||||
|
||||
def separator(self, **_options):
|
||||
return None
|
||||
|
||||
def label(self, **_options):
|
||||
return None
|
||||
|
||||
def row(self, **_options):
|
||||
return self
|
||||
|
||||
def operator(self, identifier, **_options):
|
||||
operator = SimpleNamespace(model="")
|
||||
if identifier == "anyimage.download_model":
|
||||
downloads.append(operator)
|
||||
return operator
|
||||
|
||||
preferences = preference_module.AnyImagePreferences()
|
||||
preferences.layout = Layout()
|
||||
catalog = ({
|
||||
"key": "REALESRGAN_GENERAL_X4V3",
|
||||
"label": "Real-ESRGAN General x4v3",
|
||||
"family": "upscale",
|
||||
"enabled": True,
|
||||
"ready": False,
|
||||
},)
|
||||
with (
|
||||
patch.object(
|
||||
preference_module,
|
||||
"ai_status",
|
||||
return_value={
|
||||
"environment_ready": True,
|
||||
"missing_models": (),
|
||||
"ready": True,
|
||||
},
|
||||
),
|
||||
patch.object(
|
||||
preference_module,
|
||||
"model_catalog",
|
||||
return_value=catalog,
|
||||
),
|
||||
patch.object(preference_module.runtime, "server_busy", return_value=False),
|
||||
):
|
||||
preferences.draw(None)
|
||||
|
||||
self.assertEqual(len(downloads), 1)
|
||||
self.assertEqual(downloads[0].model, "REALESRGAN_GENERAL_X4V3")
|
||||
|
||||
def test_mesh_plane_operator_is_registered_with_material_options(self):
|
||||
operator = self.mesh_plane.ConvertToMeshPlane
|
||||
self.assertIn(operator, self.addon.CLASSES)
|
||||
@@ -793,141 +907,213 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertGreater(self.balloon_geometry.polygon_area(contour), 0.0)
|
||||
self.assertLess(float(lengths.std() / lengths.mean()), 0.15)
|
||||
|
||||
def test_image_region_clears_rgba_outside_the_selection(self):
|
||||
def test_image_mask_preserves_source_rgb_outside_the_selection(self):
|
||||
import numpy as np
|
||||
|
||||
top_down = np.full((3, 3, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32)
|
||||
source = np.flipud(top_down).ravel()
|
||||
pixels, bounds = self.image_region.extract_region_pixels(
|
||||
selection = self.image_mask.polygon_mask(
|
||||
(3, 3), ((0, 0), (3, 0), (0, 3))
|
||||
)
|
||||
pixels, bounds = self.image_mask.edit_image_mask_pixels(
|
||||
source,
|
||||
(3, 3),
|
||||
((0, 0), (3, 0), (0, 3)),
|
||||
selection,
|
||||
)
|
||||
result = np.flipud(np.asarray(pixels).reshape((3, 3, 4)))
|
||||
width = bounds[2] - bounds[0]
|
||||
height = bounds[3] - bounds[1]
|
||||
result = np.flipud(np.asarray(pixels).reshape((height, width, 4)))
|
||||
|
||||
self.assertEqual(bounds, (0, 0, 3, 3))
|
||||
np.testing.assert_allclose(result[2, 2], (0.0, 0.0, 0.0, 0.0))
|
||||
self.assertEqual(bounds, (0, 0, 2, 2))
|
||||
np.testing.assert_allclose(result[1, 1], (0.2, 0.4, 0.6, 0.0))
|
||||
self.assertEqual(float(result[0, 0, 3]), 1.0)
|
||||
|
||||
def test_image_region_clears_hidden_rgb_from_source_transparency(self):
|
||||
def test_image_mask_preserves_hidden_rgb_from_source_transparency(self):
|
||||
import numpy as np
|
||||
|
||||
top_down = np.full((2, 2, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32)
|
||||
top_down[0, 0, 3] = 0.0
|
||||
pixels, _bounds = self.image_region.extract_region_pixels(
|
||||
pixels, _bounds = self.image_mask.edit_image_mask_pixels(
|
||||
np.flipud(top_down).ravel(),
|
||||
(2, 2),
|
||||
((0, 0), (2, 0), (2, 2), (0, 2)),
|
||||
np.ones((2, 2), dtype=np.float32),
|
||||
)
|
||||
result = np.flipud(np.asarray(pixels).reshape((2, 2, 4)))
|
||||
|
||||
np.testing.assert_allclose(result[0, 0], (0.0, 0.0, 0.0, 0.0))
|
||||
np.testing.assert_allclose(result[0, 0], (0.2, 0.4, 0.6, 0.0))
|
||||
|
||||
def test_image_region_keeps_out_of_image_canvas_transparent(self):
|
||||
def test_image_mask_invert_runs_after_refinement(self):
|
||||
import numpy as np
|
||||
|
||||
top_down = np.ones((2, 2, 4), dtype=np.float32)
|
||||
pixels, bounds = self.image_region.extract_region_pixels(
|
||||
top_down = np.full((2, 2, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32)
|
||||
refined = np.asarray(((1.0, 0.25), (0.0, 1.0)), dtype=np.float32)
|
||||
pixels, bounds = self.image_mask.edit_image_mask_pixels(
|
||||
np.flipud(top_down).ravel(),
|
||||
(2, 2),
|
||||
((-1, -1), (2, -1), (2, 2), (-1, 2)),
|
||||
refined,
|
||||
invert=True,
|
||||
)
|
||||
result = np.flipud(np.asarray(pixels).reshape((3, 3, 4)))
|
||||
result = np.flipud(np.asarray(pixels).reshape((2, 2, 4)))
|
||||
|
||||
self.assertEqual(bounds, (-1, -1, 2, 2))
|
||||
self.assertEqual(bounds, (0, 0, 2, 2))
|
||||
self.assertEqual(float(result[0, 0, 3]), 0.0)
|
||||
self.assertEqual(float(result[1, 1, 3]), 1.0)
|
||||
self.assertAlmostEqual(float(result[0, 1, 3]), 0.75)
|
||||
np.testing.assert_allclose(result[0, 0], (0.2, 0.4, 0.6, 0.0))
|
||||
|
||||
def test_image_region_rejects_a_fully_out_of_image_region(self):
|
||||
def test_image_mask_rejects_a_fully_transparent_result(self):
|
||||
import numpy as np
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "does not overlap"):
|
||||
self.image_region.extract_region_pixels(
|
||||
with self.assertRaisesRegex(ValueError, "no visible"):
|
||||
self.image_mask.edit_image_mask_pixels(
|
||||
np.ones(2 * 2 * 4, dtype=np.float32),
|
||||
(2, 2),
|
||||
((3, 3), (4, 3), (4, 4), (3, 4)),
|
||||
np.zeros((2, 2), dtype=np.float32),
|
||||
)
|
||||
|
||||
def test_image_region_crops_to_visible_alpha_with_pixel_padding(self):
|
||||
def test_image_mask_crops_to_visible_alpha_without_padding(self):
|
||||
import numpy as np
|
||||
|
||||
rgba = np.zeros((4, 5, 4), dtype=np.float32)
|
||||
rgba[1:3, 2:4] = (0.2, 0.4, 0.6, 1.0)
|
||||
pixels, bounds = self.image_extract_common.crop_pixels_to_alpha(
|
||||
rgba = np.full((4, 5, 4), (0.8, 0.7, 0.6, 0.0), dtype=np.float32)
|
||||
rgba[1, 2] = (0.2, 0.4, 0.6, 1.0)
|
||||
rgba[2, 3] = (0.2, 0.4, 0.6, 1.0)
|
||||
pixels, bounds = self.image_edit_common.crop_pixels_to_alpha(
|
||||
np.flipud(rgba).ravel(),
|
||||
(10, 20, 15, 24),
|
||||
padding=1,
|
||||
)
|
||||
result = np.flipud(np.asarray(pixels).reshape((4, 4, 4)))
|
||||
result = np.flipud(np.asarray(pixels).reshape((2, 2, 4)))
|
||||
|
||||
self.assertEqual(bounds, (11, 20, 15, 24))
|
||||
self.assertEqual(result.shape, (4, 4, 4))
|
||||
self.assertTrue(np.all(result[[0, -1], :, 3] == 0.0))
|
||||
self.assertTrue(np.all(result[:, [0, -1], 3] == 0.0))
|
||||
self.assertTrue(np.all(result[1:3, 1:3, 3] == 1.0))
|
||||
self.assertEqual(bounds, (12, 21, 14, 23))
|
||||
self.assertEqual(result.shape, (2, 2, 4))
|
||||
np.testing.assert_allclose(result[0, 1], (0.8, 0.7, 0.6, 0.0))
|
||||
|
||||
def test_image_extract_tools_and_operators_are_registered(self):
|
||||
def test_image_mask_restores_refined_alpha_to_source_coordinates(self):
|
||||
import numpy as np
|
||||
|
||||
top_down = np.zeros((2, 2, 4), dtype=np.float32)
|
||||
top_down[:, :, 3] = ((1.0, 0.5), (0.0, 1.0))
|
||||
mask = self.image_mask.refined_mask_from_pixels(
|
||||
np.flipud(top_down).ravel(),
|
||||
(2, 2),
|
||||
(1, 1, 3, 3),
|
||||
(4, 4),
|
||||
)
|
||||
|
||||
expected = np.zeros((4, 4), dtype=np.float32)
|
||||
expected[1:3, 1:3] = top_down[:, :, 3]
|
||||
np.testing.assert_allclose(mask, expected)
|
||||
|
||||
def test_image_edit_aligns_canvas_without_changing_object_transform(self):
|
||||
matrix = object()
|
||||
image_object = SimpleNamespace(
|
||||
matrix_world=matrix,
|
||||
empty_display_size=4.0,
|
||||
empty_image_offset=(0.0, 0.0),
|
||||
)
|
||||
with patch.object(
|
||||
self.image_edit_common,
|
||||
"plane_bounds",
|
||||
return_value=(-2.0, 2.0, -1.0, 1.0),
|
||||
):
|
||||
self.image_edit_common.align_empty_display_to_bounds(
|
||||
image_object,
|
||||
(100, 50),
|
||||
(50, 25),
|
||||
(50, 0, 100, 25),
|
||||
)
|
||||
|
||||
self.assertIs(image_object.matrix_world, matrix)
|
||||
self.assertEqual(image_object.empty_display_size, 2.0)
|
||||
self.assertEqual(image_object.empty_image_offset, (0.0, 0.0))
|
||||
|
||||
def test_image_edit_tools_and_operators_are_registered(self):
|
||||
import inspect
|
||||
|
||||
self.assertIn(self.image_region.ExtractImageRegion, self.addon.CLASSES)
|
||||
self.assertIn(
|
||||
self.image_region.RemoveImageRegionBackground, self.addon.CLASSES
|
||||
operator = self.image_mask.EditImageMask
|
||||
self.assertIn(operator, self.addon.CLASSES)
|
||||
self.assertIn(self.image_mask.RefineImageMask, self.addon.CLASSES)
|
||||
self.assertIn(self.image_mask.ActivateImageEditTool, self.addon.CLASSES)
|
||||
tool_cases = (
|
||||
(self.image_box.ImageBoxTool, "BOX", "ops.sculpt.border_mask"),
|
||||
(self.image_lasso.ImageLassoTool, "LASSO", "ops.sculpt.lasso_mask"),
|
||||
(self.image_line.ImageLineTool, "LINE", "ops.sculpt.line_mask"),
|
||||
(
|
||||
self.image_polyline.ImagePolylineTool,
|
||||
"POLYLINE",
|
||||
"ops.sculpt.polyline_mask",
|
||||
),
|
||||
)
|
||||
self.assertIn(self.image_region.ActivateImageExtractTool, self.addon.CLASSES)
|
||||
self.assertNotIn(self.image_lasso.ImageLassoTool, self.addon.CLASSES)
|
||||
self.assertNotIn(self.image_box.ImageBoxTool, self.addon.CLASSES)
|
||||
self.assertEqual(
|
||||
self.image_lasso.ImageLassoTool.bl_operator,
|
||||
self.image_region.ExtractImageRegion.bl_idname,
|
||||
for tool, gesture, icon in tool_cases:
|
||||
self.assertNotIn(tool, self.addon.CLASSES)
|
||||
self.assertEqual(tool.bl_operator, operator.bl_idname)
|
||||
self.assertEqual(tool.bl_icon, icon)
|
||||
self.assertEqual(
|
||||
tool.bl_keymap[0][2]["properties"], [("gesture", gesture)]
|
||||
)
|
||||
self.assertNotIn("bl_cursor", tool.__dict__)
|
||||
self.assertIn("draw_settings", self.image_box.ImageBoxTool.__dict__)
|
||||
self.assertIn("draw_settings", self.image_lasso.ImageLassoTool.__dict__)
|
||||
self.assertNotIn("draw_settings", self.image_line.ImageLineTool.__dict__)
|
||||
self.assertNotIn(
|
||||
"draw_settings", self.image_polyline.ImagePolylineTool.__dict__
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_lasso.ImageLassoTool.bl_icon,
|
||||
"ops.sculpt.lasso_mask",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_box.ImageBoxTool.bl_icon,
|
||||
"ops.sculpt.border_mask",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_box.ImageBoxTool.bl_keymap[0][2]["properties"],
|
||||
[("selection_shape", "BOX")],
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_lasso.ImageLassoTool.bl_keymap[0][2]["properties"],
|
||||
[("selection_shape", "LASSO")],
|
||||
)
|
||||
self.assertNotIn("bl_cursor", self.image_box.ImageBoxTool.__dict__)
|
||||
self.assertNotIn("bl_cursor", self.image_lasso.ImageLassoTool.__dict__)
|
||||
self.assertEqual(
|
||||
self.image_region.ExtractImageRegion.__annotations__["selection_shape"][
|
||||
"default"
|
||||
],
|
||||
"BOX",
|
||||
)
|
||||
operator = self.image_region.ExtractImageRegion
|
||||
self.assertIn("continuous", operator.__annotations__)
|
||||
self.assertFalse(operator.__annotations__["continuous"]["default"])
|
||||
self.assertIn("padding", operator.__annotations__)
|
||||
self.assertEqual(operator.__annotations__["padding"]["default"], 0)
|
||||
self.assertEqual(operator.__annotations__["gesture"]["default"], "BOX")
|
||||
self.assertIn("invert", operator.__annotations__)
|
||||
self.assertNotIn("continuous", operator.__annotations__)
|
||||
self.assertNotIn("padding", operator.__annotations__)
|
||||
settings = self.addon.properties.AnyImageSettings.__annotations__
|
||||
self.assertFalse(settings["image_extract_continuous"]["default"])
|
||||
self.assertEqual(settings["image_extract_padding"]["default"], 0)
|
||||
self.assertEqual(
|
||||
settings["image_extract_remove_background"]["name"],
|
||||
"Remove Background",
|
||||
)
|
||||
source_picker = inspect.getsource(
|
||||
self.image_extract_common.resolve_image_extract_click
|
||||
)
|
||||
continuous_offset = inspect.getsource(
|
||||
self.image_extract_common.offset_continuous_result
|
||||
)
|
||||
self.assertFalse(settings["image_edit_refine_selection"]["default"])
|
||||
self.assertNotIn("image_extract_continuous", settings)
|
||||
self.assertNotIn("image_extract_padding", settings)
|
||||
self.assertNotIn("image_extract_remove_background", settings)
|
||||
source_picker = inspect.getsource(self.image_edit_common.resolve_image_edit_click)
|
||||
self.assertIn("bpy.ops.view3d.select(", source_picker)
|
||||
self.assertIn("Click an Image Empty", source_picker)
|
||||
self.assertIn("_restore_selection(", source_picker)
|
||||
self.assertIn("EXTRACT_LAYER_STEP * layer", continuous_offset)
|
||||
modal_source = inspect.getsource(operator.modal)
|
||||
self.assertIn('event.type == "F"', modal_source)
|
||||
self.assertIn("self.invert = not self.invert", modal_source)
|
||||
|
||||
def test_image_extract_first_click_selects_another_image_without_extracting(self):
|
||||
def test_image_line_and_polyline_disable_refine_selection(self):
|
||||
settings = SimpleNamespace(image_edit_refine_selection=True)
|
||||
context = SimpleNamespace(scene=SimpleNamespace(anyimage_settings=settings))
|
||||
with patch.object(self.image_edit_common, "ai_ready", return_value=True):
|
||||
for gesture in ("LINE", "POLYLINE"):
|
||||
operator = SimpleNamespace(
|
||||
gesture=gesture,
|
||||
refine_selection=True,
|
||||
)
|
||||
self.image_edit_common.apply_image_mask_settings(operator, context)
|
||||
self.assertFalse(operator.refine_selection)
|
||||
|
||||
operator = SimpleNamespace(gesture="BOX", refine_selection=False)
|
||||
self.image_edit_common.apply_image_mask_settings(operator, context)
|
||||
self.assertTrue(operator.refine_selection)
|
||||
|
||||
def test_image_line_click_is_cancelled_before_projection(self):
|
||||
import inspect
|
||||
|
||||
modal_source = inspect.getsource(self.image_mask.EditImageMask.modal)
|
||||
projection_source = inspect.getsource(
|
||||
self.image_mask.image_line_half_plane_polygon
|
||||
)
|
||||
self.assertIn("math.dist(self._path[0], self._cursor) < 2.0", modal_source)
|
||||
self.assertIn("if len(line) != 2", projection_source)
|
||||
|
||||
def test_image_polyline_closes_from_start_or_enter(self):
|
||||
import inspect
|
||||
|
||||
modal_source = inspect.getsource(
|
||||
self.image_mask.EditImageMask._modal_polyline
|
||||
)
|
||||
draw_source = inspect.getsource(
|
||||
self.image_mask.EditImageMask._draw_selection
|
||||
)
|
||||
self.assertIn("return self._confirm_polyline(context)", modal_source)
|
||||
self.assertNotIn("DOUBLE_CLICK", modal_source)
|
||||
self.assertIn('event.type in {"RET", "NUMPAD_ENTER"}', modal_source)
|
||||
self.assertIn("_draw_polyline_close_indicator()", draw_source)
|
||||
|
||||
def test_image_edit_first_click_selects_another_image_without_editing(self):
|
||||
selected = []
|
||||
|
||||
def image(name):
|
||||
@@ -962,8 +1148,8 @@ class AddonTest(unittest.TestCase):
|
||||
clicked.name: clicked,
|
||||
}
|
||||
try:
|
||||
source, should_extract = (
|
||||
self.image_extract_common.resolve_image_extract_click(
|
||||
source, should_edit = (
|
||||
self.image_edit_common.resolve_image_edit_click(
|
||||
context, event, lambda *_args: None
|
||||
)
|
||||
)
|
||||
@@ -975,10 +1161,10 @@ class AddonTest(unittest.TestCase):
|
||||
self.fake_bpy.ops = previous_ops
|
||||
|
||||
self.assertIs(source, clicked)
|
||||
self.assertFalse(should_extract)
|
||||
self.assertFalse(should_edit)
|
||||
self.assertEqual(selected, [])
|
||||
|
||||
def test_image_extract_first_click_on_non_image_keeps_active_image(self):
|
||||
def test_image_edit_first_click_on_non_image_keeps_active_image(self):
|
||||
selection_calls = []
|
||||
active = SimpleNamespace(
|
||||
name="Active",
|
||||
@@ -1009,8 +1195,8 @@ class AddonTest(unittest.TestCase):
|
||||
)
|
||||
self.fake_bpy.data.objects = {active.name: active, clicked.name: clicked}
|
||||
try:
|
||||
source, should_extract = (
|
||||
self.image_extract_common.resolve_image_extract_click(
|
||||
source, should_edit = (
|
||||
self.image_edit_common.resolve_image_edit_click(
|
||||
context, event, lambda *_args: None
|
||||
)
|
||||
)
|
||||
@@ -1022,7 +1208,7 @@ class AddonTest(unittest.TestCase):
|
||||
self.fake_bpy.ops = previous_ops
|
||||
|
||||
self.assertIs(source, active)
|
||||
self.assertTrue(should_extract)
|
||||
self.assertTrue(should_edit)
|
||||
self.assertIs(view_objects.active, active)
|
||||
self.assertEqual(selection_calls, [False, True])
|
||||
|
||||
@@ -1045,7 +1231,7 @@ class AddonTest(unittest.TestCase):
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
def test_image_perspective_crops_alpha_and_applies_padding(self):
|
||||
def test_image_perspective_tightly_crops_alpha(self):
|
||||
import numpy as np
|
||||
|
||||
top_down = np.zeros((4, 4, 4), dtype=np.float32)
|
||||
@@ -1056,16 +1242,13 @@ class AddonTest(unittest.TestCase):
|
||||
(4, 4),
|
||||
((0, 0), (4, 0), (4, 4), (0, 4)),
|
||||
1.0,
|
||||
1,
|
||||
)
|
||||
)
|
||||
result = np.flipud(np.asarray(pixels).reshape((4, 4, 4)))
|
||||
result = np.flipud(np.asarray(pixels).reshape((2, 2, 4)))
|
||||
|
||||
self.assertEqual(output_size, (4, 4))
|
||||
self.assertEqual(placement_bounds, (0.0, 0.0, 4.0, 4.0))
|
||||
self.assertTrue(np.all(result[[0, -1], :, 3] == 0.0))
|
||||
self.assertTrue(np.all(result[:, [0, -1], 3] == 0.0))
|
||||
self.assertTrue(np.all(result[1:3, 1:3, 3] == 1.0))
|
||||
self.assertEqual(output_size, (2, 2))
|
||||
self.assertEqual(placement_bounds, (1.0, 1.0, 3.0, 3.0))
|
||||
self.assertTrue(np.all(result[:, :, 3] == 1.0))
|
||||
|
||||
def test_image_perspective_rejects_non_convex_corners(self):
|
||||
with self.assertRaisesRegex(ValueError, "convex"):
|
||||
@@ -1100,7 +1283,7 @@ class AddonTest(unittest.TestCase):
|
||||
|
||||
def test_image_perspective_builds_polyline_mask_fill_triangles(self):
|
||||
self.assertEqual(
|
||||
self.image_perspective.triangle_fan_vertices(
|
||||
self.image_edit_common.triangle_fan_vertices(
|
||||
((0, 0), (2, 0), (2, 2), (0, 2))
|
||||
),
|
||||
(
|
||||
@@ -1153,7 +1336,7 @@ class AddonTest(unittest.TestCase):
|
||||
|
||||
self.assertAlmostEqual(ratio, 2.0)
|
||||
|
||||
def test_image_perspective_maps_padding_to_source_plane_bounds(self):
|
||||
def test_image_perspective_maps_crop_to_source_plane_bounds(self):
|
||||
bounds = self.image_perspective.perspective_placement_bounds(
|
||||
((10, 20), (30, 20), (30, 40), (10, 40)),
|
||||
(100, 50),
|
||||
@@ -1162,39 +1345,29 @@ class AddonTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(bounds, (8.0, 18.0, 32.0, 42.0))
|
||||
|
||||
def test_image_perspective_tool_uses_four_clicks_and_shared_style(self):
|
||||
def test_image_perspective_tool_uses_four_clicks_and_deform_behavior(self):
|
||||
import inspect
|
||||
|
||||
operator = self.image_perspective.ExtractImagePerspective
|
||||
operator = self.image_perspective.EditImagePerspective
|
||||
tool = self.image_perspective.ImagePerspectiveTool
|
||||
settings = self.addon.properties.AnyImageSettings.__annotations__
|
||||
|
||||
self.assertIn(operator, self.addon.CLASSES)
|
||||
self.assertIn(
|
||||
self.image_perspective.RemoveImagePerspectiveBackground,
|
||||
self.addon.CLASSES,
|
||||
)
|
||||
self.assertNotIn(tool, self.addon.CLASSES)
|
||||
self.assertEqual(tool.bl_operator, operator.bl_idname)
|
||||
self.assertEqual(tool.bl_icon, "ops.sculpt.polyline_mask")
|
||||
self.assertEqual(tool.bl_icon, "ops.sculpt.line_project")
|
||||
self.assertNotIn("bl_cursor", tool.__dict__)
|
||||
self.assertIn("draw_settings", tool.__dict__)
|
||||
self.assertNotIn("draw_settings", tool.__dict__)
|
||||
self.assertNotIn("image_perspective_aspect_mode", settings)
|
||||
self.assertNotIn("image_perspective_aspect_ratio", settings)
|
||||
self.assertNotIn("aspect_mode", operator.__annotations__)
|
||||
self.assertIn("aspect_ratio", operator.__annotations__)
|
||||
self.assertIn("remove_background", operator.__annotations__)
|
||||
self.assertIn("continuous", operator.__annotations__)
|
||||
self.assertIn("padding", operator.__annotations__)
|
||||
self.assertNotIn("refine_selection", operator.__annotations__)
|
||||
self.assertNotIn("continuous", operator.__annotations__)
|
||||
self.assertNotIn("padding", operator.__annotations__)
|
||||
modal_source = inspect.getsource(operator.modal)
|
||||
invoke_source = inspect.getsource(operator.invoke)
|
||||
execute_source = inspect.getsource(operator.execute)
|
||||
background_request_source = inspect.getsource(
|
||||
self.image_perspective.RemoveImagePerspectiveBackground.request
|
||||
)
|
||||
background_response_source = inspect.getsource(
|
||||
self.image_perspective.RemoveImagePerspectiveBackground.response
|
||||
)
|
||||
draw_source = inspect.getsource(operator._draw_overlay)
|
||||
preview_source = inspect.getsource(operator._draw_preview)
|
||||
finish_source = inspect.getsource(operator._finish)
|
||||
@@ -1208,17 +1381,9 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertIn("active_view3d_tool_id(context)", modal_source)
|
||||
self.assertIn("drawing_in_region(", modal_source)
|
||||
self.assertNotIn("context.region_data.is_perspective", invoke_source)
|
||||
self.assertIn("resolve_image_extract_click(", invoke_source)
|
||||
self.assertIn("apply_image_extract_settings(self, context)", invoke_source)
|
||||
self.assertIn(
|
||||
"bpy.ops.anyimage.remove_image_perspective_background(",
|
||||
execute_source,
|
||||
)
|
||||
self.assertIn("resolve_image_edit_click(", invoke_source)
|
||||
self.assertIn("extract_perspective_pixels(", execute_source)
|
||||
self.assertIn("create_perspective_result(", execute_source)
|
||||
self.assertIn("image_extract_background_request(", background_request_source)
|
||||
self.assertIn("placement_polygon=polygon", background_response_source)
|
||||
self.assertIn("self.padding", background_response_source)
|
||||
self.assertIn("draw_polygon_fill(vertices)", draw_source)
|
||||
self.assertIn("self._draw_preview()", draw_source)
|
||||
self.assertIn("draw_texture_2d(", preview_source)
|
||||
@@ -1230,16 +1395,16 @@ class AddonTest(unittest.TestCase):
|
||||
|
||||
def test_image_box_builds_a_rectangular_screen_path(self):
|
||||
self.assertEqual(
|
||||
self.image_region.box_screen_path(((2, 3), (8, 9))),
|
||||
self.image_mask.box_screen_path((2, 3), (8, 9)),
|
||||
((2, 3), (8, 3), (8, 9), (2, 9)),
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_region.box_screen_path(((2, 3), (3, 9))),
|
||||
self.image_mask.box_screen_path((2, 3), (3, 9)),
|
||||
(),
|
||||
)
|
||||
|
||||
def test_image_extract_selection_uses_white_dashed_segments(self):
|
||||
segments = self.image_extract_common.dashed_line_segments(
|
||||
def test_image_edit_selection_uses_white_dashed_segments(self):
|
||||
segments = self.image_edit_common.dashed_line_segments(
|
||||
((0, 0), (10, 0)), dash_length=3.0, gap_length=2.0
|
||||
)
|
||||
|
||||
@@ -1248,22 +1413,29 @@ class AddonTest(unittest.TestCase):
|
||||
((0.0, 0.0), (3.0, 0.0), (5.0, 0.0), (8.0, 0.0)),
|
||||
)
|
||||
|
||||
def test_image_extract_activation_defaults_to_box(self):
|
||||
def test_image_edit_activation_defaults_to_box(self):
|
||||
import inspect
|
||||
|
||||
source = inspect.getsource(self.image_region.ActivateImageExtractTool.execute)
|
||||
source = inspect.getsource(self.image_mask.ActivateImageEditTool.execute)
|
||||
self.assertIn('name="anyimage.image_box"', source)
|
||||
|
||||
def test_continuous_keeps_the_source_empty_selected(self):
|
||||
import inspect
|
||||
def test_image_line_snaps_and_keeps_the_previewed_half_plane(self):
|
||||
import math
|
||||
|
||||
execute_source = inspect.getsource(self.image_region.ExtractImageRegion.execute)
|
||||
response_source = inspect.getsource(
|
||||
self.image_region.RemoveImageRegionBackground.response
|
||||
endpoint = self.image_mask.snapped_line_endpoint((0, 0), (10, 1), True)
|
||||
self.assertAlmostEqual(math.degrees(math.atan2(endpoint[1], endpoint[0])), 5.0)
|
||||
preview = self.image_mask.line_half_plane_polygon(
|
||||
(0, 5), (10, 5), (10, 10)
|
||||
)
|
||||
image = self.image_mask.image_line_half_plane_polygon(
|
||||
((0, 5), (10, 5)), (10, 10)
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_edit_common.polygon_bounds(preview), (0, 5, 10, 10)
|
||||
)
|
||||
self.assertEqual(
|
||||
self.image_edit_common.polygon_bounds(image), (0, 0, 10, 5)
|
||||
)
|
||||
expected = "select_image_extract_result("
|
||||
self.assertIn(expected, execute_source)
|
||||
self.assertIn(expected, response_source)
|
||||
|
||||
def test_balloon_interior_points_have_uniform_spacing(self):
|
||||
import numpy as np
|
||||
@@ -2160,12 +2332,25 @@ class AddonTest(unittest.TestCase):
|
||||
},
|
||||
),
|
||||
(
|
||||
self.image_perspective.ImagePerspectiveTool,
|
||||
self.image_line.ImageLineTool,
|
||||
{
|
||||
"after": {self.image_lasso.ImageLassoTool.bl_idname},
|
||||
"separator": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
self.image_polyline.ImagePolylineTool,
|
||||
{
|
||||
"after": {self.image_line.ImageLineTool.bl_idname},
|
||||
"separator": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
self.image_perspective.ImagePerspectiveTool,
|
||||
{
|
||||
"separator": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
self.balloon.BalloonTool,
|
||||
{
|
||||
|
||||
@@ -89,6 +89,35 @@ class ModelTest(unittest.TestCase):
|
||||
"BEN2-ONNX",
|
||||
)
|
||||
|
||||
def test_realesrgan_general_x4v3_uses_verified_onnx_artifact(self):
|
||||
model = self.models.get_downloadable_model(
|
||||
"REALESRGAN_GENERAL_X4V3"
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
model.repository,
|
||||
"jonathanst29/tinier-upscale-models",
|
||||
)
|
||||
self.assertEqual(model.license_name, "BSD 3-Clause")
|
||||
self.assertEqual(
|
||||
model.download_patterns,
|
||||
("realesr-general-x4v3.onnx",),
|
||||
)
|
||||
self.assertEqual(model.ready_patterns, model.download_patterns)
|
||||
self.assertEqual(model.r2_directory, "realesr-general-x4v3")
|
||||
self.assertEqual(
|
||||
model.r2_files,
|
||||
((
|
||||
"realesr-general-x4v3.onnx",
|
||||
4866396,
|
||||
"924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5",
|
||||
),),
|
||||
)
|
||||
self.assertEqual(
|
||||
self.models.model_record(model)["family"],
|
||||
"upscale",
|
||||
)
|
||||
|
||||
def test_enabled_da3_models_use_onnx_and_r2_checksums(self):
|
||||
for model in self.models.MODELS.values():
|
||||
if not model.enabled:
|
||||
|
||||
@@ -68,14 +68,14 @@ class PackTest(unittest.TestCase):
|
||||
self.assertIn("operators/mesh_plane.py", names)
|
||||
self.assertIn("operators/image.py", names)
|
||||
self.assertIn("operators/balloon/operator.py", names)
|
||||
self.assertIn("operators/image_extract/__init__.py", names)
|
||||
self.assertIn("operators/image_extract/common.py", names)
|
||||
self.assertIn("operators/image_extract/region.py", names)
|
||||
self.assertIn("operators/image_extract/box.py", names)
|
||||
self.assertIn("operators/image_extract/lasso.py", names)
|
||||
self.assertIn("operators/image_extract/perspective.py", names)
|
||||
self.assertNotIn("operators/image_lasso.py", names)
|
||||
self.assertNotIn("operators/image_perspective.py", names)
|
||||
self.assertIn("operators/image_edit/__init__.py", names)
|
||||
self.assertIn("operators/image_edit/common.py", names)
|
||||
self.assertIn("operators/image_edit/mask.py", names)
|
||||
self.assertIn("operators/image_edit/box.py", names)
|
||||
self.assertIn("operators/image_edit/lasso.py", names)
|
||||
self.assertIn("operators/image_edit/line.py", names)
|
||||
self.assertIn("operators/image_edit/polyline.py", names)
|
||||
self.assertIn("operators/image_edit/perspective.py", names)
|
||||
self.assertIn("operators/surface.py", names)
|
||||
self.assertIn(
|
||||
"assets/O_AnyImage.blend",
|
||||
@@ -86,8 +86,7 @@ class PackTest(unittest.TestCase):
|
||||
self.assertIn("server/geometry/grid_mesh.py", names)
|
||||
self.assertIn("server/geometry/surface.py", names)
|
||||
self.assertIn("server/jobs/background.py", names)
|
||||
self.assertIn("server/jobs/image_extract.py", names)
|
||||
self.assertNotIn("server/jobs/image_lasso.py", names)
|
||||
self.assertIn("server/jobs/image_mask.py", names)
|
||||
self.assertIn("server/models/ben2.py", names)
|
||||
self.assertIn("server/models/da3.py", names)
|
||||
self.assertIn("server/models/moge2.py", names)
|
||||
@@ -96,6 +95,7 @@ class PackTest(unittest.TestCase):
|
||||
self.assertIn("server/jobs/debug_da3.py", names)
|
||||
self.assertIn("server/jobs/debug_moge2.py", names)
|
||||
self.assertIn("server/jobs/debug_ben2.py", names)
|
||||
self.assertIn("server/jobs/debug_realesrgan.py", names)
|
||||
self.assertNotIn("server/processing.py", names)
|
||||
self.assertNotIn("server/progress.py", names)
|
||||
self.assertIn("model_catalog.py", names)
|
||||
@@ -103,6 +103,8 @@ class PackTest(unittest.TestCase):
|
||||
self.assertIn("server/models/onnx_ben2.py", names)
|
||||
self.assertIn("server/models/onnx_da3.py", names)
|
||||
self.assertIn("server/models/onnx_moge2.py", names)
|
||||
self.assertIn("server/models/onnx_realesrgan.py", names)
|
||||
self.assertIn("server/models/realesrgan.py", names)
|
||||
self.assertNotIn("server/da3.py", names)
|
||||
self.assertNotIn("server/moge2.py", names)
|
||||
self.assertIn("server/__init__.py", names)
|
||||
|
||||
75
tests/test_realesrgan_onnx.py
Normal file
75
tests/test_realesrgan_onnx.py
Normal file
@@ -0,0 +1,75 @@
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
SOURCE_DIR = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "src"
|
||||
/ "anyimage"
|
||||
/ "server"
|
||||
/ "models"
|
||||
)
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
import onnx_realesrgan
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def get_inputs(self):
|
||||
return [SimpleNamespace(name="input")]
|
||||
|
||||
def get_outputs(self):
|
||||
return [SimpleNamespace(name="output")]
|
||||
|
||||
def run(self, output_names, feeds):
|
||||
self.output_names = output_names
|
||||
self.feeds = feeds
|
||||
source = feeds["input"]
|
||||
return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)]
|
||||
|
||||
|
||||
class RealesrganOnnxTest(unittest.TestCase):
|
||||
def test_preprocess_matches_dynamic_rgb_layout(self):
|
||||
image = Image.new("RGB", (7, 5), (64, 128, 192))
|
||||
|
||||
result = onnx_realesrgan.preprocess(image)
|
||||
|
||||
self.assertEqual(result.shape, (1, 3, 5, 7))
|
||||
self.assertEqual(result.dtype, np.float32)
|
||||
self.assertTrue(result.flags.c_contiguous)
|
||||
np.testing.assert_allclose(
|
||||
result[0, :, 0, 0],
|
||||
np.array((64, 128, 192), dtype=np.float32) / 255.0,
|
||||
)
|
||||
|
||||
def test_infer_returns_direct_four_x_rgb_prediction(self):
|
||||
session = FakeSession()
|
||||
image = Image.new("RGB", (7, 5), (10, 20, 30))
|
||||
|
||||
result = onnx_realesrgan.infer(session, image)
|
||||
|
||||
self.assertEqual(result.mode, "RGB")
|
||||
self.assertEqual(result.size, (28, 20))
|
||||
self.assertEqual(session.output_names, ["output"])
|
||||
self.assertEqual(session.feeds["input"].shape, (1, 3, 5, 7))
|
||||
self.assertEqual(result.getpixel((0, 0)), (10, 20, 30))
|
||||
|
||||
def test_infer_rejects_non_four_x_output(self):
|
||||
class InvalidSession(FakeSession):
|
||||
def run(self, _output_names, feeds):
|
||||
return [feeds["input"]]
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "unexpected output shape"):
|
||||
onnx_realesrgan.infer(
|
||||
InvalidSession(),
|
||||
Image.new("RGB", (7, 5)),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -79,7 +79,14 @@ class ServerTest(unittest.TestCase):
|
||||
snapshot = manager.snapshot()
|
||||
self.assertEqual(
|
||||
snapshot,
|
||||
{"loaded": {"ben2": "", "da3": "", "moge2": ""}},
|
||||
{
|
||||
"loaded": {
|
||||
"ben2": "",
|
||||
"da3": "",
|
||||
"moge2": "",
|
||||
"realesrgan": "",
|
||||
}
|
||||
},
|
||||
)
|
||||
self.assertIn("server.model_catalog", sys.modules)
|
||||
self.assertNotIn("model_catalog", sys.modules)
|
||||
@@ -94,6 +101,8 @@ class ServerTest(unittest.TestCase):
|
||||
cache.da3_key = None
|
||||
cache.moge_model = None
|
||||
cache.moge_key = None
|
||||
cache.realesrgan_model = None
|
||||
cache.realesrgan_key = None
|
||||
return cache
|
||||
|
||||
def create_server(self):
|
||||
@@ -184,6 +193,25 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertIsNone(cache.da3_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
|
||||
def test_realesrgan_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
realesrgan = ModuleType("server.models.onnx_realesrgan")
|
||||
realesrgan.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"server.models.onnx_realesrgan": realesrgan},
|
||||
):
|
||||
first = cache.get_realesrgan("model", "cpu")
|
||||
second = cache.get_realesrgan("model", "cpu")
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, [("model", "cpu")])
|
||||
self.assertEqual(cache.snapshot()["realesrgan"], "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()
|
||||
@@ -220,7 +248,7 @@ class ServerTest(unittest.TestCase):
|
||||
"download-model",
|
||||
"download-required-models",
|
||||
"analyze-balloon",
|
||||
"extract-image-background",
|
||||
"refine-image-mask",
|
||||
"remove-background",
|
||||
"extract-frames",
|
||||
"generate-depth",
|
||||
@@ -228,6 +256,7 @@ class ServerTest(unittest.TestCase):
|
||||
"debug-da3",
|
||||
"debug-moge2",
|
||||
"debug-ben2",
|
||||
"debug-realesrgan-general-x4v3",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -401,16 +430,16 @@ class ServerTest(unittest.TestCase):
|
||||
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):
|
||||
def test_image_mask_canvas_clears_rgba_outside_the_selection(self):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
image_extract = sys.modules["server.jobs.image_extract"]
|
||||
image_mask = sys.modules["server.jobs.image_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)
|
||||
|
||||
canvas, _selection, bounds = image_extract.create_selection_canvas(
|
||||
canvas, _selection, bounds = image_mask.create_selection_canvas(
|
||||
source_path,
|
||||
((-1, -1), (2, -1), (-1, 2)),
|
||||
)
|
||||
@@ -421,25 +450,22 @@ class ServerTest(unittest.TestCase):
|
||||
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):
|
||||
def test_image_mask_crops_tightly_to_alpha(self):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
image_extract = sys.modules["server.jobs.image_extract"]
|
||||
image_mask = sys.modules["server.jobs.image_mask"]
|
||||
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(
|
||||
result, bounds = image_mask.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))
|
||||
self.assertEqual(bounds, (12, 21, 14, 23))
|
||||
self.assertEqual(result.shape, (2, 2, 4))
|
||||
self.assertTrue(np.all(result[:, :, 3] == 255))
|
||||
|
||||
def test_balloon_job_sends_the_masked_crop_to_ben2(self):
|
||||
import numpy as np
|
||||
|
||||
@@ -115,6 +115,7 @@ class WorkerTest(unittest.TestCase):
|
||||
cls.ben2 = load_server_module("models.ben2")
|
||||
cls.da3_model = load_server_module("models.da3")
|
||||
cls.moge2_model = load_server_module("models.moge2")
|
||||
cls.realesrgan_model = load_server_module("models.realesrgan")
|
||||
cls.surface_outputs = load_server_module("outputs.surface")
|
||||
cls.debug_da3_outputs = load_server_module("outputs.debug_da3")
|
||||
cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2")
|
||||
@@ -124,6 +125,9 @@ class WorkerTest(unittest.TestCase):
|
||||
cls.debug_da3_job = load_server_module("jobs.debug_da3")
|
||||
cls.debug_moge2_job = load_server_module("jobs.debug_moge2")
|
||||
cls.debug_ben2_job = load_server_module("jobs.debug_ben2")
|
||||
cls.debug_realesrgan_job = load_server_module(
|
||||
"jobs.debug_realesrgan"
|
||||
)
|
||||
cls.frames_job = load_server_module("jobs.frames")
|
||||
cls.foreground_job = load_server_module("jobs.foreground")
|
||||
cls.download = load_server_module("model_download")
|
||||
@@ -608,13 +612,24 @@ class WorkerTest(unittest.TestCase):
|
||||
"cpu",
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
raw_result, _load_ms, _inference_ms = self.ben2.infer_foreground(
|
||||
input_path,
|
||||
"model",
|
||||
"cpu",
|
||||
)
|
||||
|
||||
inference_pixels = received[0]
|
||||
result_pixels = np.asarray(result.convert("RGBA"), dtype=np.uint8)
|
||||
self.assertEqual(tuple(inference_pixels[0, 0]), (0, 0, 0))
|
||||
self.assertEqual(tuple(inference_pixels[0, 1]), (0, 128, 0))
|
||||
self.assertEqual(tuple(received[1][0, 0]), (0, 0, 0))
|
||||
self.assertEqual(tuple(received[1][0, 1]), (0, 128, 0))
|
||||
self.assertEqual(tuple(result_pixels[0, 0]), (255, 0, 0, 0))
|
||||
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),
|
||||
)
|
||||
|
||||
def test_ben2_background_removal_processes_each_frame(self):
|
||||
inference_options = []
|
||||
@@ -874,6 +889,39 @@ class WorkerTest(unittest.TestCase):
|
||||
(context.directory / "foreground.png").read_bytes(),
|
||||
b"foreground",
|
||||
)
|
||||
|
||||
def test_realesrgan_debug_saves_each_direct_api_prediction(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()
|
||||
context = FakeJobContext(root / "result", "debug-realesrgan")
|
||||
with patch.object(
|
||||
self.debug_realesrgan_job,
|
||||
"infer_one",
|
||||
return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
|
||||
) as infer_one:
|
||||
self.debug_realesrgan_job.run(
|
||||
context,
|
||||
{
|
||||
"input": str(frames),
|
||||
"model": "REALESRGAN_GENERAL_X4V3",
|
||||
"video_frames": 2,
|
||||
"device": "cpu",
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(infer_one.call_count, 2)
|
||||
predictions = sorted(context.directory.glob("prediction.*.png"))
|
||||
self.assertEqual(len(predictions), 2)
|
||||
with Image.open(predictions[0]) as prediction:
|
||||
self.assertEqual(prediction.size, (8, 8))
|
||||
self.assertEqual(prediction.mode, "RGB")
|
||||
|
||||
def test_da3_removes_background_after_depth_using_original_input(self):
|
||||
order = []
|
||||
inference_inputs = []
|
||||
|
||||
Reference in New Issue
Block a user