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

8
src/anyimage/server/app.py
View File

@@ -16,7 +16,7 @@ from .jobs.debug_predictions import (
)
from .jobs.cutout import run as _generate_cutout_artifacts
from .jobs.depth_plane import run as _generate_depth_plane_geometry
from .jobs.refine_image_region import run as _refine_image_region
from .jobs.refine_image_selection import run as _refine_image_selection
from .jobs.remove_background import run as _remove_background
from .jobs.upscale import run as _upscale_image
from .model_manager import ModelManager
@@ -41,9 +41,9 @@ def download_required_models(context, parameters):
return manager.download_missing(context, parameters["models"])
@server.job("refine-image-region")
def refine_image_region(context, parameters):
return _refine_image_region(context, parameters)
@server.job("refine-image-selection")
def refine_image_selection(context, parameters):
return _refine_image_selection(context, parameters)
@server.job("remove-background")

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,

24
src/anyimage/server/jobs/cutout.py
View File

@@ -1,35 +1,35 @@
"""Generate optional AI files for one Blender-rasterized Cutout Region."""
"""Generate optional AI files for one Blender-rasterized Cutout Selection."""
import json
from pathlib import Path
from ..geometry.prediction_artifacts import generate_moge2_artifacts
from ..region import refine_region
from ..selection import refine_selection
def run(context, parameters):
context.progress(0.05, "Preparing Cutout Region")
context.progress(0.05, "Preparing Cutout Selection")
context.directory.mkdir(parents=True, exist_ok=True)
input_path = Path(parameters["input"])
region_image_path = input_path
selection_image_path = input_path
result = {}
if parameters.get("refine_region"):
context.progress(0.2, "Refining Region")
region_image_path, bounds = refine_region(
if parameters.get("refine_selection"):
context.progress(0.2, "Refining Selection")
selection_image_path, bounds = refine_selection(
context,
input_path,
parameters["refine_model"],
parameters["device"],
context.directory / "region.png",
context.directory / "selection.png",
)
bounds_path = context.directory / "region_bounds.json"
bounds_path = context.directory / "selection_bounds.json"
bounds_path.write_text(
json.dumps(bounds, separators=(",", ":")),
encoding="utf-8",
)
result.update(
region_image=region_image_path.name,
region_bounds=bounds_path.name,
selection_image=selection_image_path.name,
selection_bounds=bounds_path.name,
)
depth_mode = str(parameters.get("depth_mode", "NONE"))
@@ -38,7 +38,7 @@ def run(context, parameters):
artifacts = generate_moge2_artifacts(
context,
parameters,
region_image_path,
selection_image_path,
depth_mode=depth_mode,
normal_mode=normal_mode,
)

10
src/anyimage/server/jobs/refine_image_region.py → src/anyimage/server/jobs/refine_image_selection.py
View File

@@ -1,18 +1,18 @@
import json
import time
from ..region import refine_region
from ..selection import refine_selection
def run(context, parameters):
"""Refine a gesture mask with BEN2 and return its source-image bounds."""
started = time.perf_counter()
context.progress(0.05, "Preparing image Region")
context.progress(0.05, "Preparing image Selection")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
output_path, bounds = refine_region(
output_path, bounds = refine_selection(
context,
parameters["input"],
parameters["model"],
@@ -22,6 +22,6 @@ def run(context, parameters):
bounds_path = output_directory / "bounds.json"
bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8")
elapsed_ms = (time.perf_counter() - started) * 1000.0
message = f"Image Region refined in {elapsed_ms:.0f} ms"
message = f"Image Selection refined in {elapsed_ms:.0f} ms"
context.progress(1.0, message)
return {"region_image": output_path.name, "region_bounds": bounds_path.name}
return {"selection_image": output_path.name, "selection_bounds": bounds_path.name}

36
src/anyimage/server/models/da3.py
View File

@@ -71,24 +71,26 @@ def infer(
return merge_frame_predictions(predictions)
def geometry_frame(prediction):
def geometry_frames(prediction):
import numpy as np
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
validity = np.asarray(prediction.conf, dtype=np.float32)
intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32)
if depth.shape[0] != 1:
raise ValueError("Image mesh generation currently supports one image")
points = unproject_camera_depth(depth[0], intrinsics[0])
valid = (
np.isfinite(depth[0])
& (depth[0] > 0.0)
& np.isfinite(points).all(axis=-1)
)
return GeometryFrame(
depth=depth[0],
points=points,
intrinsics=intrinsics[0],
valid_mask=valid,
confidence=confidence[0],
)
if validity.shape != depth.shape or intrinsics.shape != (len(depth), 3, 3):
raise ValueError("DA3 prediction fields do not share one frame layout")
return [
GeometryFrame(
depth=frame_depth,
validity=frame_validity,
intrinsics=frame_intrinsics,
points=unproject_camera_depth(frame_depth, frame_intrinsics).astype(
np.float32
),
)
for frame_depth, frame_validity, frame_intrinsics in zip(
depth,
validity,
intrinsics,
)
]

29
src/anyimage/server/models/geometry.py
View File

@@ -4,8 +4,29 @@ from dataclasses import dataclass
@dataclass
class GeometryFrame:
depth: object
points: object | None
validity: object
intrinsics: object
valid_mask: object
confidence: object | None = None
model_mask: object | None = None
normal: object | None = None
points: object | None = None
def __post_init__(self):
import numpy as np
depth = np.asarray(self.depth)
validity = np.asarray(self.validity)
intrinsics = np.asarray(self.intrinsics)
if depth.dtype != np.float32 or depth.ndim != 2:
raise ValueError("Geometry depth must be a 2D float32 field")
if validity.dtype != np.float32 or validity.shape != depth.shape:
raise ValueError("Geometry validity must match the depth as float32")
if intrinsics.dtype != np.float32 or intrinsics.shape != (3, 3):
raise ValueError("Geometry intrinsics must be a 3x3 float32 matrix")
for name in ("normal", "points"):
value = getattr(self, name)
if value is None:
continue
field = np.asarray(value)
if field.dtype != np.float32 or field.shape != (*depth.shape, 3):
raise ValueError(
f"Geometry {name} must match the depth as three-channel float32"
)

117
src/anyimage/server/models/moge2.py
View File

@@ -1,26 +1,37 @@
from dataclasses import dataclass
from .geometry import GeometryFrame
@dataclass
class Prediction:
depth: object
normal: object
mask: object
intrinsics: object
points: object | None
def _pixel_intrinsics(intrinsics, image_size):
import numpy as np
def archive_values(self):
values = {
"depth": self.depth,
"normal": self.normal,
"mask": self.mask,
"intrinsics": self.intrinsics,
}
if self.points is not None:
values["points"] = self.points
return values
width, height = image_size
result = np.asarray(intrinsics, dtype=np.float32).copy()
if (
0.0 <= float(result[0, 2]) <= 1.0
and 0.0 <= float(result[1, 2]) <= 1.0
):
result[0] *= width
result[1] *= height
result[2] = (0.0, 0.0, 1.0)
return result
def _geometry_frame(prediction, include_points):
import numpy as np
depth = np.asarray(prediction["depth"], dtype=np.float32)
height, width = depth.shape
return GeometryFrame(
depth=depth,
normal=np.asarray(prediction["normal"], dtype=np.float32),
validity=np.asarray(prediction["mask"], dtype=np.float32),
intrinsics=_pixel_intrinsics(prediction["intrinsics"], (width, height)),
points=(
np.asarray(prediction["points"], dtype=np.float32)
if include_points
else None
),
)
def infer(
@@ -32,41 +43,19 @@ def infer(
progress_callback=None,
cancel_check=None,
):
import numpy as np
depths = []
normals = []
masks = []
intrinsics = []
points = []
frames = []
for index, image_path in enumerate(images):
if cancel_check is not None:
cancel_check()
if progress_callback is not None:
progress_callback(index, len(images))
prediction = infer_one(
model_manager,
parameters,
image_path,
frames.append(
_geometry_frame(
infer_one(model_manager, parameters, image_path),
include_points,
)
)
depths.append(prediction["depth"])
normals.append(prediction["normal"])
masks.append(prediction["mask"].astype(bool))
intrinsics.append(prediction["intrinsics"])
if include_points:
points.append(prediction["points"])
if len({array.shape for array in depths}) != 1:
raise RuntimeError("MoGe-2 frame sizes must match")
return Prediction(
depth=np.stack(depths).astype(np.float32),
normal=np.stack(normals).astype(np.float32),
mask=np.stack(masks),
intrinsics=np.stack(intrinsics).astype(np.float32),
points=(
np.stack(points).astype(np.float32)
if include_points
else None
),
)
return frames
def infer_one(model_manager, parameters, image_path):
@@ -87,35 +76,3 @@ def infer_one(model_manager, parameters, image_path):
image,
int(parameters["resolution_level"]),
)
def geometry_frame(prediction):
import numpy as np
if prediction.depth.shape[0] != 1:
raise ValueError("Image depth generation requires one MoGe-2 frame")
depth = prediction.depth[0]
points = prediction.points[0] if prediction.points is not None else None
intrinsics = prediction.intrinsics[0].copy()
height, width = depth.shape
if (
0.0 <= float(intrinsics[0, 2]) <= 1.0
and 0.0 <= float(intrinsics[1, 2]) <= 1.0
):
intrinsics[0] *= width
intrinsics[1] *= height
intrinsics[2] = (0.0, 0.0, 1.0)
valid = (
prediction.mask[0]
& np.isfinite(depth)
& (depth > 0.0)
)
if points is not None:
valid &= np.isfinite(points).all(axis=-1)
return GeometryFrame(
depth=depth,
points=points,
intrinsics=intrinsics,
valid_mask=valid,
model_mask=prediction.mask[0],
)

5
src/anyimage/server/region/__init__.py
View File

@@ -1,5 +0,0 @@
"""AI generation for rasterized image Regions."""
from .refine_region import refine_region
__all__ = ("refine_region",)

5
src/anyimage/server/selection/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
"""AI generation for rasterized image Selections."""
from .refine_selection import refine_selection
__all__ = ("refine_selection",)

6
src/anyimage/server/region/refine_region.py → src/anyimage/server/selection/refine_selection.py
View File

@@ -7,7 +7,7 @@ def _crop_to_alpha(image):
pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
if not len(visible_x):
raise ValueError("The generated Region contains no visible pixels")
raise ValueError("The generated Selection contains no visible pixels")
bounds = (
int(visible_x.min()),
int(visible_y.min()),
@@ -17,14 +17,14 @@ def _crop_to_alpha(image):
return image.crop(bounds), bounds
def refine_region(
def refine_selection(
context,
input_path,
model,
device,
output_path,
):
"""Refine an image Region from an already rasterized Region."""
"""Refine an image Selection from an already rasterized Selection."""
model_manager = context.resource("model_manager")
foreground, _load_ms, _inference_ms = infer_foreground(
input_path,