Add image selection refinement functionality with BEN2 integration
- Implemented viewport.py for handling 3D view projections, gestures, and overlay drawing. - Created refine_image_selection.py to manage the refinement of gesture masks and output image bounds. - Added selection module with refine_selection.py to process and crop images based on alpha channel visibility. - Introduced __init__.py for selection module to expose refine_selection function.
This commit is contained in:
@@ -10,12 +10,12 @@ import numpy as np
|
||||
SOURCE_ROOT = Path(__file__).parents[1] / "src"
|
||||
sys.path.insert(0, str(SOURCE_ROOT))
|
||||
|
||||
from anyimage.common.image_region import (
|
||||
RegionMask,
|
||||
RegionPath,
|
||||
rasterize_region_path,
|
||||
from anyimage.common.selection import (
|
||||
SelectionMask,
|
||||
SelectionPath,
|
||||
rasterize_selection_path,
|
||||
)
|
||||
from anyimage.common.depth_data import (
|
||||
from anyimage.common.depth import (
|
||||
depth_uniform_scale,
|
||||
load_depth_metadata,
|
||||
)
|
||||
@@ -25,11 +25,13 @@ from anyimage.operators.cutout_tool import geometry as cutout_geometry
|
||||
from anyimage.operators.cutout_tool import object as cutout_object
|
||||
from anyimage.operators.cutout_tool import shape as cutout_shape
|
||||
from anyimage.operators.cutout_tool.geometry import (
|
||||
RegionContours,
|
||||
extract_region_contours,
|
||||
SelectionContours,
|
||||
extract_selection_contours,
|
||||
poisson_balloon_profile,
|
||||
)
|
||||
from anyimage.server.models import moge2 as moge2_model
|
||||
from anyimage.server.models import da3 as da3_model
|
||||
from anyimage.server.models.geometry import GeometryFrame
|
||||
|
||||
|
||||
def test_cutout_mesh_detail_uses_the_cropped_analysis_scale():
|
||||
@@ -193,53 +195,102 @@ def test_depth_surface_projection_uses_one_scale_for_all_camera_axes():
|
||||
assert np.isclose(reference_depth, 5.8, atol=1e-6)
|
||||
|
||||
|
||||
def test_moge_surface_field_uses_the_model_validity_mask():
|
||||
prediction = moge2_model.Prediction(
|
||||
depth=np.ones((1, 2, 2), dtype=np.float32),
|
||||
normal=np.zeros((1, 2, 2, 3), dtype=np.float32),
|
||||
mask=np.asarray((((True, False), (True, True)),)),
|
||||
intrinsics=np.asarray(
|
||||
(((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),),
|
||||
def test_geometry_frame_requires_matching_float32_fields():
|
||||
with np.testing.assert_raises_regex(ValueError, "validity"):
|
||||
GeometryFrame(
|
||||
depth=np.ones((2, 2), dtype=np.float32),
|
||||
validity=np.ones((1, 2), dtype=np.float32),
|
||||
intrinsics=np.eye(3, dtype=np.float32),
|
||||
)
|
||||
with np.testing.assert_raises_regex(ValueError, "normal"):
|
||||
GeometryFrame(
|
||||
depth=np.ones((2, 2), dtype=np.float32),
|
||||
validity=np.ones((2, 2), dtype=np.float32),
|
||||
intrinsics=np.eye(3, dtype=np.float32),
|
||||
normal=np.ones((2, 2, 2), dtype=np.float32),
|
||||
)
|
||||
|
||||
|
||||
def test_moge_infer_returns_geometry_frames_with_model_validity():
|
||||
mask = np.asarray(((True, False), (True, True)))
|
||||
prediction = {
|
||||
"depth": np.ones((2, 2), dtype=np.float32),
|
||||
"normal": np.zeros((2, 2, 3), dtype=np.float32),
|
||||
"mask": mask,
|
||||
"intrinsics": np.asarray(
|
||||
((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),
|
||||
dtype=np.float32,
|
||||
),
|
||||
points=np.ones((1, 2, 2, 3), dtype=np.float32),
|
||||
"points": np.ones((2, 2, 3), dtype=np.float32),
|
||||
}
|
||||
|
||||
with patch.object(moge2_model, "infer_one", return_value=prediction):
|
||||
frames = moge2_model.infer(None, {}, ["image"], include_points=True)
|
||||
|
||||
assert len(frames) == 1
|
||||
frame = frames[0]
|
||||
assert np.array_equal(frame.validity, mask.astype(np.float32))
|
||||
assert np.array_equal(frame.normal, prediction["normal"])
|
||||
assert np.array_equal(frame.points, prediction["points"])
|
||||
assert np.array_equal(
|
||||
frame.intrinsics,
|
||||
np.asarray(((2.0, 0.0, 1.0), (0.0, 2.0, 1.0), (0.0, 0.0, 1.0))),
|
||||
)
|
||||
|
||||
frame = moge2_model.geometry_frame(prediction)
|
||||
|
||||
assert np.array_equal(frame.valid_mask, prediction.mask[0])
|
||||
|
||||
|
||||
def test_moge_z_depth_frame_does_not_require_points():
|
||||
prediction = moge2_model.Prediction(
|
||||
depth=np.ones((1, 2, 2), dtype=np.float32),
|
||||
normal=np.zeros((1, 2, 2, 3), dtype=np.float32),
|
||||
mask=np.ones((1, 2, 2), dtype=bool),
|
||||
intrinsics=np.eye(3, dtype=np.float32)[None],
|
||||
points=None,
|
||||
)
|
||||
prediction = {
|
||||
"depth": np.ones((2, 2), dtype=np.float32),
|
||||
"normal": np.zeros((2, 2, 3), dtype=np.float32),
|
||||
"mask": np.ones((2, 2), dtype=bool),
|
||||
"intrinsics": np.eye(3, dtype=np.float32),
|
||||
"points": np.ones((2, 2, 3), dtype=np.float32),
|
||||
}
|
||||
|
||||
frame = moge2_model.geometry_frame(prediction)
|
||||
with patch.object(moge2_model, "infer_one", return_value=prediction):
|
||||
frame = moge2_model.infer(None, {}, ["image"], include_points=False)[0]
|
||||
|
||||
assert frame.points is None
|
||||
assert frame.valid_mask.all()
|
||||
assert frame.validity.all()
|
||||
|
||||
|
||||
def test_region_contours_group_holes_with_their_connected_component():
|
||||
def test_da3_geometry_frames_use_confidence_as_validity_and_derive_points():
|
||||
prediction = SimpleNamespace(
|
||||
depth=np.ones((2, 2, 3), dtype=np.float32),
|
||||
conf=np.asarray(
|
||||
(
|
||||
((0.25, 0.5, 0.75), (1.0, 0.5, 0.25)),
|
||||
((0.1, 0.2, 0.3), (0.4, 0.5, 0.6)),
|
||||
),
|
||||
dtype=np.float32,
|
||||
),
|
||||
intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
|
||||
)
|
||||
|
||||
frames = da3_model.geometry_frames(prediction)
|
||||
|
||||
assert len(frames) == 2
|
||||
assert frames[0].normal is None
|
||||
assert np.array_equal(frames[1].validity, prediction.conf[1])
|
||||
assert frames[0].points.shape == (2, 3, 3)
|
||||
assert frames[0].points.dtype == np.float32
|
||||
|
||||
|
||||
def test_selection_contours_group_holes_with_their_connected_component():
|
||||
alpha = np.zeros((120, 200), dtype=np.uint8)
|
||||
alpha[20:100, 50:160] = 255
|
||||
alpha[45:75, 80:120] = 0
|
||||
alpha[5:15, 5:15] = 255
|
||||
|
||||
region_contours = extract_region_contours(
|
||||
RegionMask(alpha, (0, 0, 200, 120)),
|
||||
selection_contours = extract_selection_contours(
|
||||
SelectionMask(alpha, (0, 0, 200, 120)),
|
||||
max_resolution=200,
|
||||
)
|
||||
|
||||
assert region_contours.image_size == (200, 120)
|
||||
assert sorted(len(component) for component in region_contours.components) == [1, 2]
|
||||
assert selection_contours.image_size == (200, 120)
|
||||
assert sorted(len(component) for component in selection_contours.components) == [1, 2]
|
||||
contours = [
|
||||
contour for component in region_contours.components for contour in component
|
||||
contour for component in selection_contours.components for contour in component
|
||||
]
|
||||
bounds = [
|
||||
(*np.min(contour, axis=0), *np.max(contour, axis=0)) for contour in contours
|
||||
@@ -249,19 +300,19 @@ def test_region_contours_group_holes_with_their_connected_component():
|
||||
assert (5.0, 5.0, 15.0, 15.0) in bounds
|
||||
|
||||
|
||||
def test_region_contours_accept_blender_float_mask_with_a_hole():
|
||||
def test_selection_contours_accept_blender_float_mask_with_a_hole():
|
||||
alpha = np.zeros((120, 200), dtype=np.float32)
|
||||
alpha[20:100, 50:160] = 1.0
|
||||
alpha[45:75, 80:120] = 0.0
|
||||
|
||||
region_contours = extract_region_contours(
|
||||
RegionMask(alpha, (0, 0, 200, 120)),
|
||||
selection_contours = extract_selection_contours(
|
||||
SelectionMask(alpha, (0, 0, 200, 120)),
|
||||
max_resolution=200,
|
||||
)
|
||||
|
||||
assert region_contours.image_size == (200, 120)
|
||||
assert len(region_contours.components) == 1
|
||||
assert len(region_contours.components[0]) == 2
|
||||
assert selection_contours.image_size == (200, 120)
|
||||
assert len(selection_contours.components) == 1
|
||||
assert len(selection_contours.components[0]) == 2
|
||||
|
||||
|
||||
def test_cutout_crop_bounds_map_to_the_source_plane():
|
||||
@@ -285,14 +336,14 @@ def test_depth_surface_cutout_uses_half_alpha_boundary():
|
||||
alpha[2:10, 2:14] = 0.49
|
||||
alpha[4:8, 5:11] = 0.51
|
||||
|
||||
region_contours = cutout_geometry.extract_cutout_contours(
|
||||
RegionMask(alpha, (0, 0, 16, 12)),
|
||||
selection_contours = cutout_geometry.extract_cutout_contours(
|
||||
SelectionMask(alpha, (0, 0, 16, 12)),
|
||||
64,
|
||||
threshold=0.5,
|
||||
)
|
||||
contour = np.asarray(region_contours.components[0][0])
|
||||
contour = np.asarray(selection_contours.components[0][0])
|
||||
|
||||
assert region_contours.image_size == (16, 12)
|
||||
assert selection_contours.image_size == (16, 12)
|
||||
np.testing.assert_allclose(contour.min(axis=0), (5.0, 4.0))
|
||||
np.testing.assert_allclose(contour.max(axis=0), (11.0, 8.0))
|
||||
|
||||
@@ -306,37 +357,37 @@ def test_blender_cutout_alpha_pipeline_does_not_import_pillow():
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=without_pillow):
|
||||
region_mask = rasterize_region_path(
|
||||
selection_mask = rasterize_selection_path(
|
||||
(48, 32),
|
||||
RegionPath(points=((6, 4), (42, 4), (42, 28), (6, 28))),
|
||||
SelectionPath(points=((6, 4), (42, 4), (42, 28), (6, 28))),
|
||||
antialias=True,
|
||||
)
|
||||
antialiased = region_mask.full_values((48, 32))
|
||||
region_contours = extract_region_contours(
|
||||
RegionMask(antialiased, (0, 0, 48, 32)),
|
||||
antialiased = selection_mask.full_values((48, 32))
|
||||
selection_contours = extract_selection_contours(
|
||||
SelectionMask(antialiased, (0, 0, 48, 32)),
|
||||
max_resolution=24,
|
||||
)
|
||||
|
||||
assert region_contours.image_size == (48, 32)
|
||||
assert region_contours.components
|
||||
assert selection_contours.image_size == (48, 32)
|
||||
assert selection_contours.components
|
||||
assert np.any((antialiased > 0.0) & (antialiased < 1.0))
|
||||
|
||||
|
||||
def test_blender_cutout_contour_resolution_tracks_mesh_detail():
|
||||
expected = RegionContours(
|
||||
expected = SelectionContours(
|
||||
components=(((((0, 0), (1, 0), (0, 1))),),),
|
||||
image_size=(1, 1),
|
||||
)
|
||||
region_mask = RegionMask(np.ones((2, 2), dtype=np.float32), (0, 0, 2, 2))
|
||||
with patch.object(cutout_geometry, "extract_region_contours", return_value=expected) as trace:
|
||||
assert cutout_geometry.extract_cutout_contours(region_mask, 32) == expected
|
||||
assert cutout_geometry.extract_cutout_contours(region_mask, 256) == expected
|
||||
selection_mask = SelectionMask(np.ones((2, 2), dtype=np.float32), (0, 0, 2, 2))
|
||||
with patch.object(cutout_geometry, "extract_selection_contours", return_value=expected) as trace:
|
||||
assert cutout_geometry.extract_cutout_contours(selection_mask, 32) == expected
|
||||
assert cutout_geometry.extract_cutout_contours(selection_mask, 256) == expected
|
||||
|
||||
assert trace.call_args_list[0].kwargs["max_resolution"] == 512
|
||||
assert trace.call_args_list[1].kwargs["max_resolution"] == 1024
|
||||
|
||||
|
||||
def test_cutout_point_spacing_uses_the_complete_source_plane():
|
||||
def test_cutout_sampling_spacing_uses_the_complete_source_plane():
|
||||
source = SimpleNamespace(matrix_world=object())
|
||||
with (
|
||||
patch.object(cutout_geometry, "image_empty_bounds", return_value=(-2.0, 2.0, -1.0, 1.0)),
|
||||
@@ -346,7 +397,7 @@ def test_cutout_point_spacing_uses_the_complete_source_plane():
|
||||
return_value=(((1.0, 0.0), (0.0, 1.0)), 1.0),
|
||||
),
|
||||
):
|
||||
spacing = cutout_geometry.cutout_point_spacing(source, 200)
|
||||
spacing = cutout_geometry.cutout_sampling_spacing(source, 200)
|
||||
|
||||
assert spacing == 0.02
|
||||
|
||||
@@ -401,7 +452,7 @@ def test_poisson_balloon_profile_scales_in_world_units():
|
||||
assert np.isclose(scaled[4], original[4] * 2.0, atol=1e-6)
|
||||
|
||||
|
||||
def test_poisson_balloon_profile_makes_wide_regions_thicker():
|
||||
def test_poisson_balloon_profile_makes_wide_selections_thicker():
|
||||
faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
|
||||
boundary = np.asarray((True, True, True, True, False))
|
||||
wide = np.asarray(
|
||||
@@ -445,11 +496,11 @@ def test_teddy_balloon_profile_uses_chordal_axis_radius():
|
||||
assert np.isclose(profile[4], 0.45, atol=1e-6)
|
||||
|
||||
|
||||
def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_region():
|
||||
regions = ((((2046, 2046), (2050, 2046), (2050, 2050), (2046, 2050)),),)
|
||||
def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_selection():
|
||||
components = ((((2046, 2046), (2050, 2046), (2050, 2050), (2046, 2050)),),)
|
||||
|
||||
domain = cutout_geometry.cutout_domain(
|
||||
RegionContours(regions, (4096, 4096)),
|
||||
SelectionContours(components, (4096, 4096)),
|
||||
(-1.0, 1.0, -1.0, 1.0),
|
||||
0.5,
|
||||
((1.0, 0.0), (0.0, 1.0)),
|
||||
@@ -460,10 +511,10 @@ def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_region():
|
||||
|
||||
|
||||
def test_cutout_domain_builds_a_full_image_at_minimum_detail():
|
||||
regions = ((((0, 0), (4096, 0), (4096, 4096), (0, 4096)),),)
|
||||
components = ((((0, 0), (4096, 0), (4096, 4096), (0, 4096)),),)
|
||||
|
||||
domain = cutout_geometry.cutout_domain(
|
||||
RegionContours(regions, (4096, 4096)),
|
||||
SelectionContours(components, (4096, 4096)),
|
||||
(-1.0, 1.0, -1.0, 1.0),
|
||||
0.5,
|
||||
((1.0, 0.0), (0.0, 1.0)),
|
||||
@@ -474,10 +525,10 @@ def test_cutout_domain_builds_a_full_image_at_minimum_detail():
|
||||
|
||||
|
||||
def test_cutout_domain_centers_vertices_on_crop_bounds():
|
||||
regions = ((((0, 0), (100, 0), (100, 50), (0, 50)),),)
|
||||
components = ((((0, 0), (100, 0), (100, 50), (0, 50)),),)
|
||||
|
||||
vertices, _faces, uv, _boundary, _planar, _profile = cutout_geometry.cutout_domain(
|
||||
RegionContours(regions, (100, 50)),
|
||||
SelectionContours(components, (100, 50)),
|
||||
(2.0, 6.0, -3.0, 1.0),
|
||||
0.5,
|
||||
((1.0, 0.0), (0.0, 1.0)),
|
||||
@@ -489,8 +540,8 @@ def test_cutout_domain_centers_vertices_on_crop_bounds():
|
||||
assert np.allclose(uv, expected_uv, atol=1e-6)
|
||||
|
||||
|
||||
def test_cutout_domain_keeps_overlapping_regions_disconnected():
|
||||
regions = (
|
||||
def test_cutout_domain_keeps_overlapping_selection_components_disconnected():
|
||||
components = (
|
||||
(
|
||||
((20, 0), (100, 0), (100, 80), (20, 80)),
|
||||
((50, 20), (50, 60), (70, 60), (70, 20)),
|
||||
@@ -498,7 +549,7 @@ def test_cutout_domain_keeps_overlapping_regions_disconnected():
|
||||
(((0, 20), (40, 20), (40, 50), (0, 50)),),
|
||||
)
|
||||
domain = cutout_geometry.cutout_domain(
|
||||
RegionContours(regions, (120, 80)),
|
||||
SelectionContours(components, (120, 80)),
|
||||
(0.0, 1.0, 0.0, 1.0),
|
||||
0.025,
|
||||
((1.0, 0.0), (0.0, 1.0)),
|
||||
@@ -511,7 +562,7 @@ def test_cutout_domain_keeps_overlapping_regions_disconnected():
|
||||
& (centers[:, 1] > 0.35)
|
||||
& (centers[:, 1] < 0.65)
|
||||
)
|
||||
in_left_region = (
|
||||
in_left_component = (
|
||||
(centers[:, 0] > 0.02)
|
||||
& (centers[:, 0] < 0.12)
|
||||
& (centers[:, 1] > 0.4)
|
||||
@@ -519,17 +570,17 @@ def test_cutout_domain_keeps_overlapping_regions_disconnected():
|
||||
)
|
||||
|
||||
assert not in_hole.any()
|
||||
assert in_left_region.any()
|
||||
assert in_left_component.any()
|
||||
assert mesh_component_count(domain[1]) == 2
|
||||
|
||||
|
||||
def test_region_calibration_uses_the_final_region_mask():
|
||||
region_mask = RegionMask(
|
||||
def test_selection_calibration_uses_the_final_selection_mask():
|
||||
selection_mask = SelectionMask(
|
||||
np.asarray(((1.0, 1.0), (1.0, 0.0)), dtype=np.float32),
|
||||
(0, 0, 2, 2),
|
||||
)
|
||||
|
||||
valid, transparent_fraction = cutout_geometry.region_calibration(region_mask, (2, 2))
|
||||
valid, transparent_fraction = cutout_geometry.selection_calibration(selection_mask, (2, 2))
|
||||
|
||||
assert np.array_equal(valid, ((True, True), (True, False)))
|
||||
assert transparent_fraction == 0.25
|
||||
|
||||
Reference in New Issue
Block a user