icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor image processing and cutout functionalities

Browse Source 
- Updated the naming convention for upscaled images to include underscores and file extensions.
- Renamed properties related to cutout pixel size to point spacing for clarity.
- Adjusted alpha region calculations to use a half-alpha threshold for better accuracy.
- Enhanced the handling of foreground mesh masks to consider half-alpha values.
- Introduced new functions for unbounded selections and cutout images, ensuring proper handling of source and selection images.
- Improved test coverage for cutout functionalities, including bounds mapping and alpha preservation.
- Updated tests to reflect changes in naming conventions and functionality adjustments.
This commit is contained in:
马南
2026-08-27 12:01:47 +08:00
parent 2b79a342b0
commit 6c8358f92f
27 changed files with 1058 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

133
tests/test_server.py
View File

@@ -562,7 +562,85 @@ class ServerTest(unittest.TestCase):
self.assertEqual(tuple(result[1, 1]), (20, 40, 60, 255))
self.assertEqual(tuple(result[5, 7]), (0, 0, 0, 0))
def test_mesh_job_keeps_original_ai_input_separate_from_geometry_mask(self):
def test_mesh_cutout_images_crop_color_and_preserve_source_alpha(self):
import json
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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
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)
selection[1, 2] = (200, 20, 40, 0)
Image.fromarray(selection, mode="RGBA").save(selection_path)
color_path, mask_path, analysis_path, bounds_path = (
mesh_job._write_cutout_images(
source_path,
selection_path,
output_directory,
)
)
color = np.asarray(Image.open(color_path).convert("RGBA"))
mask = np.asarray(Image.open(mask_path).convert("RGBA"))
analysis = np.asarray(Image.open(analysis_path).convert("RGBA"))
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(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))
def test_mesh_unrefined_cutout_keeps_selection_outside_the_source(self):
import json
import numpy as np
from PIL import Image
mesh_job = sys.modules['server.jobs.mesh']
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'
output_directory.mkdir()
Image.new('RGBA', (8, 6), (10, 80, 160, 255)).save(source_path)
polygon = ((-2, 1), (10, 1), (10, 5), (-2, 5))
_selection_path, selection_bounds = (
mesh_job._write_unbounded_selection(
source_path,
polygon,
full_selection_path,
)
)
color_path, mask_path, analysis_path, bounds_path = (
mesh_job._write_cutout_images(
source_path,
full_selection_path,
output_directory,
selection_bounds,
)
)
color = np.asarray(Image.open(color_path).convert('RGBA'))
mask = np.asarray(Image.open(mask_path).convert('RGBA'))
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.assertGreater(int(mask[1, 0, 3]), 127)
self.assertEqual(tuple(analysis[1, 0]), (0, 0, 0, 0))
def test_mesh_job_keeps_unrefined_outside_selection_with_normal(self):
import json
import numpy as np
from PIL import Image
@@ -577,7 +655,7 @@ class ServerTest(unittest.TestCase):
check_cancelled=lambda: None,
resource=lambda _name: object(),
)
normal = np.zeros((1, 6, 8, 3), dtype=np.float32)
normal = np.zeros((1, 4, 12, 3), dtype=np.float32)
normal[..., 0] = 1.0
prediction = SimpleNamespace(normal=normal)
received = []
@@ -594,7 +672,7 @@ class ServerTest(unittest.TestCase):
{
"input": str(source_path),
"shape": "BALLOON",
"polygon": ((2, 1), (7, 1), (7, 5), (2, 5)),
"polygon": ((-2, 1), (10, 1), (10, 5), (-2, 5)),
"mesh_detail": 64,
"invert": False,
"device": "cpu",
@@ -609,16 +687,36 @@ class ServerTest(unittest.TestCase):
normal_pixels = np.asarray(
Image.open(job_directory / result["normal"]).convert("RGB")
)
self.assertEqual(result, {"normal": "normal.png"})
self.assertEqual(
result,
{
"selection": "selection.png",
"bounds": "bounds.json",
"normal": "normal.png",
},
)
self.assertTrue(color_path.is_file())
self.assertFalse((job_directory / "contour.json").exists())
self.assertEqual(received[0]["model"], "MOGE2_VITS_NORMAL")
self.assertEqual(received_images, [source_path])
self.assertEqual(tuple(color_pixels[3, 4, :3]), (20, 40, 60))
self.assertGreater(int(color_pixels[3, 4, 3]), 240)
self.assertLess(int(color_pixels[0, 0, 3]), 5)
self.assertEqual(tuple(normal_pixels[0, 0]), (255, 127, 127))
self.assertEqual(tuple(normal_pixels[3, 4]), (255, 127, 127))
self.assertEqual(received_images, [job_directory / "analysis.png"])
analysis_pixels = np.asarray(
Image.open(received_images[0]).convert("RGBA")
)
self.assertEqual(
json.loads((job_directory / "bounds.json").read_text()),
[-2, 1, 10, 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))
def test_mesh_job_skips_moge_normal_when_option_is_disabled(self):
from PIL import Image
@@ -648,17 +746,22 @@ class ServerTest(unittest.TestCase):
},
)
self.assertEqual(result, {})
self.assertEqual(
result,
{"selection": "selection.png", "bounds": "bounds.json"},
)
write_normal.assert_not_called()
self.assertFalse((job_directory / "normal.png").exists())
def test_mesh_depth_balloon_analyzes_full_source_then_applies_selection(self):
def test_mesh_depth_balloon_uses_cropped_analysis_alpha(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
source_path = self.storage_root / "mesh-depth-balloon-source.png"
Image.new("RGBA", (6, 6), (20, 40, 60, 255)).save(source_path)
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)
job_directory = self.storage_root / "mesh-depth-balloon-job"
job_directory.mkdir()
context = SimpleNamespace(
@@ -684,8 +787,6 @@ class ServerTest(unittest.TestCase):
context,
{"device": "cpu"},
source_path,
((2, 2), (10, 2), (2, 10)),
selection_size=(12, 12),
include_normal=True,
)
@@ -696,7 +797,7 @@ class ServerTest(unittest.TestCase):
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]), (255, 127, 127))
self.assertEqual(tuple(normal[0, 0]), (128, 128, 255))
self.assertEqual(tuple(normal[2, 2]), (255, 127, 127))
def test_mesh_job_uses_upscale_source_only_for_analysis(self):