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

feat: Add BEN2 background removal functionality and integrate into SnapForm

Browse Source 
- Introduced BEN2 model for background removal in the panel and properties.
- Updated runtime to handle background removal commands and model downloads.
- Implemented background removal logic in worker, including image processing and saving results.
- Enhanced tests to cover new background removal features and ensure proper integration.
- Updated version to 0.9.28 to reflect new features and improvements.
This commit is contained in:
马南
2026-08-03 20:18:00 +08:00
parent 69eaefa016
commit d694980da7
19 changed files with 1523 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
src/snapform/operators.py
View File

@@ -5,11 +5,19 @@ from pathlib import Path
import bpy
from . import runtime
from .models import get_model, selected_model_key
from .models import (
BEN2_MODEL,
get_downloadable_model,
get_model,
selected_model_key,
)
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024
IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize"
@@ -33,10 +41,13 @@ class ActiveJob:
result_directory: Path
export_format: str = ""
source_object_name: str = ""
depth_data_name: str = ""
process_res: int = IMAGE_DEPTH_DEFAULT_PROCESS_RES
subdivision_level: int = 5
delete_source: bool = True
displacement_via: str = "MODIFIER"
unshaded: bool = False
remove_background: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
@@ -83,6 +94,9 @@ class RunJob(bpy.types.Operator):
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
),
options={"HIDDEN", "SKIP_SAVE"},
)
@@ -95,6 +109,9 @@ class RunJob(bpy.types.Operator):
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_data_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_model_key: bpy.props.StringProperty(
default="BASE",
options={"HIDDEN", "SKIP_SAVE"},
@@ -131,6 +148,14 @@ class RunJob(bpy.types.Operator):
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_remove_background: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
global _active_job
@@ -173,13 +198,20 @@ class RunJob(bpy.types.Operator):
else ""
),
source_object_name=self.snapform_source_object_name,
depth_data_name=self.snapform_depth_data_name,
process_res=self.snapform_process_res,
subdivision_level=self.snapform_subdivision_level,
delete_source=self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
remove_background=self.snapform_remove_background,
temporary_input_path=(
Path(self.snapform_input_path)
if self.job_type == IMAGE_DEPTH_JOB
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}
and self.snapform_delete_input
else None
),
@@ -196,6 +228,9 @@ class RunJob(bpy.types.Operator):
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
remove_background = bool(
getattr(self, "snapform_remove_background", False)
)
if self.job_type == "INSTALL_RUNTIME":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
@@ -205,20 +240,49 @@ class RunJob(bpy.types.Operator):
if self.job_type == "DOWNLOAD_MODEL":
self._require_online_access()
model = get_model(self.model_key)
model = get_downloadable_model(self.model_key)
return (
runtime.download_model_command(model.key, paths["status"]),
paths,
)
configured_input = (
self.snapform_input_path
if self.job_type == IMAGE_DEPTH_JOB
else settings.input_path
)
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}:
configured_input = self.snapform_input_path
elif self.job_type == BACKGROUND_REMOVAL_JOB:
configured_input = settings.background_input_path
else:
configured_input = settings.input_path
input_path = Path(bpy.path.abspath(configured_input))
if not input_path.is_file():
raise ValueError("Select a valid image or video file")
if self.job_type in {
BACKGROUND_REMOVAL_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}:
if input_path.suffix.lower() not in {
".jpg",
".jpeg",
".png",
".webp",
}:
raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
if not runtime.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_command(
settings,
input_path,
paths["status"],
paths["result"],
BEN2_MODEL.key,
),
paths,
)
if self.job_type not in {RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB}:
raise RuntimeError(f"Unsupported task type: {self.job_type}")
export_format = (
@@ -236,6 +300,12 @@ class RunJob(bpy.types.Operator):
model = get_model(model_key)
if not runtime.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if (
self.job_type == IMAGE_DEPTH_JOB
and remove_background
and not runtime.model_ready(BEN2_MODEL.key)
):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
if export_format == "gs_ply" and not model.supports_gs:
raise RuntimeError(
"PLY export requires a DA3 Giant or Nested Giant model"
@@ -255,6 +325,10 @@ class RunJob(bpy.types.Operator):
else None
),
use_ray_pose=use_ray_pose,
remove_background=(
self.job_type == IMAGE_DEPTH_JOB
and remove_background
),
),
paths,
)
@@ -274,8 +348,15 @@ class RunJob(bpy.types.Operator):
return f"Starting {settings.export_format} export"
if self.job_type == IMAGE_DEPTH_JOB:
return "Starting depth estimation"
if self.job_type == BACKGROUND_REMOVAL_JOB:
return "Starting BEN2 background removal"
if self.job_type == IMAGE_BACKGROUND_JOB:
return "Starting Image Empty background removal"
if self.job_type == ORIGINAL_BACKGROUND_JOB:
return "Starting background removal"
if self.job_type == "DOWNLOAD_MODEL":
return f"Starting {get_model(self.model_key).label} download"
model = get_downloadable_model(self.model_key)
return f"Starting {model.label} download"
return labels[self.job_type]
def modal(self, context, event):
@@ -330,15 +411,18 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
if job.kind == RECONSTRUCTION_JOB:
if job.kind in {RECONSTRUCTION_JOB, BACKGROUND_REMOVAL_JOB}:
try:
open_job_directory(job.result_directory)
except Exception as open_error:
self.report({"WARNING"}, str(open_error))
self.report(
{"INFO"},
f"Exported {job.export_format} to {job.result_directory}",
)
if job.kind == RECONSTRUCTION_JOB:
message = (
f"Exported {job.export_format} to {job.result_directory}"
)
else:
message = f"Removed background to {job.result_directory}"
self.report({"INFO"}, message)
if job.kind == IMAGE_DEPTH_JOB:
try:
settings.status = "Building depth plane"
@@ -353,6 +437,34 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == ORIGINAL_BACKGROUND_JOB:
try:
settings.status = "Building depth plane"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_original_depth_plane(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == IMAGE_BACKGROUND_JOB:
try:
settings.status = "Updating Image Empty"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._apply_removed_background(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Image update failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Image background is removed"
settings.status = status.get("message", "Task complete")
settings.progress = 1.0
@@ -370,9 +482,16 @@ class RunJob(bpy.types.Operator):
result = bpy.ops.snapform.build_depth_plane(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(npz_path),
snapform_depth_data_name=job.depth_data_name,
snapform_background_path=(
str(job.result_directory / "foreground.png")
if job.remove_background
else ""
),
snapform_subdivision_level=job.subdivision_level,
snapform_delete_source=job.delete_source,
snapform_displacement_via=job.displacement_via,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the depth plane")
@@ -380,6 +499,37 @@ class RunJob(bpy.types.Operator):
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _build_original_depth_plane(self, job):
result = bpy.ops.snapform.build_original_depth_plane(
snapform_source_object_name=job.source_object_name,
snapform_background_path=str(
job.result_directory / "foreground.png"
),
snapform_subdivision_level=job.subdivision_level,
snapform_delete_source=job.delete_source,
snapform_displacement_via=job.displacement_via,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the depth plane")
undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _apply_removed_background(self, job, push_undo=True):
output_path = job.result_directory / "foreground.png"
result = bpy.ops.snapform.apply_removed_background(
snapform_source_object_name=job.source_object_name,
snapform_output_path=str(output_path),
)
if "FINISHED" not in result:
raise RuntimeError("Unable to update the Image Empty")
if not push_undo:
return
undo_result = bpy.ops.ed.undo_push(message="Remove Background")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the background removal Undo step")
def _remove_timer(self, context):
timer = getattr(self, "_timer", None)
if timer is None: