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

Refactor image handling and update foreground processing

Browse Source 
- Changed the logic in `remove_background_faces` to keep only faces with all foreground vertices.
- Updated background job to save foreground images with a new naming convention (`color.0001.png`).
- Replaced `copy_frames` with `copy_color_images` in various jobs to ensure consistent image naming.
- Introduced a new function `copy_color_images` to handle the copying of images with a specific naming format.
- Updated the BEN2 model to save output images with a new prefix for better clarity.
- Modified tests to reflect changes in image naming and processing logic.
- Added a new shadeless node group for improved material handling in geometry nodes.
- Enhanced validation scripts to check for new image naming conventions and ensure proper functionality.
This commit is contained in:
马南
2026-08-19 01:45:59 +08:00
parent 835225d264
commit 20b653bb30
28 changed files with 663 additions and 383 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/snapform/server/geometry/grid_mesh.py
View File

@@ -69,9 +69,9 @@ def triangle_edge_stretch_values(points, triangles):
def remove_background_faces(faces, foreground_mask):
"""Keep faces with at least one foreground vertex."""
"""Keep only faces whose vertices are all foreground."""
foreground_points = foreground_mask.reshape(-1)
return faces[foreground_points[faces].any(axis=1)]
return faces[foreground_points[faces].all(axis=1)]
def depth_edge_mask(depth, threshold=12.0 / 255.0):

10
src/snapform/server/jobs/background.py
View File

@@ -1,7 +1,7 @@
import shutil
import time
from ..media import copy_frames, prepare_background_input
from ..media import prepare_background_input
from ..models.ben2 import infer_foreground, infer_foreground_frames
@@ -39,7 +39,7 @@ def run(context, parameters):
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
output_path = output_directory / "color.0001.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
@@ -58,8 +58,8 @@ def run(context, parameters):
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -73,6 +73,4 @@ def run(context, parameters):
)
print(message)
context.progress(1.0, message)
if tuple(output_directory.glob("foreground.*.png")):
return {"foreground": "."}
return {"foreground": "foreground.png"}
return {"color": "."}

7
src/snapform/server/jobs/depth.py
View File

@@ -1,4 +1,4 @@
from ..media import copy_frames
from ..media import copy_color_images
from ..models import da3, geometry_adapter, moge2
from ..outputs.depth import write_da3_depth, write_moge_depth
from .common import input_frames
@@ -36,14 +36,15 @@ def run(context, parameters):
cancel_check=context.check_cancelled,
)
write_moge_depth(context.directory, prediction)
write_optional(
color_paths = write_optional(
context,
parameters,
images,
progress_start=0.84,
progress_span=0.08,
)
copy_frames(images, context.directory)
if not color_paths:
copy_color_images(images, context.directory)
context.check_cancelled()
context.progress(1.0, f"Exported depth for {len(images)} frame(s)")
return {"depth": "results.npz"}

5
src/snapform/server/jobs/foreground.py
View File

@@ -24,7 +24,7 @@ def write_optional(
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
path = context.directory / "foreground.png"
path = context.directory / "color.0001.png"
foreground.convert("RGBA").save(
path,
format="PNG",
@@ -42,8 +42,9 @@ def write_optional(
progress_span,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
)
paths = sorted(context.directory.glob("foreground.*.png"))
paths = sorted(context.directory.glob("color.*.png"))
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"

4
src/snapform/server/jobs/frames.py
View File

@@ -1,6 +1,6 @@
import shutil
from ..media import copy_frames, prepare_input
from ..media import copy_color_images, prepare_input
def run(context, parameters):
@@ -13,7 +13,7 @@ def run(context, parameters):
cancel_check=context.check_cancelled,
)
try:
copy_frames(images, context.directory)
copy_color_images(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

33
src/snapform/server/jobs/surface.py
View File

@@ -1,15 +1,18 @@
from ..geometry import foreground_mesh_mask
from ..geometry import build_surface, foreground_mesh_mask
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from ..geometry import build_surface
from ..outputs.surface import (
write_normal_texture,
write_surface_archive,
write_surface_texture,
)
from ..outputs.surface import write_normal_texture, write_surface_archive
from .common import input_frames
from .foreground import write_optional
def inference_progress_message(model_name, index, total):
message = f"Estimating {model_name} geometry"
if total > 1:
message += f" ({index + 1}/{total})"
return message
def run(context, parameters):
context.progress(0.02, "Preparing surface input")
context.directory.mkdir(parents=True, exist_ok=True)
@@ -27,7 +30,7 @@ def run(context, parameters):
include_points=True,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating MoGe-2 geometry ({index + 1}/{total})",
inference_progress_message("MoGe-2", index, total),
),
cancel_check=context.check_cancelled,
)
@@ -38,13 +41,13 @@ def run(context, parameters):
images,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating DA3 geometry ({index + 1}/{total})",
inference_progress_message("DA3", index, total),
),
cancel_check=context.check_cancelled,
)
context.check_cancelled()
frame = adapter.surface_frame(prediction)
foreground_paths = write_optional(
color_paths = write_optional(
context,
parameters,
images,
@@ -52,19 +55,17 @@ def run(context, parameters):
progress_span=0.08,
)
foreground_mask = None
if foreground_paths:
if color_paths:
foreground_mask = foreground_mesh_mask(
foreground_paths[0],
color_paths[0],
frame.depth.shape,
)
context.progress(0.82, "Building surface")
mesh = build_surface(frame, foreground_mask=foreground_mask)
context.progress(0.92, "Writing surface")
write_surface_archive(mesh, context.directory / "mesh.npz")
write_surface_texture(
foreground_paths[0] if foreground_paths else images[0],
context.directory / "texture.png",
)
if not color_paths:
copy_color_images(images, context.directory)
if adapter is moge2:
write_normal_texture(
prediction.normal[0],

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

@@ -4,6 +4,7 @@ from .input import (
TaskCancelled,
check_cancelled,
collect_images,
copy_color_images,
copy_frames,
extract_video_frames,
first_frames,
@@ -18,6 +19,7 @@ __all__ = (
"TaskCancelled",
"check_cancelled",
"collect_images",
"copy_color_images",
"copy_frames",
"extract_video_frames",
"first_frames",

20
src/snapform/server/media/input.py
View File

@@ -143,16 +143,26 @@ def prepare_background_input(
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
def copy_color_images(images, output_directory):
"""Copy input images as the canonical numbered color layer."""
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
Path(output_directory) / f"color.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def copy_frames(images, output_directory):
"""Copy multiple input images as numbered raw Debug frame files."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)

5
src/snapform/server/models/ben2.py
View File

@@ -48,8 +48,9 @@ def infer_foreground_frames(
progress_span,
model_cache=None,
cancel_check=None,
output_prefix="foreground",
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
"""Run BEN2 on every frame and write numbered transparent PNG files."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
@@ -65,7 +66,7 @@ def infer_foreground_frames(
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
output_path = output_directory / f"{output_prefix}.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",

3
src/snapform/server/outputs/__init__.py
View File

@@ -1,5 +1,5 @@
from .depth import write_da3_depth, write_moge_depth
from .surface import write_normal_texture, write_surface_archive, write_surface_texture
from .surface import write_normal_texture, write_surface_archive
__all__ = (
@@ -7,5 +7,4 @@ __all__ = (
"write_moge_depth",
"write_normal_texture",
"write_surface_archive",
"write_surface_texture",
)

9
src/snapform/server/outputs/surface.py
View File

@@ -20,15 +20,6 @@ def write_surface_archive(mesh, output_path):
return output_path
def write_surface_texture(image_path, output_path):
from PIL import Image
output_path = Path(output_path)
with Image.open(image_path) as image:
image.convert("RGBA").save(output_path, format="PNG", compress_level=1)
return output_path
def write_normal_texture(normal, output_path):
import numpy as np
from PIL import Image