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

Add new tests and functionality for image processing and geometry handling

Browse Source 
- Introduced color_image.py for preparing standalone color images for tests.
- Added hdr_image.py to create HDR textures for testing.
- Implemented test_image_depth_panorama.py to evaluate panorama topology, UVs, and depth discontinuities.
- Created test_shadeless.py to validate the opacity and HDR color endpoints of the Shadeless group.
- Developed moge3_export.py for exporting dynamic MoGe-3 backbone and ONNX graphs.
- Added material.py to inherit material slots onto generated geometry.
- Built image_depth_panorama.py to construct a quad panorama with depth separation and sampling.
This commit is contained in:
马南
2026-09-10 14:42:17 +08:00
parent 27ca25224a
commit 489003c2cc
136 changed files with 5758 additions and 1083 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,7 @@ from blendjob import JobServer
from .jobs.cutout import run as _generate_cutout_artifacts
from .jobs.depth_plane import run as _generate_depth_plane_geometry
from .jobs.panorama import run as _generate_panorama_geometry
from .jobs.remove_background import run as _remove_background
from .jobs.upscale import run as _upscale_image
from .model_manager import ModelManager
@@ -46,3 +47,8 @@ def generate_depth_plane_geometry(context, parameters):
@server.job("generate-cutout-artifacts")
def generate_cutout_artifacts(context, parameters):
return _generate_cutout_artifacts(context, parameters)
@server.job("generate-panorama-geometry")
def generate_panorama_geometry(context, parameters):
return _generate_panorama_geometry(context, parameters)

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

@@ -25,7 +25,7 @@ def _exr_channels():
) + b"\0"
def _write_float_exr(rgba, output_path):
def write_float_exr(rgba, output_path):
"""Write one uncompressed scanline RGBA OpenEXR without extra dependencies."""
import numpy as np
@@ -87,7 +87,7 @@ def write_depth_texture(frame, output_path, *, alpha):
rgba = np.empty((*validity.shape, 4), dtype=np.float32)
rgba[..., :3] = points
rgba[..., 3] = alpha * validity
return _write_float_exr(rgba, output_path)
return write_float_exr(rgba, output_path)
def write_depth_metadata(frame, reference_valid, output_path):

156
src/anyimage/server/geometry/panorama.py Normal file
View File

@@ -0,0 +1,156 @@
"""Spherical projection and periodic radial-distance fusion."""
import numpy as np
from scipy.ndimage import map_coordinates
from scipy.sparse import coo_matrix, diags
from scipy.sparse.linalg import spsolve
def panorama_directions(width, height):
u, v = np.meshgrid((np.arange(width) + 0.5) / width,
(np.arange(height) + 0.5) / height)
return uv_directions(u, v)
def uv_directions(u, v):
theta, phi = 2 * np.pi * u, np.pi * v
return np.stack((np.sin(phi) * np.cos(theta),
np.sin(phi) * np.sin(theta), np.cos(phi)), axis=-1)
def camera_bases():
ratio = (1 + np.sqrt(5)) / 2
vertices = []
for a in (-1, 1):
for b in (-ratio, ratio):
vertices.extend(((0, a, b), (a, b, 0), (b, 0, a)))
bases = []
for forward in np.asarray(vertices):
forward = forward / np.linalg.norm(forward)
right = np.cross(forward, (0, 0, 1))
right /= np.linalg.norm(right)
down = np.cross(forward, right)
bases.append(np.stack((right, down, forward)))
return bases
def view_rays(size):
x, y = np.meshgrid(2 * (np.arange(size) + 0.5) / size - 1,
2 * (np.arange(size) + 0.5) / size - 1)
return np.stack((x, y, np.ones_like(x)), axis=-1)
def sample_panorama(image, directions, *, order=1):
directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True)
u = np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1
v = np.arccos(np.clip(directions[..., 2], -1, 1)) / np.pi
height, width = image.shape[:2]
padded = np.concatenate((image[:, -1:], image, image[:, :1]), axis=1)
coords = (np.clip(v * height - 0.5, 0, height - 1), u * width + 0.5)
if image.ndim == 2:
return map_coordinates(padded, coords, order=order, mode="nearest")
return np.stack([map_coordinates(padded[..., c], coords, order=order,
mode="nearest") for c in range(image.shape[2])], -1)
def warp_distances(distances, masks, bases, width, height, *, directions=None):
if directions is None:
directions = panorama_directions(width, height)
logs, validities = [], []
for distance, mask, basis in zip(distances, masks, bases):
rays = directions @ basis.T
projected = rays[..., :2] / np.maximum(rays[..., 2:], 1e-8)
valid = (rays[..., 2] > 0) & (np.abs(projected) < 0.98).all(axis=-1)
size = distance.shape[0]
coords = ((projected[..., 1] + 1) * size / 2 - 0.5,
(projected[..., 0] + 1) * size / 2 - 0.5)
valid &= map_coordinates(mask.astype(float), coords, order=0, mode="nearest") > 0.5
logs.append(map_coordinates(np.log(np.maximum(distance, 1e-6)), coords,
order=1, mode="nearest"))
validities.append(valid)
return np.asarray(logs), np.asarray(validities)
def merge_distances(distances, masks, bases, width, height, *, cancel_check=None):
"""Fuse log-distance derivatives with a cyclic longitude and a scale anchor.
Algorithm reference: microsoft/MoGe, moge/utils/panorama.py.
This independent implementation solves normal equations on a bounded panorama grid.
"""
if cancel_check is not None:
cancel_check()
if width > 512:
return _refine_distances(distances, masks, bases, width, height, cancel_check)
logs, valid = warp_distances(distances, masks, bases, width, height)
view_count = len(logs)
count = width * height
index = np.arange(count).reshape(height, width)
start = np.concatenate((index.ravel(), index[:-1].ravel()))
end = np.concatenate((np.roll(index, -1, axis=1).ravel(), index[1:].ravel()))
rows = np.arange(len(start))
gradient = coo_matrix((np.concatenate((-np.ones(len(start)), np.ones(len(start)))),
(np.tile(rows, 2), np.concatenate((start, end)))),
shape=(len(start), count)).tocsr()
laplacian = -(gradient.T @ gradient).tocsr()
edge_valid = valid.reshape(view_count, -1)[:, start] & valid.reshape(view_count, -1)[:, end]
grad_samples = logs.reshape(view_count, -1)[:, end] - logs.reshape(view_count, -1)[:, start]
grad_target = np.sum(grad_samples * edge_valid, axis=0) / np.maximum(edge_valid.sum(0), 1)
lap_valid = valid.copy()
lap_valid &= np.roll(valid, 1, axis=2) & np.roll(valid, -1, axis=2)
lap_valid[:, 1:] &= valid[:, :-1]
lap_valid[:, :-1] &= valid[:, 1:]
lap_valid = lap_valid.reshape(view_count, -1)
lap_samples = np.asarray((laplacian @ logs.reshape(view_count, -1).T).T)
lap_target = np.sum(lap_samples * lap_valid, axis=0) / np.maximum(lap_valid.sum(0), 1)
g = gradient[edge_valid.any(0)]
l = laplacian[lap_valid.any(0)]
coverage = valid.any(0)
if not coverage.any():
raise ValueError("No valid panorama predictions")
baseline = np.sum(logs * valid, axis=0) / np.maximum(valid.sum(0), 1)
matrix = g.T @ g + l.T @ l + diags(np.full(count, 1e-7))
target = g.T @ grad_target[edge_valid.any(0)] + l.T @ lap_target[lap_valid.any(0)]
if cancel_check is not None:
cancel_check()
solution = spsolve(matrix.tocsc(), target + 1e-7 * baseline.ravel()).reshape(height, width)
if cancel_check is not None:
cancel_check()
solution += np.median(baseline[coverage] - solution[coverage])
return np.exp(solution).astype(np.float32), coverage
def _refine_distances(distances, masks, bases, width, height, cancel_check):
"""Apply coarse periodic alignment to view detail in bounded output tiles."""
coarse_width, coarse_height = 512, 256
coarse, _ = merge_distances(distances, masks, bases, coarse_width, coarse_height,
cancel_check=cancel_check)
logs, valid = warp_distances(distances, masks, bases, coarse_width, coarse_height)
corrections = np.log(coarse)[None] - logs
result = np.empty((height, width), np.float32)
mask = np.empty((height, width), bool)
tile_height = max(1, 65536 // width)
for start in range(0, height, tile_height):
if cancel_check is not None:
cancel_check()
stop = min(start + tile_height, height)
u, v = np.meshgrid((np.arange(width) + 0.5) / width,
(np.arange(start, stop) + 0.5) / height)
directions = uv_directions(u, v)
detail, visible = warp_distances(distances, masks, bases, width, height,
directions=directions)
total = np.zeros(u.shape)
weights = np.zeros(u.shape)
for index, basis in enumerate(bases):
support = sample_panorama(valid[index].astype(float), directions)
correction = sample_panorama(corrections[index] * valid[index], directions)
correction /= np.maximum(support, 1e-8)
forward = directions @ basis[2]
weight = visible[index] * np.maximum(forward, 0) ** 4 * (support > 1e-8)
total += (detail[index] + correction) * weight
weights += weight
fallback = sample_panorama(np.log(coarse), directions)
np.divide(total, weights, out=fallback, where=weights > 0)
result[start:stop] = np.exp(fallback)
mask[start:stop] = visible.any(axis=0)
return result, mask

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

@@ -1,7 +1,8 @@
"""Export requested depth and normal artifacts from one MoGe-2 prediction."""
"""Export requested depth and normal artifacts from one MoGe prediction."""
from ..media.input import open_image
from ..media import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
from ..models import moge2
from ..models import moge
from .depth_texture import (
write_depth_metadata,
write_depth_texture,
@@ -11,9 +12,7 @@ from .edge_depth import extend_edge_depth
def _limited_model_input(source_path, output_path, max_input_size):
from PIL import Image
with Image.open(source_path) as opened:
with open_image(source_path) as opened:
limited = limit_image(opened.convert("RGB"), max_input_size)
limited.save(output_path, format="PNG", compress_level=1)
return output_path
@@ -23,7 +22,7 @@ def _image_alpha(image_path, shape):
import numpy as np
from PIL import Image
with Image.open(image_path) as opened:
with open_image(image_path) as opened:
alpha = opened.convert("RGBA").getchannel("A")
alpha = alpha.resize(
(int(shape[1]), int(shape[0])),
@@ -33,10 +32,8 @@ def _image_alpha(image_path, shape):
def _write_normal(frame, color_path, output_path, normal_mode):
from PIL import Image
valid = _image_alpha(color_path, frame.depth.shape) >= 0.1
with Image.open(color_path) as source:
with open_image(color_path) as source:
output_size = source.size
return write_camera_normal_texture(
frame,
@@ -49,7 +46,7 @@ def _write_normal(frame, color_path, output_path, normal_mode):
)
def generate_moge2_artifacts(
def generate_moge_artifacts(
context,
parameters,
input_path,
@@ -58,7 +55,7 @@ def generate_moge2_artifacts(
normal_mode="NONE",
alpha_threshold=0.9,
):
"""Run MoGe-2 once and write only the requested file artifacts."""
"""Run MoGe once and write only the requested file artifacts."""
normal_mode = str(normal_mode).upper()
if normal_mode not in {"NONE", "TANGENT", "OBJECT"}:
raise ValueError(f"Unsupported normal mode: {normal_mode}")
@@ -68,7 +65,7 @@ def generate_moge2_artifacts(
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
)
try:
frames = moge2.infer(
frames = moge.infer(
context.resource("model_manager"),
parameters,
[model_path],

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

@@ -2,7 +2,7 @@
from pathlib import Path
from ..geometry.prediction_artifacts import generate_moge2_artifacts
from ..geometry.prediction_artifacts import generate_moge_artifacts
def run(context, parameters):
@@ -13,7 +13,7 @@ def run(context, parameters):
generate_depth = bool(parameters.get("generate_depth", False))
normal_mode = str(parameters.get("normal_mode", "NONE"))
if generate_depth or normal_mode != "NONE":
artifacts = generate_moge2_artifacts(
artifacts = generate_moge_artifacts(
context,
parameters,
input_path,

9
src/anyimage/server/jobs/depth_plane.py
View File

@@ -1,7 +1,7 @@
"""Generate model-space depth and preserve color for one complete image."""
"""Generate model-space depth and normals for one complete image."""
from ..geometry.prediction_artifacts import generate_moge2_artifacts
from ..media import copy_color_images, input_frames
from ..geometry.prediction_artifacts import generate_moge_artifacts
from ..media import input_frames
def run(context, parameters):
@@ -13,8 +13,7 @@ def run(context, parameters):
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Depth conversion currently supports one image")
copy_color_images(images, context.directory)
artifacts = generate_moge2_artifacts(
artifacts = generate_moge_artifacts(
context,
parameters,
images[0],

82
src/anyimage/server/jobs/panorama.py Normal file
View File

@@ -0,0 +1,82 @@
"""Generate a radial depth texture from the current RGBA panorama."""
import json
from pathlib import Path
from ..geometry.depth_texture import write_float_exr
from ..geometry.panorama import camera_bases, merge_distances, sample_panorama, view_rays
from ..media import VIDEO_SUFFIXES, input_frames
from ..media.input import open_image
from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, normalized_max_input_size
from ..models.moge import infer_image as infer
def run(context, parameters):
import numpy as np
context.check_cancelled()
input_path = Path(parameters["input"])
if input_path.is_dir() or input_path.suffix.lower() in VIDEO_SUFFIXES:
raise ValueError("Panorama conversion requires one static 2:1 image")
maximum = normalized_max_input_size(parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE))
context.directory.mkdir(parents=True, exist_ok=True)
outputs = [context.directory / name for name in ("depth.exr", "depth.json")]
try:
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Panorama conversion requires one static 2:1 image")
with open_image(images[0]) as opened:
if getattr(opened, "n_frames", 1) != 1 or opened.width != opened.height * 2:
raise ValueError("Panorama conversion requires one static 2:1 image")
source = opened.convert("RGBA")
width, height = source.size
size = min(maximum, max(512, width // 4))
image = np.asarray(source, dtype=np.float32) / 255.0
manager = context.resource("model_manager")
bases = camera_bases()
rays = view_rays(size)
distances, masks = [], []
session = None
for index, basis in enumerate(bases):
context.check_cancelled()
context.progress(0.05 + 0.65 * index / len(bases), f"Generating panorama view {index + 1}/12")
directions = rays @ basis
alpha = sample_panorama(image[..., 3], directions, order=0)
visible = alpha > 0
if not visible.any():
distances.append(np.ones((size, size), dtype=np.float32))
masks.append(visible)
continue
if session is None:
session = manager.get_moge(manager.directory(parameters["model"]), parameters["device"])
color = sample_panorama(image[..., :3], directions) * alpha[..., None]
prediction = infer(session, color, parameters["resolution_level"],
fov_x=90.0, source_valid=visible)
distances.append(np.linalg.norm(prediction["points"], axis=-1))
masks.append((prediction["mask"] > 0.5) & visible)
del prediction
context.progress(0.75, "Merging panoramic distance")
distance, valid = merge_distances(distances, masks, bases, width, height,
cancel_check=context.check_cancelled)
alpha = np.asarray(source.getchannel("A"),
dtype=np.float32) / 255.0
alpha *= valid
reference = (alpha >= 0.9) & np.isfinite(distance) & (distance > 0)
if not reference.any():
raise ValueError("The panorama contains no valid depth pixels")
if not np.isfinite(distance).all() or (distance <= 0).any():
raise ValueError("Panorama fusion produced invalid distances")
rgba = np.empty((height, width, 4), dtype=np.float32)
rgba[..., :3], rgba[..., 3] = distance[..., None], alpha
write_float_exr(rgba, outputs[0])
outputs[1].write_text(json.dumps({
"projection": "equirectangular", "image_size": [width, height],
"reference_distance": float(np.median(distance[reference])),
}, separators=(",", ":")), encoding="utf-8")
context.check_cancelled()
context.progress(1.0, "Panorama conversion data is ready")
return {"depth": outputs[0].name, "depth_metadata": outputs[1].name}
except BaseException:
for path in outputs:
path.unlink(missing_ok=True)
raise

22
src/anyimage/server/jobs/remove_background.py
View File

@@ -1,7 +1,8 @@
import time
from ..media import input_frames
from ..models.ben2 import infer_foreground, infer_foreground_frames
from ..media.input import open_image
from ..models.ben2 import infer_alpha, infer_foreground, infer_foreground_frames
def run(context, parameters):
@@ -15,10 +16,25 @@ def run(context, parameters):
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
output_kind = parameters.get("output_kind", "foreground")
if output_kind not in {"foreground", "alpha"}:
raise ValueError(f"Unknown background removal output: {output_kind}")
with input_frames(context, parameters) as images:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
if output_kind == "alpha":
import numpy as np
if len(images) != 1:
raise ValueError("Alpha output requires one still image")
with open_image(images[0]) as source:
alpha, load_ms, inference_ms = infer_alpha(
source, model_directory, device, model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
np.save(output_directory / "alpha.npy", alpha, allow_pickle=False)
elif len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
@@ -62,4 +78,4 @@ def run(context, parameters):
)
context.log(message)
context.progress(1.0, message)
return {"image": "."}
return {"alpha": "alpha.npy"} if output_kind == "alpha" else {"image": "."}

2
src/anyimage/server/media/__init__.py
View File

@@ -5,7 +5,6 @@ from .input import (
analysis_frames,
check_cancelled,
collect_images,
copy_color_images,
extract_video_frames,
first_frames,
input_frames,
@@ -29,7 +28,6 @@ __all__ = (
"analysis_frames",
"check_cancelled",
"collect_images",
"copy_color_images",
"extract_video_frames",
"first_frames",
"input_frames",

89
src/anyimage/server/media/input.py
View File

@@ -16,6 +16,52 @@ VIDEO_SUFFIXES = {
".mpg",
".mpeg",
}
def validate_input_path(value, *, images_only=False):
"""Validate supported media without importing image or video libraries."""
path = Path(value)
if path.is_dir() and not images_only:
if not collect_images(path):
supported = ", ".join(sorted(IMAGE_SUFFIXES))
raise RuntimeError(
f"No supported image frames found in: {path}. Supported images: {supported}."
)
return path
if not path.is_file():
kind = "video" if path.suffix.lower() in VIDEO_SUFFIXES else "image"
raise RuntimeError(f"Source {kind} file does not exist: {path}")
suffix = path.suffix.lower()
supported = IMAGE_SUFFIXES if images_only else IMAGE_SUFFIXES | VIDEO_SUFFIXES
if suffix not in supported:
message = (
f"Unsupported input format: {suffix or '(no extension)'}. "
f"Supported images: {', '.join(sorted(IMAGE_SUFFIXES))}."
)
if not images_only:
message += f" Supported videos: {', '.join(sorted(VIDEO_SUFFIXES))}."
raise RuntimeError(message)
return path
@contextmanager
def open_image(path):
"""Decode an image before handing it to processing code."""
from PIL import Image
opened = None
try:
try:
opened = Image.open(path)
opened.load()
except (OSError, ValueError, SyntaxError) as error:
raise RuntimeError(f"Unable to read image: {path}") from error
yield opened
finally:
if opened is not None:
opened.close()
class TaskCancelled(RuntimeError):
pass
@@ -55,18 +101,32 @@ def extract_video_frames(
try:
container = av.open(str(video_path))
except Exception as error:
raise RuntimeError(f"Unable to open video: {video_path}") from error
raise RuntimeError(f"Unable to read video: {video_path}") from error
image_paths = []
skipped = 0
try:
stream = container.streams.video[0]
for frame in container.decode(stream):
try:
stream = container.streams.video[0]
frames = iter(container.decode(stream))
except Exception as error:
raise RuntimeError(f"Unable to read video: {video_path}") from error
while True:
try:
frame = next(frames)
except StopIteration:
break
except Exception as error:
raise RuntimeError(f"Unable to read video: {video_path}") from error
check_cancelled(cancel_check)
if skipped < start_frame - 1:
skipped += 1
continue
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
frame.to_image().save(image_path, format="JPEG")
try:
decoded = frame.to_image()
except Exception as error:
raise RuntimeError(f"Unable to read video: {video_path}") from error
decoded.save(image_path, format="JPEG")
image_paths.append(image_path)
if len(image_paths) >= maximum_frames:
break
@@ -74,7 +134,7 @@ def extract_video_frames(
container.close()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
raise RuntimeError(f"Unable to read video: {video_path}")
return image_paths
@@ -86,18 +146,13 @@ def prepare_input(
cancel_check=None,
):
check_cancelled(cancel_check)
input_path = Path(input_path)
input_path = validate_input_path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES:
return [input_path], None
if suffix not in VIDEO_SUFFIXES:
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
import av
@@ -145,15 +200,3 @@ def analysis_frames(context, parameters, fallback_images):
return
with input_frames(context, dict(parameters, input=analysis_input)) as images:
yield images
def copy_color_images(images, output_directory):
"""Copy input images as the canonical numbered image layer."""
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"image.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)

24
src/anyimage/server/model_catalog.py
View File

@@ -178,10 +178,28 @@ MOGE2_MODELS = {
),),
),
}
MOGE3_VITL = ModelSpec(
key="MOGE3_VITL",
label="MoGe-3 ViT-L",
directory_name="moge-3-vitl-onnx",
repository="Ruicheng/moge-3-vitl",
license_name="MIT",
description="MoGe-3 geometry with three refinement steps using ONNX Runtime",
download_patterns=("backbone.onnx", "refiner.onnx"),
ready_patterns=("backbone.onnx", "refiner.onnx"),
r2_directory="moge-3-vitl-onnx",
r2_files=(
("backbone.onnx", 1324190318, "f21246c6921373c8c79ecdb8c78818a79bd2506c2fb88f4d4cdec21b92f725e3"),
("refiner.onnx", 157729628, "67ad7297aba18f8ef7939e8f700a5b25c41a19ddb68656bc2cd08b54c248fe2a"),
),
)
MOGE_MODELS = {**MOGE2_MODELS, MOGE3_VITL.key: MOGE3_VITL}
DOWNLOADABLE_MODELS = {
BEN2_MODEL.key: BEN2_MODEL,
**UPSCALE_MODELS,
**MOGE2_MODELS,
**MOGE_MODELS,
}
@@ -200,8 +218,8 @@ def model_record(model, ready=False):
"description": model.description,
"license": model.license_name,
"family": (
"moge2"
if model.key.startswith("MOGE2_")
"moge"
if model.key in MOGE_MODELS
else "ben2"
if model.key == BEN2_MODEL.key
else "upscale"

11
src/anyimage/server/model_manager.py
View File

@@ -38,9 +38,12 @@ class ModelSessionCache:
return model, elapsed_ms
def get_moge2(self, model_dir, requested_device):
def get_moge(self, model_dir, requested_device):
from .models.onnx_moge2 import create_session
if Path(model_dir).name == "moge-3-vitl-onnx":
from .models.onnx_moge3 import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.moge_model is not None and self.moge_key == key:
return self.moge_model
@@ -122,7 +125,7 @@ class ModelManager:
def ready(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
directory = self.directory(model_key)
if model_key == "PISA_SR":
if model_key in {"PISA_SR", "MOGE3_VITL"}:
for filename, size, checksum in model.r2_files:
path = directory / filename
try:
@@ -149,8 +152,8 @@ class ModelManager:
return self.sessions.get_ben2(model_dir, device)
def get_moge2(self, model_dir, device):
return self.sessions.get_moge2(model_dir, device)
def get_moge(self, model_dir, device):
return self.sessions.get_moge(model_dir, device)
def get_upscale(self, model_dir, device):
return self.sessions.get_upscale(model_dir, device)

4
src/anyimage/server/models/__init__.py
View File

@@ -1,7 +1,7 @@
from . import ben2, moge2, upscale
from . import ben2, moge, upscale
__all__ = (
"ben2",
"moge2",
"moge",
"upscale",
)

61
src/anyimage/server/models/ben2.py
View File

@@ -1,25 +1,22 @@
import time
from pathlib import Path
from ..media import IMAGE_SUFFIXES, check_cancelled
from ..media import check_cancelled
from ..media.input import open_image, validate_input_path
def infer_foreground(
input_path,
def infer_alpha(
image,
model_dir,
device,
model_cache=None,
cancel_check=None,
preserve_input_alpha=False,
release_memory=True,
):
from PIL import Image, ImageChops
from PIL import Image
input_path = Path(input_path)
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
from .onnx_ben2 import create_session, infer
from .onnx_ben2 import create_session, infer_alpha as predict_alpha
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
@@ -31,8 +28,7 @@ def infer_foreground(
model, load_ms = model_cache.get_ben2(model_dir, device)
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
source_rgba = opened_image.convert("RGBA")
source_rgba = image.convert("RGBA")
image = Image.new("RGB", source_rgba.size, (0, 0, 0))
image.paste(
source_rgba.convert("RGB"),
@@ -40,23 +36,44 @@ def infer_foreground(
source_rgba.getchannel("A"),
)
inference_started = time.perf_counter()
foreground = infer(
alpha = predict_alpha(
model,
image,
release_memory=release_memory,
)
inference_ms = (time.perf_counter() - inference_started) * 1000.0
if preserve_input_alpha:
foreground_alpha = foreground.convert("RGBA").getchannel("A")
source_alpha = source_rgba.getchannel("A")
if source_alpha.size != foreground_alpha.size:
source_alpha = source_alpha.resize(
foreground_alpha.size,
Image.Resampling.BILINEAR,
)
source_rgba.putalpha(ImageChops.multiply(source_alpha, foreground_alpha))
foreground = source_rgba
check_cancelled(cancel_check)
return alpha, load_ms, inference_ms
def infer_foreground(
input_path,
model_dir,
device,
model_cache=None,
cancel_check=None,
preserve_input_alpha=False,
release_memory=True,
):
import numpy as np
from PIL import Image, ImageChops
check_cancelled(cancel_check)
input_path = validate_input_path(Path(input_path), images_only=True)
with open_image(input_path) as opened_image:
source = opened_image.convert("RGBA")
alpha, load_ms, inference_ms = infer_alpha(
source, model_dir, device, model_cache=model_cache,
cancel_check=cancel_check, release_memory=release_memory,
)
mask = Image.fromarray((alpha * 255).astype(np.uint8))
if preserve_input_alpha:
foreground = source
mask = ImageChops.multiply(source.getchannel("A"), mask)
else:
foreground = Image.new("RGBA", source.size, (0, 0, 0, 255))
foreground.paste(source.convert("RGB"), (0, 0), source.getchannel("A"))
foreground.putalpha(mask)
return foreground, load_ms, inference_ms

18
src/anyimage/server/models/moge2.py → src/anyimage/server/models/moge.py
View File

@@ -1,3 +1,4 @@
from ..media.input import open_image
from .geometry import GeometryFrame
@@ -71,20 +72,25 @@ def infer_one(
release_memory=True,
):
import numpy as np
from PIL import Image
from .onnx_moge2 import infer as infer_onnx
model_directory = model_manager.directory(parameters["model"])
model = model_manager.get_moge2(
model = model_manager.get_moge(
model_directory,
parameters["device"],
)
with Image.open(image_path) as opened:
with open_image(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
return infer_onnx(
return infer_image(
model,
image,
int(parameters["resolution_level"]),
release_memory=release_memory,
)
def infer_image(session, image, resolution_level, **options):
"""Dispatch one RGB array to the selected MoGe ONNX model."""
from . import onnx_moge2, onnx_moge3
backend = onnx_moge3 if isinstance(session, onnx_moge3.Moge3Session) else onnx_moge2
return backend.infer(session, image, resolution_level, **options)

16
src/anyimage/server/models/onnx_ben2.py
View File

@@ -34,7 +34,7 @@ def preprocess(image):
return np.ascontiguousarray(array.transpose(2, 0, 1)[None])
def infer(session, image, *, release_memory=True):
def infer_alpha(session, image, *, release_memory=True):
import numpy as np
from PIL import Image
@@ -46,12 +46,14 @@ def infer(session, image, *, release_memory=True):
{input_name: preprocess(image)},
release_memory=release_memory,
)[0]
alpha = np.asarray(alpha, dtype=np.float32).squeeze()
alpha = np.asarray(alpha, dtype=np.float32)
if (alpha.ndim != 4 or alpha.shape[:2] != (1, 1)
or not all(alpha.shape[2:]) or not np.isfinite(alpha).all()):
raise RuntimeError("BEN2 returned an invalid alpha mask")
alpha = alpha[0, 0]
minimum = float(alpha.min())
span = max(float(alpha.max()) - minimum, np.finfo(np.float32).eps)
alpha = np.clip((alpha - minimum) / span * 255.0, 0.0, 255.0)
mask = Image.fromarray(alpha.astype(np.uint8), mode="L")
alpha = np.clip((alpha - minimum) / span, 0.0, 1.0)
mask = Image.fromarray(alpha)
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
foreground = image.convert("RGBA")
foreground.putalpha(mask)
return foreground
return np.clip(np.asarray(mask, dtype=np.float32), 0.0, 1.0)

55
src/anyimage/server/models/onnx_moge2.py
View File

@@ -109,7 +109,33 @@ def depth_to_points(depth, intrinsics):
return np.stack((x, y, depth), axis=-1)
def postprocess(raw_output):
def recover_known_fov_shift(points, mask, fx, fy):
"""Fit camera Z translation with fixed intrinsics on visible predictions."""
import numpy as np
from scipy.optimize import least_squares
height, width = mask.shape
intrinsics = np.array(((fx, 0, 0.5), (0, fy, 0.5), (0, 0, 1)))
rays = depth_to_points(np.ones((height, width)), intrinsics)
valid = (mask > 0.5) & np.isfinite(points).all(-1)
if valid.sum() < 16:
return 0.0, np.zeros_like(valid)
stride = max(1, min(height, width) // 64)
selected = valid[::stride, ::stride]
p = points[::stride, ::stride][selected].astype(float)
q = rays[::stride, ::stride, :2][selected]
if len(p) < 2:
p, q = points[valid].astype(float), rays[..., :2][valid]
initial = np.sum(q * (p[:, :2] - q * p[:, 2:])) / max(np.sum(q * q), 1e-8)
lower = -float(points[..., 2][valid].min()) + 1e-5
fitted = least_squares(
lambda shift: (p[:, :2] / (p[:, 2:] + shift[0]) - q).ravel(),
[max(initial, lower + 1e-4)], bounds=([lower], [np.inf]),
)
return fitted.x[0], valid
def postprocess(raw_output, *, fov_x=None, source_valid=None):
import numpy as np
points = np.asarray(raw_output["points"][0], dtype=np.float32)
@@ -118,16 +144,30 @@ def postprocess(raw_output):
metric_scale = np.float32(np.asarray(raw_output["metric_scale"]).reshape(-1)[0])
height, width = mask.shape
aspect_ratio = width / height
focal, shift = recover_focal_shift(points, mask > 0.5)
diagonal_factor = (1.0 + aspect_ratio**2) ** 0.5
fx = focal / 2.0 * diagonal_factor / aspect_ratio
fy = focal / 2.0 * diagonal_factor
if fov_x is None:
focal, shift = recover_focal_shift(points, mask > 0.5)
diagonal_factor = (1.0 + aspect_ratio**2) ** 0.5
fx = focal / 2.0 * diagonal_factor / aspect_ratio
fy = focal / 2.0 * diagonal_factor
else:
if not 0.0 < fov_x < 180.0:
raise ValueError("Horizontal FOV must be between 0 and 180 degrees")
fx = 0.5 / np.tan(np.deg2rad(fov_x) / 2)
fy = fx * aspect_ratio
if source_valid is not None:
mask = mask * np.asarray(source_valid, dtype=bool)
shift, valid = recover_known_fov_shift(points, mask, fx, fy)
mask = np.where(valid, mask, 0.0)
intrinsics = np.array(
((fx, 0.0, 0.5), (0.0, fy, 0.5), (0.0, 0.0, 1.0)),
dtype=np.float32,
)
depth = points[..., 2] + shift
depth = depth * metric_scale
if fov_x is not None:
valid = np.isfinite(depth) & (depth > 0)
mask = np.where(valid, mask, 0.0)
depth = np.where(valid, depth, 1.0)
points = depth_to_points(depth, intrinsics)
return {
"points": points.astype(np.float32),
@@ -138,7 +178,8 @@ def postprocess(raw_output):
}
def infer(session, image, resolution_level, *, release_memory=True):
def infer(session, image, resolution_level, *, release_memory=True,
fov_x=None, source_valid=None):
import numpy as np
level = min(max(int(resolution_level), 0), 9)
@@ -159,4 +200,4 @@ def infer(session, image, resolution_level, *, release_memory=True):
)
raw_output = dict(zip(names, values))
raw_output["metric_scale"] = raw_output.pop(scale_name)
return postprocess(raw_output)
return postprocess(raw_output, fov_x=fov_x, source_valid=source_valid)

159
src/anyimage/server/models/onnx_moge3.py Normal file
View File

@@ -0,0 +1,159 @@
"""MoGe-3 ONNX backbone and sparse refinement inference."""
from dataclasses import dataclass
from itertools import product
from pathlib import Path
from .onnx_moge2 import MAX_TOKENS, MIN_TOKENS, postprocess
from .onnx_runtime import create_session as create_runtime_session
from .onnx_runtime import run_session
@dataclass
class Moge3Session:
backbone: object
refiner: object
def create_session(model_directory, requested_device="auto"):
directory = Path(model_directory)
backbone = create_runtime_session(directory / "backbone.onnx", requested_device)
refiner = create_runtime_session(directory / "refiner.onnx", requested_device)
return Moge3Session(backbone, refiner)
def build_sparse_inputs(coord):
"""Build dynamic convolution, pooling and upsampling indices for five levels."""
import numpy as np
height, width = coord.shape[1:3]
z = np.rint(coord[0, ..., 2] * 256).astype(np.int64)
z -= z.min()
y, x = np.indices((height, width))
coords = np.stack((y, x, z), axis=-1).reshape(-1, 3)
neighbors, pools, parents = [], [], []
for level in range(5):
maximum = coords.max(axis=0)
dimensions = maximum + 3
def keys(values):
return (
((values[:, 0] + 1) * dimensions[1] + values[:, 1] + 1) * dimensions[2]
+ values[:, 2]
+ 1
)
order = np.argsort(keys(coords))
sorted_keys = keys(coords)[order]
indices = []
for offset in product((-1, 0, 1), repeat=3):
query = coords + offset
query_keys = keys(query)
positions = np.minimum(
np.searchsorted(sorted_keys, query_keys), len(coords) - 1
)
valid = (
(sorted_keys[positions] == query_keys)
& (query >= 0).all(axis=1)
& (query <= maximum).all(axis=1)
)
indices.append(np.where(valid, order[positions], len(coords)))
neighbors.append(np.stack(indices, axis=1).astype(np.int64))
if level < 4:
coarse, inverse = np.unique(coords // 2, axis=0, return_inverse=True)
children = np.full((len(coarse), 8), len(coords), dtype=np.int64)
slot = (coords[:, 0] % 2) * 4 + (coords[:, 1] % 2) * 2 + coords[:, 2] % 2
children[inverse, slot] = np.arange(len(coords))
pools.append(children)
parents.append(inverse.astype(np.int64))
coords = coarse
encoder_index = (coords[:, 0] * (width // 16) + coords[:, 1]).astype(np.int64)
return {
**{f"neighbor_{i}": value for i, value in enumerate(neighbors)},
**{f"pool_{i}": value for i, value in enumerate(pools)},
**{f"up_{i}": value for i, value in enumerate(parents)},
"encoder_index": encoder_index,
}
def _resize(array, height, width):
"""Resize an HWC array with half-pixel bilinear sampling."""
import numpy as np
from scipy.ndimage import map_coordinates
y = (np.arange(height) + 0.5) * array.shape[0] / height - 0.5
x = (np.arange(width) + 0.5) * array.shape[1] / width - 0.5
coordinates = np.meshgrid(y, x, indexing="ij")
if array.ndim == 2:
return map_coordinates(
array, coordinates, order=1, mode="nearest", prefilter=False
)
return np.stack(
[
map_coordinates(
array[..., channel],
coordinates,
order=1,
mode="nearest",
prefilter=False,
)
for channel in range(array.shape[-1])
],
axis=-1,
)
def infer(
session,
image,
resolution_level,
*,
release_memory=True,
fov_x=None,
source_valid=None,
):
import numpy as np
level = min(max(int(resolution_level), 0), 9)
tokens = int(MIN_TOKENS + level / 9.0 * (MAX_TOKENS - MIN_TOKENS))
coord, encoder, normal, mask, scale = run_session(
session.backbone,
["coord", "encoder", "normal", "mask", "scale"],
{
"image": np.ascontiguousarray(
image.transpose(2, 0, 1)[None], dtype=np.float32
),
"num_tokens": np.asarray(tokens, dtype=np.int64),
},
release_memory=release_memory,
)
for step in range(3):
if not np.isfinite(coord).all():
raise ValueError("MoGe-3 produced non-finite coordinates")
inputs = build_sparse_inputs(coord)
inputs.update(coord=coord, encoder=encoder)
(delta,) = run_session(
session.refiner,
["delta"],
inputs,
release_memory=release_memory and step == 2,
)
coord[..., 2] += delta
del inputs, delta
height, width = image.shape[:2]
coord = _resize(coord[0], height, width)
depth = np.exp(coord[..., 2:])
points = np.concatenate((coord[..., :2] * depth, depth), axis=-1)
normal = _resize(normal[0].transpose(1, 2, 0), height, width)
normal /= np.maximum(np.linalg.norm(normal, axis=-1, keepdims=True), 1e-12)
mask = 1.0 / (1.0 + np.exp(-np.clip(_resize(mask[0, 0], height, width), -80, 80)))
return postprocess(
{
"points": points[None],
"normal": normal[None],
"mask": mask[None],
"metric_scale": np.exp(scale),
},
fov_x=fov_x,
source_valid=source_valid,
)

3
src/anyimage/server/models/upscale.py
View File

@@ -1,3 +1,4 @@
from ..media.input import open_image
from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
@@ -27,7 +28,7 @@ def infer_one(
parameters["device"],
)
try:
with Image.open(image_path) as opened:
with open_image(image_path) as opened:
source = limit_image(
opened,
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),