from types import SimpleNamespace from unittest.mock import patch import numpy as np import pytest from PIL import Image from server.geometry import prediction_artifacts as artifacts from server.geometry.edge_depth import extend_edge_depth from server.models.geometry import GeometryFrame from server.models.onnx_moge2 import depth_to_points def make_frame(depth): depth = np.asarray(depth, dtype=np.float32) height, width = depth.shape normalized = np.array([[1.2, 0, 0.4], [0, 1.4, 0.6], [0, 0, 1]], np.float32) intrinsics = normalized.copy() intrinsics[0] *= width intrinsics[1] *= height return GeometryFrame( depth=depth, validity=np.ones_like(depth), intrinsics=intrinsics, points=depth_to_points(depth, normalized), normal=np.zeros((*depth.shape, 3), np.float32), ) def silhouette(): alpha = np.zeros((25, 35), np.float32) alpha[4:21, 4:22] = 1 alpha[4:21, 22:25] = 0.2 depth = np.full(alpha.shape, 20, np.float32) depth[alpha == 1] = 2 return make_frame(depth), alpha def test_hair_extension_preserves_core_alpha_and_projection(): frame, alpha = silhouette() original_alpha = alpha.copy() result = extend_edge_depth(frame, alpha) assert np.all(result.depth[alpha >= 0.1] == 2) np.testing.assert_array_equal(result.points[9:16, 9:16], frame.points[9:16, 9:16]) np.testing.assert_array_equal(alpha, original_alpha) np.testing.assert_array_equal(frame.depth[alpha == 0.2], 20) np.testing.assert_allclose( result.points[..., :2] / result.points[..., 2:], frame.points[..., :2] / frame.points[..., 2:], atol=1e-6, ) assert result.validity is frame.validity assert result.normal is frame.normal assert np.isfinite(result.points).all() def test_sampling_padding_is_bounded(): frame, alpha = silhouette() result = extend_edge_depth(frame, alpha, alpha_threshold=0.1) assert np.all(result.depth[12, 24:29] == 2) assert result.depth[12, 29] == 20 assert (result.depth[12, 24] + result.depth[12, 25]) / 2 == 2 def test_edge_padding_tracks_selected_threshold(): frame, alpha = silhouette() loose = extend_edge_depth(frame, alpha, alpha_threshold=0.1) strict = extend_edge_depth(frame, alpha) assert loose.depth[12, 28] == 2 assert strict.depth[12, 28] == 20 def test_strict_threshold_excludes_lower_alpha_donors(): alpha = np.zeros((10, 10), np.float32) alpha[3:7, 3:7] = 0.96 alpha[5, 5] = 1 depth = np.full(alpha.shape, 10, np.float32) depth[5, 5] = 2 result = extend_edge_depth(make_frame(depth), alpha, alpha_threshold=1) assert result.depth[5, 6] == 2 def test_image_rectangle_does_not_erode_subject(): alpha = np.ones((15, 20), np.float32) alpha[:, 15:] = 0 depth = np.full(alpha.shape, 2, np.float32) depth[0, 0] = 3 frame = make_frame(depth) result = extend_edge_depth(frame, alpha) assert result.depth[0, 0] == 3 np.testing.assert_array_equal(result.points[0, 0], frame.points[0, 0]) def test_small_components_fall_back_without_borrowing_other_depths(): alpha = np.zeros((15, 30), np.float32) alpha[5, 5:7] = [0.95, 0.1] alpha[5, 10:12] = [1, 0.2] alpha[5, 14] = 0.94 depth = np.full(alpha.shape, 30, np.float32) depth[5, 5] = 2 depth[5, 10] = 8 result = extend_edge_depth(make_frame(depth), alpha) assert result.depth[5, 6] == 2 assert result.depth[5, 11] == 8 assert result.depth[5, 14] == 30 assert result.depth[5, 7] == 2 assert result.depth[5, 9] == 8 def test_diagonal_pixels_share_a_component(): alpha = np.zeros((10, 10), np.float32) alpha[4, 4] = 1 alpha[5, 5] = 0.2 depth = np.full(alpha.shape, 10, np.float32) depth[4, 4] = 2 result = extend_edge_depth(make_frame(depth), alpha) assert result.depth[5, 5] == 2 @pytest.mark.parametrize("validity", [0.0, 0.2, 0.5]) def test_invalid_donors_are_excluded(validity): frame, alpha = silhouette() frame.depth[10, 6] = 100 frame.validity[10, 6] = validity result = extend_edge_depth(frame, alpha) assert result.depth[10, 4] == 2 assert result.depth[10, 6] == 100 assert result.validity[10, 6] == np.float32(validity) def test_internal_mask_hole_does_not_expand_alpha_repair(): frame, alpha = silhouette() frame.depth[9:14, 10:15] = 100 frame.validity[9:14, 10:15] = 0 original = frame.depth.copy() result = extend_edge_depth(frame, alpha) np.testing.assert_array_equal(result.depth[6:19, 6:20], original[6:19, 6:20]) assert result.depth[12, 24] == 2 def test_opaque_image_keeps_low_mask_depth(): frame = make_frame(np.ones((10, 10))) frame.depth[5, 5] = 100 frame.validity[5, 5] = 0 assert extend_edge_depth(frame, np.ones((10, 10))) is frame def test_component_without_valid_donors_keeps_its_depth(): frame, alpha = silhouette() frame.validity[alpha > 0] = 0 assert extend_edge_depth(frame, alpha) is frame def test_local_smoothing_reduces_donor_step_and_preserves_core(): frame, alpha = silhouette() frame.depth[12:, alpha.shape[1] // 4:22] = 4 result = extend_edge_depth(frame, alpha) assert 2 < result.depth[11, 23] < 4 assert 2 < result.depth[12, 23] < 4 np.testing.assert_array_equal(result.depth[8:17, 8:17], frame.depth[8:17, 8:17]) @pytest.mark.parametrize("alpha_value", [0, 0.3, 1]) def test_images_without_repair_return_original_frame(alpha_value): frame = make_frame(np.ones((5, 8))) assert extend_edge_depth(frame, np.full((5, 8), alpha_value)) is frame def test_alpha_resize_retains_threshold_values(tmp_path): rgba = np.zeros((1, 4, 4), np.uint8) rgba[0, :, 3] = [25, 26, 242, 243] path = tmp_path / "alpha.png" Image.fromarray(rgba).save(path) alpha = artifacts._image_alpha(path, (2, 8)) np.testing.assert_array_equal( alpha >= 0.1, [[False] * 2 + [True] * 6] * 2 ) np.testing.assert_array_equal( alpha >= 0.95, [[False] * 6 + [True] * 2] * 2 ) @pytest.mark.parametrize( "generate_depth,normal_mode", [ (True, "OBJECT"), (True, "TANGENT"), (True, "NONE"), (False, "OBJECT"), ], ) def test_artifacts_use_corrected_depth_and_original_normals( tmp_path, generate_depth, normal_mode ): frame, alpha = silhouette() rgba = np.full((*alpha.shape, 4), 255, np.uint8) rgba[..., 3] = np.round(alpha * 255).astype(np.uint8) path = tmp_path / "source.png" Image.fromarray(rgba).save(path) context = SimpleNamespace( directory=tmp_path, resource=lambda _: None, progress=lambda *_: None, check_cancelled=lambda: None, ) with ( patch.object(artifacts.moge, "infer", return_value=[frame]), patch.object( artifacts, "write_depth_texture", return_value=tmp_path / "depth.exr", ) as depth, patch.object( artifacts, "write_depth_metadata", return_value=tmp_path / "depth.json" ) as metadata, patch.object( artifacts, "_write_normal", return_value=tmp_path / "normal.png" ) as normal, patch.object(artifacts, "extend_edge_depth", wraps=extend_edge_depth) as repair, ): artifacts.generate_moge_artifacts( context, {}, path, generate_depth=generate_depth, normal_mode=normal_mode, ) if normal_mode != "NONE": assert normal.call_args.args[0] is frame else: normal.assert_not_called() assert repair.call_count == int(generate_depth) if generate_depth: written_frame = depth.call_args.args[0] np.testing.assert_array_equal(depth.call_args.kwargs["alpha"], alpha) assert metadata.call_args.args[0] is written_frame assert written_frame.depth[12, 24] == 2 assert written_frame.validity is frame.validity else: metadata.assert_not_called() depth.assert_not_called()