nan
489003c2cc
- 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.
82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
import time
|
|
|
|
from ..media import input_frames
|
|
from ..media.input import open_image
|
|
from ..models.ben2 import infer_alpha, infer_foreground, infer_foreground_frames
|
|
|
|
|
|
def run(context, parameters):
|
|
"""Run BEN2 and write transparent foreground images into the Job directory."""
|
|
total_started = time.perf_counter()
|
|
context.progress(0.02, "Starting Environment")
|
|
context.check_cancelled()
|
|
|
|
output_directory = context.directory
|
|
output_directory.mkdir(parents=True, exist_ok=True)
|
|
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 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],
|
|
model_directory,
|
|
device,
|
|
model_cache=model_manager,
|
|
cancel_check=context.check_cancelled,
|
|
preserve_input_alpha=True,
|
|
)
|
|
context.check_cancelled()
|
|
context.progress(0.92, "Writing transparent PNG")
|
|
output_path = output_directory / "image.0001.png"
|
|
foreground.convert("RGBA").save(
|
|
output_path,
|
|
format="PNG",
|
|
compress_level=1,
|
|
)
|
|
if not output_path.is_file():
|
|
raise RuntimeError("BEN2 did not write a foreground image")
|
|
else:
|
|
load_ms, inference_ms = infer_foreground_frames(
|
|
images,
|
|
model_directory,
|
|
device,
|
|
output_directory,
|
|
context.progress,
|
|
0.12,
|
|
0.7,
|
|
model_cache=model_manager,
|
|
cancel_check=context.check_cancelled,
|
|
output_prefix="image",
|
|
preserve_input_alpha=True,
|
|
)
|
|
|
|
context.check_cancelled()
|
|
total_ms = (time.perf_counter() - total_started) * 1000.0
|
|
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
|
|
message = (
|
|
f"Background removed{frame_text} in {inference_ms:.0f} ms "
|
|
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
|
|
)
|
|
context.log(message)
|
|
context.progress(1.0, message)
|
|
return {"alpha": "alpha.npy"} if output_kind == "alpha" else {"image": "."}
|