nan
d08d0e9c1a
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency. - Added tests to ensure proper host termination and command execution in `test_addon.py`. - Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks. - Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting. - Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints. - Bumped version in `uv.lock` to 0.9.30.
747 lines
22 KiB
Python
747 lines
22 KiB
Python
import argparse
|
||
import json
|
||
import os
|
||
import shutil
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
|
||
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")
|
||
HF_ENDPOINT = "https://huggingface.co/"
|
||
HF_MIRRORS = ("https://hf-mirror.com",)
|
||
HF_PROBE_TIMEOUT = 5.0
|
||
|
||
|
||
class TaskCancelled(RuntimeError):
|
||
pass
|
||
|
||
|
||
def check_cancelled(cancel_check):
|
||
if cancel_check is not None and cancel_check():
|
||
raise TaskCancelled("Task cancelled")
|
||
|
||
|
||
def reachable(url, timeout=HF_PROBE_TIMEOUT):
|
||
request = urllib.request.Request(
|
||
url,
|
||
method="HEAD",
|
||
headers={"User-Agent": "SnapForm-Runtime-Installer"},
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout):
|
||
return True
|
||
except urllib.error.HTTPError:
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def resolve_hf_endpoint():
|
||
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
|
||
configured = os.environ.get("HF_ENDPOINT")
|
||
if configured:
|
||
return configured, False
|
||
if reachable(HF_ENDPOINT):
|
||
return HF_ENDPOINT, False
|
||
for mirror in HF_MIRRORS:
|
||
if reachable(mirror):
|
||
return mirror, True
|
||
return HF_ENDPOINT, False
|
||
|
||
|
||
def write_status(
|
||
path,
|
||
progress,
|
||
message,
|
||
error=None,
|
||
state=None,
|
||
stage=None,
|
||
stages=None,
|
||
stage_label=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 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 extract_video_frames(
|
||
video_path,
|
||
output_directory,
|
||
maximum_frames,
|
||
cv2,
|
||
cancel_check=None,
|
||
):
|
||
capture = cv2.VideoCapture(str(video_path))
|
||
if not capture.isOpened():
|
||
raise RuntimeError(f"Unable to open video: {video_path}")
|
||
|
||
output_directory = Path(output_directory)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
if total_frames <= 0:
|
||
capture.release()
|
||
raise RuntimeError("The video does not report a readable frame count")
|
||
|
||
import numpy as np
|
||
|
||
sample_count = min(maximum_frames, total_frames)
|
||
frame_indices = np.linspace(
|
||
0,
|
||
total_frames - 1,
|
||
sample_count,
|
||
dtype=np.int64,
|
||
)
|
||
image_paths = []
|
||
try:
|
||
for output_index, frame_index in enumerate(frame_indices):
|
||
check_cancelled(cancel_check)
|
||
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
|
||
success, frame = capture.read()
|
||
if not success:
|
||
raise RuntimeError(
|
||
f"Unable to read video frame {int(frame_index)}"
|
||
)
|
||
image_path = output_directory / f"frame-{output_index: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()
|
||
|
||
return image_paths
|
||
|
||
|
||
def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
||
check_cancelled(cancel_check)
|
||
input_path = Path(input_path)
|
||
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,
|
||
cancel_check=cancel_check,
|
||
)
|
||
except Exception:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
raise
|
||
return images, temporary_directory
|
||
|
||
|
||
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
|
||
|
||
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
|
||
speed = format_rate(rate) if rate else ""
|
||
if self.stage is not None:
|
||
message = f"downloading {speed}" if speed else "downloading"
|
||
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,
|
||
)
|
||
|
||
|
||
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),
|
||
}
|
||
write_status(
|
||
args.status,
|
||
start,
|
||
"checking download source",
|
||
**stage_kwargs,
|
||
)
|
||
endpoint, used_mirror = resolve_hf_endpoint()
|
||
if endpoint != os.environ.get("HF_ENDPOINT"):
|
||
os.environ["HF_ENDPOINT"] = endpoint
|
||
from huggingface_hub import snapshot_download
|
||
|
||
message = f"downloading {label}"
|
||
if used_mirror:
|
||
message = f"{message} (mirror)"
|
||
write_status(args.status, start, message, **stage_kwargs)
|
||
progress_class = create_download_progress_class(
|
||
args.status,
|
||
args.repository,
|
||
start=start,
|
||
span=span,
|
||
label=label,
|
||
**stage_kwargs,
|
||
)
|
||
snapshot_download(
|
||
repo_id=args.repository,
|
||
local_dir=model_dir,
|
||
tqdm_class=progress_class,
|
||
allow_patterns=args.allow_pattern or None,
|
||
)
|
||
write_status(
|
||
args.status,
|
||
start + span,
|
||
"model ready",
|
||
**stage_kwargs,
|
||
)
|
||
|
||
|
||
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,
|
||
)
|
||
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 remove_background(args, model_cache=None, cancel_check=None):
|
||
total_started = time.perf_counter()
|
||
write_status(args.status, 0.02, "Starting BEN2 Runtime")
|
||
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=False)
|
||
device = select_device(args.device, torch)
|
||
|
||
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
|
||
write_status(args.status, 0.35, "Removing background")
|
||
foreground, load_ms, inference_ms = infer_foreground(
|
||
args.input,
|
||
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")
|
||
|
||
check_cancelled(cancel_check)
|
||
total_ms = (time.perf_counter() - total_started) * 1000.0
|
||
message = (
|
||
f"Background removed 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 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 == "mini_npz":
|
||
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 reconstruct(args, model_cache=None, cancel_check=None):
|
||
write_status(args.status, 0.02, "Starting AI runtime")
|
||
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=False)
|
||
write_status(args.status, 0.08, "Preparing input")
|
||
images, temporary_directory = prepare_input(
|
||
args.input,
|
||
args.video_frames,
|
||
output_directory.parent,
|
||
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",
|
||
"",
|
||
)
|
||
if remove_background_model_dir:
|
||
write_status(args.status, 0.12, "Removing image background")
|
||
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")
|
||
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}")
|
||
export(
|
||
prediction,
|
||
args.export_format,
|
||
str(output_directory),
|
||
**export_kwargs,
|
||
)
|
||
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"
|
||
)
|
||
check_cancelled(cancel_check)
|
||
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("--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,
|
||
)
|
||
|
||
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(
|
||
"--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)
|
||
else:
|
||
remove_background(args)
|
||
except Exception as error:
|
||
write_status(args.status, 1.0, "Task failed", str(error))
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|