nan
caeb6d706c
- Introduced new enum items for depth plane and surface models with caching. - Added MOGE2_MESH_JOB to operators for converting images to surfaces. - Implemented surface mesh building and exporting in RunJob operator. - Enhanced UI to include export format options for MoGe-2. - Created validation tools for MoGe-2 surface assets and depth planes. - Updated geometry nodes to support new surface generation features. - Refactored existing tests to accommodate new functionality and ensure stability.
2059 lines
66 KiB
Python
2059 lines
66 KiB
Python
import argparse
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import shutil
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from urllib.parse import quote
|
||
from urllib.request import Request, urlopen
|
||
|
||
|
||
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
||
VIDEO_SUFFIXES = {
|
||
".mp4",
|
||
".mov",
|
||
".avi",
|
||
".mkv",
|
||
".webm",
|
||
".m4v",
|
||
".mts",
|
||
".mpg",
|
||
".mpeg",
|
||
}
|
||
SCENE_FILENAME = "scene.glb"
|
||
EXPORT_FORMATS = (
|
||
"glb",
|
||
"gs_ply",
|
||
"depth_vis",
|
||
"mini_npz",
|
||
"mesh",
|
||
"mesh_data",
|
||
)
|
||
HF_ENDPOINT = "https://huggingface.co"
|
||
R2_ENDPOINT = "https://models.omoolab.xyz"
|
||
DOWNLOAD_TIMEOUT = 60.0
|
||
|
||
|
||
class TaskCancelled(RuntimeError):
|
||
pass
|
||
|
||
|
||
def check_cancelled(cancel_check):
|
||
if cancel_check is not None and cancel_check():
|
||
raise TaskCancelled("Task cancelled")
|
||
|
||
|
||
def write_status(
|
||
path,
|
||
progress,
|
||
message,
|
||
error=None,
|
||
state=None,
|
||
stage=None,
|
||
stages=None,
|
||
stage_label=None,
|
||
eta_seconds=None,
|
||
):
|
||
payload = {"progress": progress, "message": message}
|
||
if error:
|
||
payload["error"] = error
|
||
if stage is not None:
|
||
payload["stage"] = stage
|
||
payload["stages"] = stages
|
||
payload["stage_label"] = stage_label
|
||
if eta_seconds is not None:
|
||
payload["eta_seconds"] = max(int(eta_seconds), 0)
|
||
if state is None:
|
||
if error:
|
||
state = "failed"
|
||
elif progress >= 1.0:
|
||
state = "succeeded"
|
||
else:
|
||
state = "running"
|
||
payload["state"] = state
|
||
Path(path).write_text(
|
||
json.dumps(payload, ensure_ascii=False),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def collect_images(directory):
|
||
directory = Path(directory)
|
||
return sorted(
|
||
path
|
||
for path in directory.iterdir()
|
||
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
|
||
)
|
||
|
||
|
||
def first_frames(images, maximum_frames, start_frame=1):
|
||
"""Return up to maximum_frames image paths in order from start_frame."""
|
||
if maximum_frames <= 0:
|
||
return images
|
||
start_index = max(start_frame - 1, 0)
|
||
return images[start_index:start_index + maximum_frames]
|
||
|
||
|
||
def extract_video_frames(
|
||
video_path,
|
||
output_directory,
|
||
maximum_frames,
|
||
cv2,
|
||
start_frame=1,
|
||
cancel_check=None,
|
||
):
|
||
output_directory = Path(output_directory)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
capture = cv2.VideoCapture(str(video_path))
|
||
if not capture.isOpened():
|
||
raise RuntimeError(f"Unable to open video: {video_path}")
|
||
image_paths = []
|
||
skipped = 0
|
||
try:
|
||
while len(image_paths) < maximum_frames:
|
||
check_cancelled(cancel_check)
|
||
success, frame = capture.read()
|
||
if not success:
|
||
break
|
||
if skipped < start_frame - 1:
|
||
skipped += 1
|
||
continue
|
||
image_path = (
|
||
output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
|
||
)
|
||
if not cv2.imwrite(str(image_path), frame):
|
||
raise RuntimeError(f"Unable to write video frame: {image_path}")
|
||
image_paths.append(image_path)
|
||
finally:
|
||
capture.release()
|
||
|
||
if not image_paths:
|
||
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
|
||
return image_paths
|
||
|
||
|
||
def prepare_input(
|
||
input_path,
|
||
video_frames,
|
||
temporary_root,
|
||
start_frame=1,
|
||
cancel_check=None,
|
||
):
|
||
check_cancelled(cancel_check)
|
||
input_path = 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 cv2
|
||
|
||
temporary_directory = Path(
|
||
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
|
||
)
|
||
try:
|
||
images = extract_video_frames(
|
||
input_path,
|
||
temporary_directory,
|
||
video_frames,
|
||
cv2,
|
||
start_frame=start_frame,
|
||
cancel_check=cancel_check,
|
||
)
|
||
except Exception:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
raise
|
||
return images, temporary_directory
|
||
|
||
|
||
def prepare_background_input(
|
||
input_path,
|
||
video_frames,
|
||
temporary_root,
|
||
start_frame=1,
|
||
cancel_check=None,
|
||
):
|
||
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
|
||
input_path = 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
|
||
if input_path.suffix.lower() in IMAGE_SUFFIXES:
|
||
return [input_path], None
|
||
return prepare_input(
|
||
input_path,
|
||
video_frames,
|
||
temporary_root,
|
||
start_frame=start_frame,
|
||
cancel_check=cancel_check,
|
||
)
|
||
|
||
|
||
def copy_frames(images, output_directory):
|
||
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
|
||
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)
|
||
|
||
|
||
def unproject_depth_frame(depth, intrinsics, extrinsics):
|
||
import numpy as np
|
||
|
||
height, width = depth.shape
|
||
x, y = np.meshgrid(
|
||
np.arange(width, dtype=np.float32),
|
||
np.arange(height, dtype=np.float32),
|
||
)
|
||
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
|
||
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
|
||
camera_points = rays * depth.reshape(-1, 1)
|
||
camera_points = np.column_stack(
|
||
(camera_points, np.ones(camera_points.shape[0], dtype=np.float32))
|
||
)
|
||
world_to_camera = np.eye(4, dtype=np.float32)
|
||
world_to_camera[: extrinsics.shape[0], :4] = extrinsics
|
||
world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T
|
||
return world_points[:, :3].reshape(height, width, 3)
|
||
|
||
|
||
def unproject_camera_depth(depth, intrinsics):
|
||
import numpy as np
|
||
|
||
height, width = depth.shape
|
||
x, y = np.meshgrid(
|
||
np.arange(width, dtype=np.float32),
|
||
np.arange(height, dtype=np.float32),
|
||
)
|
||
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
|
||
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
|
||
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
|
||
|
||
|
||
def smooth_depth_map(depth, smoothing):
|
||
import numpy as np
|
||
|
||
factor = min(max(float(smoothing) / 100.0, 0.0), 1.0)
|
||
if factor <= 0.0:
|
||
return np.asarray(depth, dtype=np.float32).copy()
|
||
import cv2
|
||
|
||
values = np.asarray(depth, dtype=np.float32)
|
||
valid = np.isfinite(values) & (values > 0.0)
|
||
if not valid.any():
|
||
raise ValueError("DA3 returned no valid depth values")
|
||
low, high = np.percentile(values[valid], (5.0, 95.0))
|
||
depth_range = max(float(high - low), 1e-6)
|
||
diameter = 3 + 2 * int(round(3.0 * factor))
|
||
working_values = values.copy()
|
||
working_values[~valid] = float(np.median(values[valid]))
|
||
smoothed = cv2.bilateralFilter(
|
||
working_values,
|
||
diameter,
|
||
depth_range * (0.01 + 0.09 * factor),
|
||
1.0 + 4.0 * factor,
|
||
)
|
||
smoothed[~valid] = values[~valid]
|
||
return smoothed
|
||
|
||
|
||
def compact_mesh(vertices, faces):
|
||
import numpy as np
|
||
|
||
vertices = np.asarray(vertices, dtype=np.float32)
|
||
faces = np.asarray(faces, dtype=np.int32)
|
||
used_vertices = np.unique(faces)
|
||
index_map = np.full(len(vertices), -1, dtype=np.int32)
|
||
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
|
||
vertices = vertices[used_vertices]
|
||
faces = index_map[faces]
|
||
return vertices, faces
|
||
|
||
|
||
def project_camera_uv(vertices, intrinsics, image_size):
|
||
import numpy as np
|
||
|
||
width, height = image_size
|
||
z = vertices[:, 2]
|
||
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
|
||
raise ValueError("Simplified mesh contains invalid camera depth")
|
||
u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2]
|
||
v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2]
|
||
denominator_x = max(width - 1, 1)
|
||
denominator_y = max(height - 1, 1)
|
||
return np.column_stack(
|
||
(
|
||
np.clip(u / denominator_x, 0.0, 1.0),
|
||
np.clip(1.0 - v / denominator_y, 0.0, 1.0),
|
||
)
|
||
).astype(np.float32)
|
||
|
||
|
||
def project_camera_pixels(vertices, intrinsics):
|
||
import numpy as np
|
||
|
||
z = vertices[:, 2]
|
||
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
|
||
raise ValueError("Mesh contains invalid camera depth")
|
||
return np.column_stack(
|
||
(
|
||
intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2],
|
||
intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2],
|
||
)
|
||
)
|
||
|
||
|
||
def sample_image_nearest(values, pixels):
|
||
import numpy as np
|
||
|
||
height, width = values.shape
|
||
x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1)
|
||
y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1)
|
||
return np.asarray(values)[y, x]
|
||
|
||
|
||
def confidence_percentile_map(depth, confidence):
|
||
import numpy as np
|
||
|
||
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
|
||
result = np.zeros(depth.shape, dtype=np.float32)
|
||
if not valid.any():
|
||
return result
|
||
sorted_values = np.sort(np.asarray(confidence)[valid])
|
||
ranks = np.searchsorted(
|
||
sorted_values,
|
||
np.asarray(confidence)[valid],
|
||
side="right",
|
||
)
|
||
result[valid] = ranks.astype(np.float32) / len(sorted_values)
|
||
return result
|
||
|
||
|
||
def triangle_edge_stretch_values(points, triangles):
|
||
import numpy as np
|
||
|
||
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
|
||
triangle_points = flat_points[triangles]
|
||
triangle_depth = triangle_points[..., 2]
|
||
safe_depth = np.where(
|
||
np.abs(triangle_depth) > 1e-8,
|
||
triangle_depth,
|
||
np.nan,
|
||
)
|
||
triangle_rays = triangle_points / safe_depth[..., None]
|
||
next_points = np.roll(triangle_points, -1, axis=1)
|
||
next_rays = np.roll(triangle_rays, -1, axis=1)
|
||
next_depth = np.roll(triangle_depth, -1, axis=1)
|
||
local_depth = (triangle_depth + next_depth) * 0.5
|
||
flat_edges = (
|
||
triangle_rays * local_depth[..., None]
|
||
- next_rays * local_depth[..., None]
|
||
)
|
||
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
|
||
actual_lengths = np.linalg.norm(
|
||
triangle_points - next_points,
|
||
axis=-1,
|
||
)
|
||
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
|
||
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
|
||
return stretch
|
||
|
||
|
||
def image_mesh_attributes(
|
||
vertices,
|
||
faces,
|
||
intrinsics,
|
||
depth,
|
||
confidence,
|
||
confidence_percentiles,
|
||
foreground_mask,
|
||
):
|
||
import numpy as np
|
||
|
||
pixels = project_camera_pixels(vertices, intrinsics)
|
||
point_confidence = sample_image_nearest(confidence, pixels).astype(
|
||
np.float32
|
||
)
|
||
point_confidence = np.nan_to_num(point_confidence, nan=0.0)
|
||
point_percentile = sample_image_nearest(
|
||
confidence_percentiles,
|
||
pixels,
|
||
).astype(np.float32)
|
||
point_background = ~sample_image_nearest(foreground_mask, pixels).astype(
|
||
bool
|
||
)
|
||
face_confidence = point_percentile[faces].min(axis=1)
|
||
face_edge_stretch = triangle_edge_stretch_values(
|
||
vertices,
|
||
faces,
|
||
).max(axis=1)
|
||
face_background = point_background[faces].any(axis=1)
|
||
return {
|
||
"point_depth": vertices[:, 2].astype(np.float32),
|
||
"point_confidence": point_confidence,
|
||
"point_confidence_percentile": point_percentile,
|
||
"point_background": point_background,
|
||
"face_confidence_percentile": face_confidence.astype(np.float32),
|
||
"face_edge_stretch": face_edge_stretch.astype(np.float32),
|
||
"face_background": face_background,
|
||
}
|
||
|
||
|
||
def export_prediction_attributes(prediction, output_directory):
|
||
import numpy as np
|
||
|
||
values = {}
|
||
for name in ("depth", "conf", "intrinsics", "extrinsics"):
|
||
value = getattr(prediction, name, None)
|
||
if value is not None:
|
||
values[name] = np.asarray(value)
|
||
if "depth" in values and "conf" in values:
|
||
values["confidence_percentile"] = np.stack(
|
||
[
|
||
confidence_percentile_map(depth, confidence)
|
||
for depth, confidence in zip(
|
||
values["depth"],
|
||
values["conf"],
|
||
)
|
||
]
|
||
)
|
||
output_path = Path(output_directory) / "prediction.attributes.npz"
|
||
np.savez(output_path, **values)
|
||
return output_path
|
||
|
||
|
||
def remove_background_faces(faces, foreground_mask):
|
||
"""Return only triangles whose vertices all belong to the foreground."""
|
||
foreground_points = foreground_mask.reshape(-1)
|
||
return faces[foreground_points[faces].all(axis=1)]
|
||
|
||
|
||
def export_image_mesh_data(
|
||
prediction,
|
||
image_path,
|
||
output_directory,
|
||
foreground_path=None,
|
||
progress_callback=None,
|
||
):
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
total_steps = 6
|
||
|
||
def progress(step, label):
|
||
if progress_callback is not None:
|
||
progress_callback(step, total_steps, label)
|
||
|
||
progress(1, "Validating DA3 data")
|
||
required = ("depth", "conf", "intrinsics")
|
||
missing = [
|
||
name
|
||
for name in required
|
||
if getattr(prediction, name, None) is None
|
||
]
|
||
if missing:
|
||
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
|
||
depth_frames = np.asarray(prediction.depth, dtype=np.float32)
|
||
confidence_frames = np.asarray(prediction.conf, dtype=np.float32)
|
||
if depth_frames.shape[0] != 1:
|
||
raise ValueError("Convert to Surface currently supports one image")
|
||
raw_depth = depth_frames[0]
|
||
confidence = confidence_frames[0]
|
||
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
|
||
progress(2, "Unprojecting the depth grid")
|
||
points = unproject_camera_depth(raw_depth, intrinsics)
|
||
progress(3, "Preparing confidence and background data")
|
||
foreground_mask = np.ones(raw_depth.shape, dtype=bool)
|
||
if foreground_path is not None:
|
||
foreground_mask = foreground_mesh_mask(
|
||
foreground_path,
|
||
raw_depth.shape,
|
||
)
|
||
confidence_percentiles = confidence_percentile_map(
|
||
raw_depth,
|
||
confidence,
|
||
)
|
||
progress(4, "Building source triangles")
|
||
faces = mesh_triangle_faces(
|
||
points,
|
||
raw_depth,
|
||
confidence,
|
||
None,
|
||
False,
|
||
0.0,
|
||
)
|
||
if foreground_path is not None:
|
||
faces = remove_background_faces(faces, foreground_mask)
|
||
if not len(faces):
|
||
raise ValueError(
|
||
"DA3 output does not contain any valid surface faces"
|
||
)
|
||
progress(5, "Preparing source UV and attributes")
|
||
vertices, faces = compact_mesh(
|
||
points.reshape(-1, 3),
|
||
faces,
|
||
)
|
||
height, width = raw_depth.shape
|
||
uv = project_camera_uv(vertices, intrinsics, (width, height))
|
||
attributes = image_mesh_attributes(
|
||
vertices,
|
||
faces,
|
||
intrinsics,
|
||
raw_depth,
|
||
confidence,
|
||
confidence_percentiles,
|
||
foreground_mask,
|
||
)
|
||
valid_depth = raw_depth[np.isfinite(raw_depth) & (raw_depth > 0.0)]
|
||
reference_depth = float(np.median(valid_depth))
|
||
output_directory = Path(output_directory)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
progress(6, "Writing source surface data")
|
||
result_path = output_directory / "mesh.npz"
|
||
np.savez_compressed(
|
||
result_path,
|
||
vertices=vertices,
|
||
faces=faces,
|
||
uv=uv,
|
||
intrinsics=intrinsics,
|
||
image_size=np.asarray((width, height), dtype=np.int32),
|
||
reference_depth=np.asarray(reference_depth, dtype=np.float32),
|
||
**attributes,
|
||
)
|
||
texture_path = output_directory / "texture.png"
|
||
texture_source = foreground_path or image_path
|
||
with Image.open(texture_source) as image:
|
||
image.convert("RGBA").save(texture_path, format="PNG", compress_level=1)
|
||
return result_path, texture_path
|
||
|
||
|
||
def depth_edge_mask(depth, threshold=12.0 / 255.0):
|
||
import numpy as np
|
||
|
||
values = np.asarray(depth, dtype=np.float32)
|
||
edges = np.zeros(values.shape, dtype=bool)
|
||
for axis in (0, 1):
|
||
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
|
||
second = np.take(values, range(1, values.shape[axis]), axis=axis)
|
||
valid = (
|
||
np.isfinite(first)
|
||
& np.isfinite(second)
|
||
& (first > 0.0)
|
||
& (second > 0.0)
|
||
)
|
||
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
|
||
discontinuity = valid & (np.abs(first - second) / scale > threshold)
|
||
lower = [slice(None), slice(None)]
|
||
upper = [slice(None), slice(None)]
|
||
lower[axis] = slice(0, -1)
|
||
upper[axis] = slice(1, None)
|
||
edges[tuple(lower)] |= discontinuity
|
||
edges[tuple(upper)] |= discontinuity
|
||
return edges
|
||
|
||
|
||
def edge_culling_stretch_limit(culling):
|
||
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
|
||
return 1.0 / factor
|
||
|
||
|
||
def foreground_mesh_mask(foreground_path, depth_shape):
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter
|
||
|
||
height, width = depth_shape
|
||
with Image.open(foreground_path) as source:
|
||
alpha = source.convert("RGBA").getchannel("A")
|
||
alpha = alpha.resize(
|
||
(width, height),
|
||
resample=Image.Resampling.BILINEAR,
|
||
)
|
||
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
|
||
return np.asarray(alpha) >= 128
|
||
|
||
|
||
def confidence_percentile_threshold(depth, confidence, percentile):
|
||
import numpy as np
|
||
|
||
if float(percentile) <= 0.0:
|
||
return None
|
||
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
|
||
if not valid.any():
|
||
raise ValueError("DA3 returned no valid depth confidence values")
|
||
return float(np.percentile(confidence[valid], percentile))
|
||
|
||
|
||
def mesh_triangle_faces(
|
||
points,
|
||
depth,
|
||
confidence,
|
||
confidence_threshold,
|
||
filter_edges,
|
||
max_edge_length,
|
||
edge_threshold=12.0 / 255.0,
|
||
edge_stretch_limit=None,
|
||
vertex_mask=None,
|
||
):
|
||
import numpy as np
|
||
|
||
height, width = depth.shape
|
||
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
|
||
quads = np.stack(
|
||
(
|
||
indices[:-1, :-1],
|
||
indices[1:, :-1],
|
||
indices[1:, 1:],
|
||
indices[:-1, 1:],
|
||
),
|
||
axis=-1,
|
||
).reshape(-1, 4)
|
||
valid = (
|
||
np.isfinite(points).all(axis=-1)
|
||
& np.isfinite(depth)
|
||
& (depth > 0.0)
|
||
)
|
||
if confidence_threshold is not None:
|
||
valid &= np.isfinite(confidence) & (
|
||
confidence >= confidence_threshold
|
||
)
|
||
if filter_edges:
|
||
valid &= ~depth_edge_mask(depth, edge_threshold)
|
||
if vertex_mask is not None:
|
||
if np.shape(vertex_mask) != depth.shape:
|
||
raise ValueError("Foreground mask and depth dimensions do not match")
|
||
valid &= np.asarray(vertex_mask, dtype=bool)
|
||
triangles = np.concatenate(
|
||
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
|
||
axis=0,
|
||
)
|
||
face_valid = valid.reshape(-1)[triangles].all(axis=1)
|
||
if edge_stretch_limit is not None:
|
||
face_valid &= triangle_edge_stretch_mask(
|
||
points,
|
||
triangles,
|
||
edge_stretch_limit,
|
||
)
|
||
if max_edge_length > 0.0:
|
||
flat_points = points.reshape(-1, 3)
|
||
triangle_points = flat_points[triangles]
|
||
next_points = np.roll(triangle_points, -1, axis=1)
|
||
edge_lengths = np.linalg.norm(
|
||
triangle_points - next_points,
|
||
axis=-1,
|
||
)
|
||
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
|
||
triangles = triangles[face_valid]
|
||
if not len(triangles):
|
||
return np.empty((0, 3), dtype=np.int32)
|
||
return triangles
|
||
|
||
|
||
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
|
||
stretch = triangle_edge_stretch_values(points, triangles)
|
||
tolerance = float(stretch_limit) * (1.0 + 1e-5)
|
||
return (stretch <= tolerance).all(axis=1)
|
||
|
||
|
||
def export_textured_mesh_glb(
|
||
prediction,
|
||
image_paths,
|
||
output_directory,
|
||
confidence_percentile,
|
||
filter_edges,
|
||
max_edge_length,
|
||
):
|
||
import numpy as np
|
||
import trimesh
|
||
from PIL import Image
|
||
|
||
required = ("depth", "conf", "intrinsics", "extrinsics")
|
||
missing = [name for name in required if getattr(prediction, name) is None]
|
||
if missing:
|
||
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
|
||
depth = np.asarray(prediction.depth, dtype=np.float32)
|
||
confidence = np.asarray(prediction.conf, dtype=np.float32)
|
||
threshold = confidence_percentile_threshold(
|
||
depth,
|
||
confidence,
|
||
confidence_percentile,
|
||
)
|
||
meshes = []
|
||
referenced_points = []
|
||
for index in range(depth.shape[0]):
|
||
points = unproject_depth_frame(
|
||
depth[index],
|
||
prediction.intrinsics[index],
|
||
prediction.extrinsics[index],
|
||
)
|
||
faces = mesh_triangle_faces(
|
||
points,
|
||
depth[index],
|
||
confidence[index],
|
||
threshold,
|
||
filter_edges,
|
||
max_edge_length,
|
||
)
|
||
if not len(faces):
|
||
continue
|
||
height, width = depth[index].shape
|
||
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
|
||
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
|
||
uu, vv = np.meshgrid(u, v)
|
||
uv = np.stack((uu, vv), axis=-1).reshape(-1, 2)
|
||
with Image.open(image_paths[index]) as source_image:
|
||
texture_image = source_image.convert("RGB").copy()
|
||
visual = trimesh.visual.texture.TextureVisuals(
|
||
uv=uv,
|
||
image=texture_image,
|
||
)
|
||
mesh = trimesh.Trimesh(
|
||
vertices=points.reshape(-1, 3),
|
||
faces=faces,
|
||
visual=visual,
|
||
process=False,
|
||
)
|
||
mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1)
|
||
mesh.vertex_attributes["o_confidence"] = confidence[
|
||
index
|
||
].reshape(-1)
|
||
mesh.vertex_attributes["o_confidence_percentile"] = (
|
||
confidence_percentile_map(
|
||
depth[index],
|
||
confidence[index],
|
||
).reshape(-1)
|
||
)
|
||
mesh.remove_unreferenced_vertices()
|
||
meshes.append(mesh)
|
||
referenced_points.append(mesh.vertices)
|
||
if not meshes:
|
||
raise ValueError(
|
||
"Surface filters removed every face; lower Confidence Threshold or disable filters"
|
||
)
|
||
|
||
world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64)
|
||
if world_to_camera.shape == (3, 4):
|
||
homogeneous = np.eye(4, dtype=np.float64)
|
||
homogeneous[:3, :4] = world_to_camera
|
||
world_to_camera = homogeneous
|
||
cv_to_gltf = np.eye(4, dtype=np.float64)
|
||
cv_to_gltf[1, 1] = -1.0
|
||
cv_to_gltf[2, 2] = -1.0
|
||
alignment = cv_to_gltf @ world_to_camera
|
||
aligned_points = trimesh.transform_points(
|
||
np.concatenate(referenced_points, axis=0),
|
||
alignment,
|
||
)
|
||
center = np.median(aligned_points, axis=0)
|
||
alignment[:3, 3] -= center
|
||
|
||
scene = trimesh.Scene()
|
||
for index, mesh in enumerate(meshes):
|
||
mesh.apply_transform(alignment)
|
||
scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}")
|
||
output_directory = Path(output_directory)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
output_path = output_directory / "mesh.glb"
|
||
scene.export(output_path)
|
||
return output_path
|
||
|
||
|
||
class DownloadProgressReporter:
|
||
def __init__(
|
||
self,
|
||
status_path,
|
||
repository,
|
||
clock=time.monotonic,
|
||
start=0.05,
|
||
span=0.90,
|
||
stage=None,
|
||
stages=None,
|
||
stage_label=None,
|
||
label=None,
|
||
):
|
||
self.status_path = status_path
|
||
self.repository = repository
|
||
self.clock = clock
|
||
self.lock = threading.Lock()
|
||
self.start = start
|
||
self.span = span
|
||
self.stage = stage
|
||
self.stages = stages
|
||
self.stage_label = stage_label
|
||
self.label = label
|
||
self.progress = start
|
||
self.last_write_time = 0.0
|
||
self.started_at = self.clock()
|
||
|
||
def report(self, current, total, rate=None):
|
||
if not total or current <= 0:
|
||
return
|
||
|
||
factor = min(max(current / total, 0.0), 1.0)
|
||
progress = self.start + factor * self.span
|
||
now = self.clock()
|
||
with self.lock:
|
||
if progress <= self.progress:
|
||
return
|
||
if now - self.last_write_time < 0.20 and progress < 0.95:
|
||
return
|
||
self.progress = progress
|
||
self.last_write_time = now
|
||
effective_rate = rate
|
||
elapsed = now - self.started_at
|
||
if not effective_rate and elapsed > 0:
|
||
effective_rate = current / elapsed
|
||
speed = format_rate(effective_rate) if effective_rate else ""
|
||
eta_seconds = None
|
||
if effective_rate and effective_rate > 0:
|
||
eta_seconds = max(total - current, 0) / effective_rate
|
||
if self.stage is not None:
|
||
message = (
|
||
f"Downloading packages | {speed}"
|
||
if speed
|
||
else "Downloading packages"
|
||
)
|
||
else:
|
||
prefix = self.label or self.repository
|
||
message = f"Downloading {prefix}"
|
||
if speed:
|
||
message = f"{message} | {speed}"
|
||
write_status(
|
||
self.status_path,
|
||
progress,
|
||
message,
|
||
stage=self.stage,
|
||
stages=self.stages,
|
||
stage_label=self.stage_label,
|
||
eta_seconds=eta_seconds,
|
||
)
|
||
|
||
|
||
def create_download_progress_class(
|
||
status_path,
|
||
repository,
|
||
tqdm_type=None,
|
||
clock=time.monotonic,
|
||
start=0.05,
|
||
span=0.90,
|
||
stage=None,
|
||
stages=None,
|
||
stage_label=None,
|
||
label=None,
|
||
):
|
||
if tqdm_type is None:
|
||
from tqdm.auto import tqdm as tqdm_type
|
||
|
||
reporter = DownloadProgressReporter(
|
||
status_path,
|
||
repository,
|
||
clock,
|
||
start,
|
||
span,
|
||
stage,
|
||
stages,
|
||
stage_label,
|
||
label,
|
||
)
|
||
|
||
class DownloadProgress(tqdm_type):
|
||
def __init__(self, *args, **kwargs):
|
||
description = str(kwargs.get("desc", ""))
|
||
self.snapform_reports_progress = (
|
||
description == "Downloading bytes"
|
||
or description.startswith("Fetching ")
|
||
)
|
||
super().__init__(*args, **kwargs)
|
||
|
||
def update(self, amount=1):
|
||
result = super().update(amount)
|
||
if self.snapform_reports_progress:
|
||
reporter.report(self.n, self.total, self._download_rate())
|
||
return result
|
||
|
||
def refresh(self, *args, **kwargs):
|
||
result = super().refresh(*args, **kwargs)
|
||
if self.snapform_reports_progress:
|
||
reporter.report(self.n, self.total, self._download_rate())
|
||
return result
|
||
|
||
def _download_rate(self):
|
||
try:
|
||
return getattr(self, "format_dict", {}).get("rate")
|
||
except Exception:
|
||
return None
|
||
|
||
return DownloadProgress
|
||
|
||
|
||
def progress_range(args, staged):
|
||
start = getattr(args, "progress_start", None)
|
||
end = getattr(args, "progress_end", None)
|
||
if start is not None and end is not None:
|
||
return float(start), float(end) - float(start)
|
||
if staged:
|
||
return 0.0, 1.0
|
||
return 0.05, 0.90
|
||
|
||
|
||
def download_model(args):
|
||
label = getattr(args, "label", None) or args.repository
|
||
model_dir = Path(args.model_dir)
|
||
model_dir.mkdir(parents=True, exist_ok=True)
|
||
staged = getattr(args, "stage", None) is not None
|
||
start, span = progress_range(args, staged)
|
||
stage_kwargs = {
|
||
"stage": getattr(args, "stage", None),
|
||
"stages": getattr(args, "stages", None),
|
||
"stage_label": getattr(args, "stage_label", None),
|
||
}
|
||
failures = []
|
||
try:
|
||
from huggingface_hub import snapshot_download
|
||
except Exception as error:
|
||
snapshot_download = None
|
||
failures.append((HF_ENDPOINT, error))
|
||
progress_class = None
|
||
if snapshot_download:
|
||
progress_class = create_download_progress_class(
|
||
args.status,
|
||
args.repository,
|
||
start=start,
|
||
span=span,
|
||
label=label,
|
||
**stage_kwargs,
|
||
)
|
||
for endpoint in download_endpoints() if snapshot_download else ():
|
||
os.environ["HF_ENDPOINT"] = endpoint
|
||
message = f"Downloading {label}"
|
||
if endpoint != HF_ENDPOINT:
|
||
message = f"{message} (mirror)"
|
||
write_status(args.status, start, message, **stage_kwargs)
|
||
try:
|
||
snapshot_download(
|
||
repo_id=args.repository,
|
||
local_dir=model_dir,
|
||
tqdm_class=progress_class,
|
||
allow_patterns=args.allow_pattern or None,
|
||
endpoint=endpoint,
|
||
)
|
||
_verify_model_complete(
|
||
model_dir,
|
||
getattr(args, "required_pattern", None),
|
||
)
|
||
except Exception as error:
|
||
failures.append((endpoint, error))
|
||
continue
|
||
write_status(
|
||
args.status,
|
||
start + span,
|
||
f"{label} ready",
|
||
**stage_kwargs,
|
||
)
|
||
return
|
||
# Hugging Face 失败后使用由项目维护的静态模型镜像。
|
||
r2_directory = getattr(args, "r2_directory", "")
|
||
r2_files = getattr(args, "r2_file", [])
|
||
if not r2_directory or not r2_files:
|
||
detail = "; ".join(
|
||
f"{source}: {error}" for source, error in failures
|
||
)
|
||
raise RuntimeError(f"All download sources failed: {detail}")
|
||
try:
|
||
write_status(
|
||
args.status,
|
||
start,
|
||
f"Downloading {label} | OmooLab mirror",
|
||
**stage_kwargs,
|
||
)
|
||
download_from_r2(
|
||
args,
|
||
model_dir,
|
||
start=start,
|
||
span=span,
|
||
label=label,
|
||
stage_kwargs=stage_kwargs,
|
||
)
|
||
_verify_model_complete(
|
||
model_dir,
|
||
getattr(args, "required_pattern", None),
|
||
)
|
||
except Exception as error:
|
||
failures.append((R2_ENDPOINT, error))
|
||
else:
|
||
write_status(
|
||
args.status,
|
||
start + span,
|
||
f"{label} ready",
|
||
**stage_kwargs,
|
||
)
|
||
return
|
||
detail = "; ".join(
|
||
f"{source}: {error}" for source, error in failures
|
||
)
|
||
raise RuntimeError(f"All download sources failed: {detail}")
|
||
|
||
|
||
def download_endpoints():
|
||
"""Return the configured Hugging Face endpoint or official source."""
|
||
endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT
|
||
return [endpoint.rstrip("/")]
|
||
|
||
|
||
def _verify_model_complete(model_dir, required_patterns=None):
|
||
"""Refuse to report success when the required model files are missing."""
|
||
model_dir = Path(model_dir)
|
||
required_patterns = required_patterns or (
|
||
"config.json",
|
||
"*.safetensors",
|
||
)
|
||
missing = [
|
||
pattern
|
||
for pattern in required_patterns
|
||
if not any(model_dir.glob(pattern))
|
||
]
|
||
if missing:
|
||
raise RuntimeError(
|
||
"Model download reported success but files are missing: "
|
||
+ ", ".join(missing)
|
||
)
|
||
|
||
|
||
def _download_r2_file(
|
||
directory,
|
||
filename,
|
||
destination,
|
||
expected_size,
|
||
expected_checksum,
|
||
progress=None,
|
||
retries=2,
|
||
):
|
||
"""Download and verify one OmooLab R2 model file atomically."""
|
||
url = f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}"
|
||
request = Request(url, headers={"User-Agent": "SnapForm"})
|
||
last_error = None
|
||
for _attempt in range(retries):
|
||
temporary = destination.with_name(f"{destination.name}.part")
|
||
try:
|
||
checksum = hashlib.sha256()
|
||
downloaded = 0
|
||
with urlopen(request, timeout=DOWNLOAD_TIMEOUT) as response:
|
||
with open(temporary, "wb") as output:
|
||
while True:
|
||
chunk = response.read(1024 * 1024)
|
||
if not chunk:
|
||
break
|
||
output.write(chunk)
|
||
checksum.update(chunk)
|
||
downloaded += len(chunk)
|
||
if progress is not None:
|
||
progress(len(chunk))
|
||
if downloaded != expected_size:
|
||
raise RuntimeError(
|
||
f"Size mismatch for {filename}: "
|
||
f"{downloaded} != {expected_size}"
|
||
)
|
||
if checksum.hexdigest().lower() != expected_checksum.lower():
|
||
raise RuntimeError(f"SHA-256 mismatch for {filename}")
|
||
temporary.replace(destination)
|
||
return
|
||
except Exception as error:
|
||
last_error = error
|
||
try:
|
||
temporary.unlink()
|
||
except OSError:
|
||
pass
|
||
raise last_error
|
||
|
||
|
||
def download_from_r2(
|
||
args,
|
||
model_dir,
|
||
start=0.05,
|
||
span=0.90,
|
||
label=None,
|
||
stage_kwargs=None,
|
||
):
|
||
"""Download the declared model files from OmooLab R2."""
|
||
repository = args.repository
|
||
selected = []
|
||
for value in args.r2_file:
|
||
filename, size, checksum = value.split("|", 2)
|
||
selected.append((filename, int(size), checksum))
|
||
reporter = DownloadProgressReporter(
|
||
args.status,
|
||
repository,
|
||
start=start,
|
||
span=span,
|
||
label=label or repository,
|
||
**(stage_kwargs or {}),
|
||
)
|
||
total_size = sum(size for _path, size, _checksum in selected)
|
||
downloaded = 0
|
||
|
||
def on_chunk(count):
|
||
nonlocal downloaded
|
||
downloaded += count
|
||
reporter.report(downloaded, total_size)
|
||
|
||
for path, size, checksum in selected:
|
||
destination = Path(model_dir) / path
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
_download_r2_file(
|
||
args.r2_directory,
|
||
path,
|
||
destination,
|
||
size,
|
||
checksum,
|
||
progress=on_chunk,
|
||
)
|
||
|
||
|
||
def format_rate(rate):
|
||
size = float(rate)
|
||
if size <= 0:
|
||
return ""
|
||
units = ("B/s", "KB/s", "MB/s", "GB/s")
|
||
unit_index = 0
|
||
while size >= 1024 and unit_index < len(units) - 1:
|
||
size /= 1024
|
||
unit_index += 1
|
||
if unit_index == 0:
|
||
return f"{size:.0f} {units[unit_index]}"
|
||
return f"{size:.1f} {units[unit_index]}"
|
||
|
||
|
||
def select_device(requested, torch):
|
||
if requested == "cpu":
|
||
return torch.device("cpu")
|
||
if requested == "cuda":
|
||
if not torch.cuda.is_available():
|
||
raise RuntimeError("CUDA was requested but is not available")
|
||
return torch.device("cuda")
|
||
if requested == "mps":
|
||
if not torch.backends.mps.is_available():
|
||
raise RuntimeError("MPS was requested but is not available")
|
||
return torch.device("mps")
|
||
if torch.cuda.is_available():
|
||
return torch.device("cuda")
|
||
if torch.backends.mps.is_available():
|
||
return torch.device("mps")
|
||
return torch.device("cpu")
|
||
|
||
|
||
def infer_foreground(
|
||
input_path,
|
||
model_dir,
|
||
device,
|
||
torch,
|
||
model_cache=None,
|
||
cancel_check=None,
|
||
):
|
||
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")
|
||
owns_model = model_cache is None
|
||
if owns_model:
|
||
from ben2 import BEN_Base
|
||
|
||
load_started = time.perf_counter()
|
||
model = BEN_Base.from_pretrained(model_dir)
|
||
model = model.to(device=device)
|
||
model.eval()
|
||
load_ms = (time.perf_counter() - load_started) * 1000.0
|
||
else:
|
||
model, load_ms = model_cache.get_ben2(model_dir, device)
|
||
|
||
check_cancelled(cancel_check)
|
||
with Image.open(input_path) as opened_image:
|
||
image = opened_image.convert("RGB")
|
||
if device.type == "cuda":
|
||
torch.cuda.synchronize()
|
||
inference_started = time.perf_counter()
|
||
with torch.inference_mode():
|
||
foreground = model.inference(
|
||
image,
|
||
refine_foreground=False,
|
||
)
|
||
foreground = preserve_source_rgb(image, foreground)
|
||
if device.type == "cuda":
|
||
torch.cuda.synchronize()
|
||
inference_ms = (time.perf_counter() - inference_started) * 1000.0
|
||
if owns_model:
|
||
del model
|
||
if device.type == "cuda":
|
||
torch.cuda.empty_cache()
|
||
check_cancelled(cancel_check)
|
||
return foreground, load_ms, inference_ms
|
||
|
||
|
||
def preserve_source_rgb(source, foreground):
|
||
"""Combine the untouched source RGB with BEN2's foreground alpha."""
|
||
source_rgba = source.convert("RGBA")
|
||
alpha = foreground.convert("RGBA").getchannel("A")
|
||
if alpha.size != source_rgba.size:
|
||
from PIL import Image
|
||
|
||
alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS)
|
||
source_rgba.putalpha(alpha)
|
||
return source_rgba
|
||
|
||
|
||
def infer_foreground_frames(
|
||
images,
|
||
model_dir,
|
||
device,
|
||
torch,
|
||
output_directory,
|
||
status_path,
|
||
start_progress,
|
||
progress_span,
|
||
model_cache=None,
|
||
cancel_check=None,
|
||
):
|
||
"""Run BEN2 on every frame and write foreground.0001.png etc."""
|
||
total_load_ms = 0.0
|
||
total_inference_ms = 0.0
|
||
for index, image_path in enumerate(images):
|
||
check_cancelled(cancel_check)
|
||
write_status(
|
||
status_path,
|
||
start_progress + progress_span * index / len(images),
|
||
f"Removing background ({index + 1}/{len(images)})",
|
||
)
|
||
foreground, load_ms, inference_ms = infer_foreground(
|
||
image_path,
|
||
model_dir,
|
||
device,
|
||
torch,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
output_path = output_directory / f"foreground.{index + 1:04d}.png"
|
||
foreground.convert("RGBA").save(
|
||
output_path,
|
||
format="PNG",
|
||
compress_level=1,
|
||
)
|
||
if not output_path.is_file():
|
||
raise RuntimeError(f"BEN2 did not write: {output_path}")
|
||
total_load_ms += load_ms
|
||
total_inference_ms += inference_ms
|
||
return total_load_ms, total_inference_ms
|
||
|
||
|
||
def remove_background(args, model_cache=None, cancel_check=None):
|
||
total_started = time.perf_counter()
|
||
write_status(args.status, 0.02, "Starting Environment")
|
||
check_cancelled(cancel_check)
|
||
|
||
if model_cache is None:
|
||
import torch
|
||
else:
|
||
torch = model_cache.torch
|
||
|
||
output_directory = Path(args.output_dir)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
device = select_device(args.device, torch)
|
||
|
||
video_frames = getattr(args, "video_frames", 16)
|
||
start_frame = getattr(args, "start_frame", 1)
|
||
images, temporary_directory = prepare_background_input(
|
||
args.input,
|
||
video_frames,
|
||
output_directory,
|
||
start_frame=start_frame,
|
||
cancel_check=cancel_check,
|
||
)
|
||
try:
|
||
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
|
||
if len(images) == 1:
|
||
write_status(args.status, 0.35, "Removing background")
|
||
foreground, load_ms, inference_ms = infer_foreground(
|
||
images[0],
|
||
args.model_dir,
|
||
device,
|
||
torch,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
check_cancelled(cancel_check)
|
||
write_status(args.status, 0.92, "Writing transparent PNG")
|
||
output_path = output_directory / "foreground.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,
|
||
args.model_dir,
|
||
device,
|
||
torch,
|
||
output_directory,
|
||
args.status,
|
||
0.12,
|
||
0.7,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
copy_frames(images, output_directory)
|
||
finally:
|
||
if temporary_directory is not None:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
|
||
check_cancelled(cancel_check)
|
||
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)"
|
||
)
|
||
print(message)
|
||
write_status(args.status, 1.0, message)
|
||
|
||
|
||
def extract_frames(args, cancel_check=None):
|
||
"""Copy source frames into a sequence directory without loading any AI model."""
|
||
output_directory = Path(args.output_dir)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
images, temporary_directory = prepare_input(
|
||
args.input,
|
||
getattr(args, "video_frames", 16),
|
||
output_directory,
|
||
start_frame=getattr(args, "start_frame", 1),
|
||
cancel_check=cancel_check,
|
||
)
|
||
try:
|
||
copy_frames(images, output_directory)
|
||
finally:
|
||
if temporary_directory is not None:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
message = f"Extracted {len(images)} frame(s)"
|
||
print(message)
|
||
write_status(args.status, 1.0, message)
|
||
|
||
|
||
def exported_paths(output_directory, export_format):
|
||
output_directory = Path(output_directory)
|
||
if export_format == "glb":
|
||
scene_path = output_directory / SCENE_FILENAME
|
||
return [scene_path] if scene_path.is_file() else []
|
||
if export_format == "gs_ply":
|
||
return sorted((output_directory / "gs_ply").glob("*.ply"))
|
||
if export_format == "depth_vis":
|
||
return sorted((output_directory / "depth_vis").glob("*.jpg"))
|
||
if export_format == "mesh":
|
||
result_path = output_directory / "mesh.glb"
|
||
return [result_path] if result_path.is_file() else []
|
||
if export_format == "mesh_data":
|
||
result_path = output_directory / "mesh.npz"
|
||
return [result_path] if result_path.is_file() else []
|
||
if export_format == "mini_npz":
|
||
result_path = output_directory / "results.npz"
|
||
if result_path.is_file():
|
||
return [result_path]
|
||
result_path = (
|
||
output_directory / "exports" / "mini_npz" / "results.npz"
|
||
)
|
||
return [result_path] if result_path.is_file() else []
|
||
return []
|
||
|
||
|
||
def npz_is_complete(path):
|
||
"""Return True only when the npz archive can be fully read."""
|
||
try:
|
||
import numpy as np
|
||
|
||
with np.load(path, allow_pickle=False) as archive:
|
||
return "depth" in archive.files
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def wait_for_exported_paths(
|
||
output_directory,
|
||
export_format,
|
||
timeout=30.0,
|
||
cancel_check=None,
|
||
):
|
||
# DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
|
||
# 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
|
||
deadline = time.monotonic() + timeout
|
||
while time.monotonic() < deadline:
|
||
check_cancelled(cancel_check)
|
||
result_paths = exported_paths(output_directory, export_format)
|
||
if result_paths:
|
||
if (
|
||
export_format in ("mini_npz", "npz")
|
||
and not npz_is_complete(result_paths[0])
|
||
):
|
||
time.sleep(0.1)
|
||
continue
|
||
return result_paths
|
||
time.sleep(0.05)
|
||
return []
|
||
|
||
|
||
def move_exported_file(
|
||
source,
|
||
destination,
|
||
timeout=10.0,
|
||
cancel_check=None,
|
||
):
|
||
"""Move an async export after Windows releases its file handle."""
|
||
source = Path(source)
|
||
destination = Path(destination)
|
||
deadline = time.monotonic() + timeout
|
||
while True:
|
||
check_cancelled(cancel_check)
|
||
try:
|
||
source.replace(destination)
|
||
return destination
|
||
except OSError as error:
|
||
is_file_lock = isinstance(error, PermissionError) or getattr(
|
||
error,
|
||
"winerror",
|
||
None,
|
||
) in {32, 33}
|
||
if not is_file_lock or time.monotonic() >= deadline:
|
||
raise
|
||
time.sleep(0.1)
|
||
|
||
|
||
def flatten_exported_paths(
|
||
output_directory,
|
||
result_paths,
|
||
cancel_check=None,
|
||
):
|
||
"""Move files produced in exporter subdirectories to the job root."""
|
||
output_directory = Path(output_directory)
|
||
flattened_paths = []
|
||
for source in result_paths:
|
||
source = Path(source)
|
||
check_cancelled(cancel_check)
|
||
if source.parent == output_directory:
|
||
flattened_paths.append(source)
|
||
continue
|
||
destination = output_directory / source.name
|
||
if destination.exists():
|
||
raise RuntimeError(f"Export filename already exists: {destination.name}")
|
||
flattened_paths.append(
|
||
move_exported_file(
|
||
source,
|
||
destination,
|
||
cancel_check=cancel_check,
|
||
)
|
||
)
|
||
for directory in sorted(
|
||
{path.parent for path in map(Path, result_paths)},
|
||
key=lambda path: len(path.parts),
|
||
reverse=True,
|
||
):
|
||
while directory != output_directory:
|
||
try:
|
||
directory.rmdir()
|
||
except OSError:
|
||
break
|
||
directory = directory.parent
|
||
return flattened_paths
|
||
|
||
|
||
def _save_moge_preview(output_directory, index, depth, normal, mask):
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
suffix = "" if index is None else f".{index + 1:04d}"
|
||
valid = mask & np.isfinite(depth) & (depth > 0)
|
||
preview = np.zeros(depth.shape, dtype=np.uint8)
|
||
if np.any(valid):
|
||
low, high = np.percentile(depth[valid], (2.0, 98.0))
|
||
span = max(float(high - low), 1e-6)
|
||
preview[valid] = np.clip(
|
||
(depth[valid] - low) / span * 255.0,
|
||
0.0,
|
||
255.0,
|
||
).astype(np.uint8)
|
||
Image.fromarray(preview, mode="L").save(
|
||
output_directory / f"depth{suffix}.png"
|
||
)
|
||
normal_rgb = np.clip(
|
||
(normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0)
|
||
* 127.5,
|
||
0.0,
|
||
255.0,
|
||
).astype(np.uint8)
|
||
Image.fromarray(normal_rgb, mode="RGB").save(
|
||
output_directory / f"normal{suffix}.png"
|
||
)
|
||
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
|
||
output_directory / f"mask{suffix}.png"
|
||
)
|
||
|
||
|
||
def export_moge_debug_mesh(
|
||
output_directory,
|
||
points,
|
||
depth,
|
||
mask,
|
||
normal,
|
||
image,
|
||
export_format,
|
||
):
|
||
import numpy as np
|
||
import utils3d
|
||
from moge.utils.io import save_glb, save_ply
|
||
|
||
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
|
||
uv = utils3d.np.uv_map(*depth.shape)
|
||
result = utils3d.np.build_mesh_from_map(
|
||
points,
|
||
image,
|
||
uv,
|
||
normal,
|
||
mask=clean_mask,
|
||
tri=True,
|
||
)
|
||
faces, vertices, colors, vertex_uvs, vertex_normals = result
|
||
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
|
||
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||
if export_format == "glb":
|
||
save_glb(
|
||
output_directory / "mesh.glb",
|
||
vertices,
|
||
faces,
|
||
vertex_uvs,
|
||
(image * 255.0).astype(np.uint8),
|
||
vertex_normals,
|
||
)
|
||
return
|
||
if export_format == "ply":
|
||
save_ply(
|
||
output_directory / "pointcloud.ply",
|
||
vertices,
|
||
np.zeros((0, 3), dtype=np.int32),
|
||
colors,
|
||
vertex_normals,
|
||
)
|
||
|
||
|
||
def export_moge_surface_archive(
|
||
output_path,
|
||
points,
|
||
depth,
|
||
mask,
|
||
normal,
|
||
intrinsics,
|
||
):
|
||
import numpy as np
|
||
import utils3d
|
||
|
||
height, width = depth.shape
|
||
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
|
||
result = utils3d.np.build_mesh_from_map(
|
||
points,
|
||
utils3d.np.uv_map(height, width),
|
||
normal,
|
||
mask=clean_mask,
|
||
tri=True,
|
||
)
|
||
faces, vertices, vertex_uvs, vertex_normals = result
|
||
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
|
||
pixel_intrinsics = intrinsics.copy()
|
||
pixel_intrinsics[0] *= width
|
||
pixel_intrinsics[1] *= height
|
||
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
|
||
point_depth = vertices[:, 2].astype(np.float32)
|
||
np.savez(
|
||
output_path,
|
||
vertices=vertices.astype(np.float32),
|
||
faces=faces.astype(np.int32),
|
||
uv=vertex_uvs.astype(np.float32),
|
||
intrinsics=pixel_intrinsics.astype(np.float32),
|
||
image_size=np.array((width, height), dtype=np.int32),
|
||
reference_depth=np.float32(np.median(point_depth)),
|
||
point_depth=point_depth,
|
||
vertex_normal=vertex_normals.astype(np.float32),
|
||
face_mask=np.ones(len(faces), dtype=bool),
|
||
)
|
||
|
||
|
||
def run_moge(args, model_cache=None, cancel_check=None):
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
prepare_moge_runtime()
|
||
write_status(args.status, 0.02, "Starting MoGe-2")
|
||
if model_cache is None:
|
||
import torch
|
||
from moge.model.v2 import MoGeModel
|
||
else:
|
||
torch = model_cache.torch
|
||
output_directory = Path(args.output_dir)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
images, temporary_directory = prepare_input(
|
||
args.input,
|
||
args.video_frames,
|
||
output_directory,
|
||
start_frame=getattr(args, "start_frame", 1),
|
||
cancel_check=cancel_check,
|
||
)
|
||
try:
|
||
device = select_device(args.device, torch)
|
||
write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}")
|
||
if model_cache is None:
|
||
checkpoint = Path(args.model_dir) / "model.pt"
|
||
model = MoGeModel.from_pretrained(checkpoint)
|
||
model = model.to(device=device)
|
||
model.eval()
|
||
else:
|
||
model = model_cache.get_moge(args.model_dir, device)
|
||
|
||
depths = []
|
||
normals = []
|
||
masks = []
|
||
points = []
|
||
intrinsics = []
|
||
source_images = []
|
||
for index, image_path in enumerate(images):
|
||
check_cancelled(cancel_check)
|
||
write_status(
|
||
args.status,
|
||
0.15 + 0.68 * index / len(images),
|
||
f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})",
|
||
)
|
||
with Image.open(image_path) as opened:
|
||
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
|
||
image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device)
|
||
source_images.append(image)
|
||
with torch.inference_mode():
|
||
prediction = model.infer(
|
||
image_tensor,
|
||
resolution_level=int(args.resolution_level),
|
||
force_projection=True,
|
||
apply_mask=False,
|
||
use_fp16=device.type == "cuda",
|
||
)
|
||
depth = prediction["depth"].detach().float().cpu().numpy()
|
||
normal = prediction["normal"].detach().float().cpu().numpy()
|
||
mask = prediction["mask"].detach().cpu().numpy().astype(bool)
|
||
point = prediction["points"].detach().float().cpu().numpy()
|
||
intrinsic = prediction["intrinsics"].detach().float().cpu().numpy()
|
||
depths.append(depth)
|
||
normals.append(normal)
|
||
masks.append(mask)
|
||
intrinsics.append(intrinsic)
|
||
if bool(getattr(args, "include_points", False)):
|
||
points.append(point)
|
||
debug_format = getattr(args, "debug_export_format", "")
|
||
if not debug_format or debug_format == "maps":
|
||
_save_moge_preview(
|
||
output_directory,
|
||
None if len(images) == 1 else index,
|
||
depth,
|
||
normal,
|
||
mask,
|
||
)
|
||
|
||
shapes = {array.shape for array in depths}
|
||
if len(shapes) != 1:
|
||
raise RuntimeError("MoGe-2 frame sizes must match")
|
||
result = {
|
||
"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),
|
||
}
|
||
if points:
|
||
result["points"] = np.stack(points).astype(np.float32)
|
||
debug_format = getattr(args, "debug_export_format", "")
|
||
if debug_format in {"glb", "ply"} and len(images) == 1:
|
||
write_status(args.status, 0.84, f"Exporting MoGe-2 {debug_format.upper()}")
|
||
export_moge_debug_mesh(
|
||
output_directory,
|
||
points[0],
|
||
depths[0],
|
||
masks[0],
|
||
normals[0],
|
||
source_images[0],
|
||
debug_format,
|
||
)
|
||
|
||
remove_background_model_dir = getattr(
|
||
args, "remove_background_model_dir", ""
|
||
)
|
||
if remove_background_model_dir:
|
||
write_status(args.status, 0.87, "Removing image background")
|
||
if len(images) == 1:
|
||
foreground, _, _ = infer_foreground(
|
||
images[0],
|
||
remove_background_model_dir,
|
||
device,
|
||
torch,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
foreground_path = output_directory / "foreground.png"
|
||
foreground.save(foreground_path)
|
||
if points:
|
||
foreground_mask = foreground_mesh_mask(
|
||
foreground_path,
|
||
result["depth"][0].shape,
|
||
)
|
||
result["mask"][0] &= foreground_mask
|
||
else:
|
||
infer_foreground_frames(
|
||
images,
|
||
remove_background_model_dir,
|
||
device,
|
||
torch,
|
||
output_directory,
|
||
args.status,
|
||
0.87,
|
||
0.08,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
if bool(getattr(args, "surface_mesh", False)):
|
||
if len(images) != 1 or not points:
|
||
raise RuntimeError("MoGe-2 Surface requires one point map")
|
||
write_status(args.status, 0.94, "Building official MoGe-2 surface")
|
||
export_moge_surface_archive(
|
||
output_directory / "mesh.npz",
|
||
points[0],
|
||
result["depth"][0],
|
||
result["mask"][0],
|
||
result["normal"][0],
|
||
result["intrinsics"][0],
|
||
)
|
||
np.savez(output_directory / "results.npz", **result)
|
||
copy_frames(images, output_directory)
|
||
write_status(
|
||
args.status,
|
||
1.0,
|
||
f"Exported MoGe-2 data for {len(images)} frame(s)",
|
||
)
|
||
finally:
|
||
if temporary_directory is not None:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
|
||
|
||
def prepare_moge_runtime():
|
||
"""Use PyTorch SDPA because xFormers does not support every CUDA GPU."""
|
||
os.environ["XFORMERS_DISABLED"] = "1"
|
||
|
||
|
||
def reconstruct(args, model_cache=None, cancel_check=None):
|
||
write_status(args.status, 0.02, "Starting Environment")
|
||
check_cancelled(cancel_check)
|
||
|
||
if model_cache is None:
|
||
import torch
|
||
from depth_anything_3.api import DepthAnything3
|
||
else:
|
||
torch = model_cache.torch
|
||
from depth_anything_3.utils.export import export
|
||
|
||
output_directory = Path(args.output_dir)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
write_status(args.status, 0.08, "Preparing input")
|
||
images, temporary_directory = prepare_input(
|
||
args.input,
|
||
args.video_frames,
|
||
output_directory,
|
||
start_frame=getattr(args, "start_frame", 1),
|
||
cancel_check=cancel_check,
|
||
)
|
||
|
||
try:
|
||
check_cancelled(cancel_check)
|
||
device = select_device(args.device, torch)
|
||
remove_background_model_dir = getattr(
|
||
args,
|
||
"remove_background_model_dir",
|
||
"",
|
||
)
|
||
foreground_images = []
|
||
if remove_background_model_dir:
|
||
write_status(args.status, 0.12, "Removing image background")
|
||
if len(images) == 1:
|
||
foreground, load_ms, inference_ms = infer_foreground(
|
||
images[0],
|
||
remove_background_model_dir,
|
||
device,
|
||
torch,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
check_cancelled(cancel_check)
|
||
foreground.convert("RGBA").save(
|
||
output_directory / "foreground.png",
|
||
format="PNG",
|
||
compress_level=1,
|
||
)
|
||
if not (output_directory / "foreground.png").is_file():
|
||
raise RuntimeError("BEN2 did not write a foreground image")
|
||
foreground_images = [output_directory / "foreground.png"]
|
||
else:
|
||
load_ms, inference_ms = infer_foreground_frames(
|
||
images,
|
||
remove_background_model_dir,
|
||
device,
|
||
torch,
|
||
output_directory,
|
||
args.status,
|
||
0.12,
|
||
0.05,
|
||
model_cache=model_cache,
|
||
cancel_check=cancel_check,
|
||
)
|
||
foreground_images = sorted(
|
||
output_directory.glob("foreground.*.png")
|
||
)
|
||
if remove_background_model_dir:
|
||
print(
|
||
f"BEN2 inference {inference_ms:.0f} ms; "
|
||
f"model load {load_ms:.0f} ms"
|
||
)
|
||
write_status(args.status, 0.15, f"Loading model on {device.type}")
|
||
if model_cache is None:
|
||
model = DepthAnything3.from_pretrained(args.model_dir)
|
||
model = model.to(device=device)
|
||
model.eval()
|
||
else:
|
||
model = model_cache.get_da3(args.model_dir, device)
|
||
|
||
check_cancelled(cancel_check)
|
||
write_status(
|
||
args.status,
|
||
0.25,
|
||
f"Estimating depth for {len(images)} frame(s)",
|
||
)
|
||
export_kwargs = {}
|
||
if args.export_format == "glb":
|
||
export_kwargs["glb"] = {
|
||
"conf_thresh_percentile": args.conf_thresh_percentile,
|
||
"num_max_points": args.num_max_points,
|
||
"show_cameras": bool(args.show_cameras),
|
||
"export_depth_vis": False,
|
||
}
|
||
if args.export_format == "gs_ply":
|
||
export_kwargs["gs_ply"] = {
|
||
"gs_views_interval": args.gs_views_interval,
|
||
}
|
||
with torch.inference_mode():
|
||
prediction = model.inference(
|
||
[str(path) for path in images],
|
||
infer_gs=args.export_format == "gs_ply",
|
||
process_res=args.process_res,
|
||
process_res_method=args.process_res_method,
|
||
use_ray_pose=bool(args.use_ray_pose),
|
||
ref_view_strategy=args.ref_view_strategy,
|
||
)
|
||
check_cancelled(cancel_check)
|
||
write_status(args.status, 0.88, f"Exporting {args.export_format}")
|
||
if args.export_format == "mesh":
|
||
export_textured_mesh_glb(
|
||
prediction,
|
||
images,
|
||
output_directory,
|
||
getattr(args, "conf_thresh_percentile", 40.0),
|
||
bool(getattr(args, "mesh_filter_edges", True)),
|
||
float(getattr(args, "mesh_max_edge_length", 0.1)),
|
||
)
|
||
elif args.export_format == "mesh_data":
|
||
def mesh_progress(step, total, label):
|
||
write_status(
|
||
args.status,
|
||
0.88 + 0.09 * step / total,
|
||
f"Preparing surface data {step}/{total}: {label}",
|
||
)
|
||
|
||
export_image_mesh_data(
|
||
prediction,
|
||
images[0],
|
||
output_directory,
|
||
foreground_path=(
|
||
foreground_images[0] if foreground_images else None
|
||
),
|
||
progress_callback=mesh_progress,
|
||
)
|
||
else:
|
||
export(
|
||
prediction,
|
||
args.export_format,
|
||
str(output_directory),
|
||
**export_kwargs,
|
||
)
|
||
if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
|
||
write_status(
|
||
args.status,
|
||
0.97,
|
||
"Writing structured debug attributes",
|
||
)
|
||
export_prediction_attributes(prediction, output_directory)
|
||
check_cancelled(cancel_check)
|
||
|
||
result_paths = wait_for_exported_paths(
|
||
output_directory,
|
||
args.export_format,
|
||
cancel_check=cancel_check,
|
||
)
|
||
if not result_paths:
|
||
raise RuntimeError(
|
||
f"No {args.export_format} files were exported"
|
||
)
|
||
result_paths = flatten_exported_paths(
|
||
output_directory,
|
||
result_paths,
|
||
cancel_check=cancel_check,
|
||
)
|
||
check_cancelled(cancel_check)
|
||
copy_frames(images, output_directory)
|
||
write_status(
|
||
args.status,
|
||
1.0,
|
||
f"Exported {len(result_paths)} {args.export_format} file(s)",
|
||
)
|
||
finally:
|
||
if temporary_directory is not None:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
|
||
|
||
def parser():
|
||
root = argparse.ArgumentParser()
|
||
commands = root.add_subparsers(dest="command", required=True)
|
||
|
||
download = commands.add_parser("download")
|
||
download.add_argument("--repository", required=True)
|
||
download.add_argument("--model-dir", required=True)
|
||
download.add_argument("--status", required=True)
|
||
download.add_argument("--allow-pattern", action="append", default=[])
|
||
download.add_argument("--required-pattern", action="append", default=[])
|
||
download.add_argument("--r2-directory", default="")
|
||
download.add_argument("--r2-file", action="append", default=[])
|
||
download.add_argument("--stage", type=int, default=None)
|
||
download.add_argument("--stages", type=int, default=None)
|
||
download.add_argument("--stage-label", default=None)
|
||
download.add_argument("--label", default=None)
|
||
download.add_argument("--progress-start", type=float, default=None)
|
||
download.add_argument("--progress-end", type=float, default=None)
|
||
|
||
background = commands.add_parser("remove-background")
|
||
background.add_argument("--input", required=True)
|
||
background.add_argument("--model-dir", required=True)
|
||
background.add_argument("--output-dir", required=True)
|
||
background.add_argument("--status", required=True)
|
||
background.add_argument(
|
||
"--device",
|
||
choices=("auto", "cuda", "mps", "cpu"),
|
||
required=True,
|
||
)
|
||
background.add_argument("--video-frames", type=int, default=16)
|
||
background.add_argument("--start-frame", type=int, default=1)
|
||
|
||
frames = commands.add_parser("extract-frames")
|
||
frames.add_argument("--input", required=True)
|
||
frames.add_argument("--output-dir", required=True)
|
||
frames.add_argument("--status", required=True)
|
||
frames.add_argument("--video-frames", type=int, default=16)
|
||
frames.add_argument("--start-frame", type=int, default=1)
|
||
|
||
reconstruction = commands.add_parser("reconstruct")
|
||
reconstruction.add_argument("--input", required=True)
|
||
reconstruction.add_argument("--model-dir", required=True)
|
||
reconstruction.add_argument("--output-dir", required=True)
|
||
reconstruction.add_argument(
|
||
"--export-format",
|
||
choices=EXPORT_FORMATS,
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument("--status", required=True)
|
||
reconstruction.add_argument("--remove-background-model-dir", default="")
|
||
reconstruction.add_argument("--process-res", type=int, required=True)
|
||
reconstruction.add_argument(
|
||
"--process-res-method",
|
||
choices=("upper_bound_resize", "lower_bound_resize"),
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument(
|
||
"--conf-thresh-percentile",
|
||
type=float,
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument("--num-max-points", type=int, required=True)
|
||
reconstruction.add_argument(
|
||
"--show-cameras",
|
||
type=int,
|
||
choices=(0, 1),
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument(
|
||
"--gs-views-interval",
|
||
type=int,
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument("--video-frames", type=int, required=True)
|
||
reconstruction.add_argument("--start-frame", type=int, default=1)
|
||
reconstruction.add_argument(
|
||
"--device",
|
||
choices=("auto", "cuda", "mps", "cpu"),
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument(
|
||
"--use-ray-pose",
|
||
type=int,
|
||
choices=(0, 1),
|
||
required=True,
|
||
)
|
||
reconstruction.add_argument(
|
||
"--ref-view-strategy",
|
||
choices=("first", "middle", "saddle_balanced", "saddle_sim_range"),
|
||
required=True,
|
||
)
|
||
return root
|
||
|
||
|
||
def main():
|
||
args = parser().parse_args()
|
||
try:
|
||
if args.command == "download":
|
||
download_model(args)
|
||
elif args.command == "reconstruct":
|
||
reconstruct(args)
|
||
elif args.command == "extract-frames":
|
||
extract_frames(args)
|
||
else:
|
||
remove_background(args)
|
||
except Exception as error:
|
||
write_status(args.status, 1.0, "Task failed", str(error))
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|