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

Add image selection refinement functionality with BEN2 integration

Browse Source 
- 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:
马南
2026-08-31 22:10:05 +08:00
parent 14d63001d2
commit c97b4a3e41
76 changed files with 2127 additions and 1137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
src/anyimage/server/geometry/depth_texture.py
View File

@@ -71,31 +71,33 @@ def _write_float_exr(rgba, output_path):
return output_path
def write_vector_depth_texture(points, valid, output_path):
"""Write camera-space XYZ and validity as a float RGBA texture."""
def write_vector_depth_texture(frame, output_path):
"""Write a frame's camera-space XYZ and validity as float RGBA."""
import numpy as np
points = np.asarray(points, dtype=np.float32)
valid = np.asarray(valid, dtype=bool)
if points.ndim != 3 or points.shape[2] != 3 or valid.shape != points.shape[:2]:
raise ValueError("Vector depth dimensions do not match")
rgba = np.zeros((*valid.shape, 4), dtype=np.float32)
rgba[valid, :3] = points[valid]
rgba[..., 3] = valid
if frame.points is None:
raise ValueError("Vector depth requires a point field")
points = np.asarray(frame.points, dtype=np.float32)
validity = np.asarray(frame.validity, dtype=np.float32)
if not np.isfinite(points).all():
raise ValueError("Point field contains non-finite values")
rgba = np.empty((*validity.shape, 4), dtype=np.float32)
rgba[..., :3] = points
rgba[..., 3] = validity
return _write_float_exr(rgba, output_path)
def write_z_depth_texture(depth, valid, output_path):
"""Write camera Z and validity as a float RGBA texture."""
def write_z_depth_texture(frame, output_path):
"""Write a frame's camera Z and validity as float RGBA."""
import numpy as np
depth = np.asarray(depth, dtype=np.float32)
valid = np.asarray(valid, dtype=bool)
if depth.ndim != 2 or valid.shape != depth.shape:
raise ValueError("Z depth dimensions do not match")
rgba = np.zeros((*depth.shape, 4), dtype=np.float32)
rgba[..., :3] = np.where(valid, depth, 0.0)[..., None]
rgba[..., 3] = valid
depth = np.asarray(frame.depth, dtype=np.float32)
validity = np.asarray(frame.validity, dtype=np.float32)
if not np.isfinite(depth).all():
raise ValueError("Depth field contains non-finite values")
rgba = np.empty((*depth.shape, 4), dtype=np.float32)
rgba[..., :3] = depth[..., None]
rgba[..., 3] = validity
return _write_float_exr(rgba, output_path)
@@ -104,15 +106,19 @@ def write_depth_metadata(frame, reference_valid, output_path):
import numpy as np
depth = np.asarray(frame.depth, dtype=np.float32)
valid = np.asarray(frame.valid_mask, dtype=bool)
validity = np.asarray(frame.validity, dtype=np.float32)
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
reference_valid = np.asarray(reference_valid, dtype=bool) & valid
if depth.ndim != 2 or valid.shape != depth.shape:
reference_valid = np.asarray(reference_valid, dtype=bool)
if depth.ndim != 2 or validity.shape != depth.shape:
raise ValueError("Depth metadata dimensions do not match")
if reference_valid.shape != depth.shape:
raise ValueError("Depth metadata reference mask dimensions do not match")
if intrinsics.shape != (3, 3):
raise ValueError("Depth metadata requires 3x3 camera intrinsics")
valid = (validity > 0.0) & np.isfinite(depth) & (depth > 0.0)
reference_valid &= valid
if not valid.any():
raise ValueError("MoGe-2 returned no valid depth pixels")
raise ValueError("Geometry frame contains no valid depth pixels")
if not reference_valid.any():
raise ValueError("The image contains no valid depth pixels")
metadata = {

30
src/anyimage/server/geometry/normal_texture.py
View File

@@ -3,14 +3,14 @@
from pathlib import Path
def resized_valid_mask(valid_mask, shape):
def resized_visibility(visibility, shape):
import numpy as np
from PIL import Image
valid_mask = np.asarray(valid_mask, dtype=bool)
if valid_mask.shape == tuple(shape):
return valid_mask
resized = Image.fromarray(valid_mask.astype(np.uint8) * 255, mode="L")
visibility = np.asarray(visibility, dtype=bool)
if visibility.shape == tuple(shape):
return visibility
resized = Image.fromarray(visibility.astype(np.uint8) * 255, mode="L")
resized = resized.resize(
(int(shape[1]), int(shape[0])),
Image.Resampling.NEAREST,
@@ -19,8 +19,8 @@ def resized_valid_mask(valid_mask, shape):
def write_camera_normal_texture(
normal,
valid,
frame,
visibility,
origin,
crop_size,
source_size,
@@ -31,12 +31,14 @@ def write_camera_normal_texture(
import numpy as np
from PIL import Image
normal = np.asarray(normal, dtype=np.float32)
valid = np.asarray(valid, dtype=bool)
if normal.ndim != 3 or normal.shape[2] != 3:
raise ValueError("Depth Balloon normal data must have three components")
if valid.shape != normal.shape[:2]:
valid = resized_valid_mask(valid, normal.shape[:2])
if frame.normal is None:
raise ValueError("Normal texture requires a normal field")
normal = np.asarray(frame.normal, dtype=np.float32)
visibility = np.asarray(visibility, dtype=bool)
if not np.isfinite(normal).all():
raise ValueError("Normal field contains non-finite values")
if visibility.shape != normal.shape[:2]:
visibility = resized_visibility(visibility, normal.shape[:2])
normal_mode = str(normal_mode).upper()
if normal_mode == "TANGENT":
converted = normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32)
@@ -55,7 +57,7 @@ def write_camera_normal_texture(
0.0,
255.0,
).astype(np.uint8)
encoded[~valid] = empty_normal
encoded[~visibility] = empty_normal
crop_width, crop_height = (int(value) for value in crop_size)
crop = Image.fromarray(encoded, mode="RGB")
if crop.size != (crop_width, crop_height):

34
src/anyimage/server/geometry/prediction_artifacts.py
View File

@@ -14,10 +14,10 @@ def _limited_model_input(source_path, output_path, max_input_size):
from PIL import Image
with Image.open(source_path) as opened:
limited = limit_image(opened, max_input_size)
if limited is opened:
return source_path
limited.save(output_path, format="PNG", compress_level=1)
limited = limit_image(opened.convert("RGBA"), max_input_size)
background = Image.new("RGBA", limited.size, (0, 0, 0, 255))
normalized = Image.alpha_composite(background, limited).convert("RGB")
normalized.save(output_path, format="PNG", compress_level=1)
return output_path
@@ -34,16 +34,14 @@ def _alpha_mask(image_path, shape):
return np.asarray(alpha, dtype=np.uint8) >= 26
def _write_normal(prediction, color_path, output_path, normal_mode):
import numpy as np
def _write_normal(frame, color_path, output_path, normal_mode):
from PIL import Image
normal = np.asarray(prediction.normal[0], dtype=np.float32)
valid = _alpha_mask(color_path, normal.shape[:2])
valid = _alpha_mask(color_path, frame.depth.shape)
with Image.open(color_path) as source:
output_size = source.size
return write_camera_normal_texture(
normal,
frame,
valid,
(0, 0),
output_size,
@@ -74,36 +72,34 @@ def generate_moge2_artifacts(
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
)
try:
prediction = moge2.infer(
frames = moge2.infer(
context.resource("model_manager"),
parameters,
[model_path],
include_points=depth_mode == "VECTOR",
cancel_check=context.check_cancelled,
)
if len(frames) != 1:
raise ValueError("Artifact generation requires one geometry frame")
frame = frames[0]
result = {}
if depth_mode != "NONE":
context.progress(0.45, "Generating depth texture")
frame = moge2.geometry_frame(prediction)
if depth_mode == "VECTOR":
if frame.points is None:
raise ValueError("Vector depth requires a MoGe-2 point field")
depth_path = write_vector_depth_texture(
frame.points,
frame.valid_mask,
frame,
context.directory / "vector-depth.exr",
)
result["vector_depth"] = depth_path
else:
depth_path = write_z_depth_texture(
frame.depth,
frame.valid_mask,
frame,
context.directory / "z-depth.exr",
)
result["z_depth"] = depth_path
result["depth_metadata"] = write_depth_metadata(
frame,
_alpha_mask(input_path, frame.valid_mask.shape),
_alpha_mask(input_path, frame.depth.shape),
context.directory / "depth.json",
)
@@ -111,7 +107,7 @@ def generate_moge2_artifacts(
context.progress(0.72, "Generating surface normal")
key = f"{normal_mode.lower()}_normal"
result[key] = _write_normal(
prediction,
frame,
input_path,
context.directory / f"{normal_mode.lower()}-normal.png",
normal_mode,