import argparse import fnmatch import json import os import shutil import tempfile import threading import time from pathlib import Path from urllib.parse import quote, urlencode 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") HF_ENDPOINT = "https://huggingface.co" MODELSCOPE_ENDPOINT = "https://modelscope.cn" MODELSCOPE_REVISION = "master" MODELSCOPE_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" shutil.copy2( image_path, Path(output_directory) / f"frame.{index + 1:04d}{suffix}", ) 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), } from huggingface_hub import snapshot_download progress_class = create_download_progress_class( args.status, args.repository, start=start, span=span, label=label, **stage_kwargs, ) failures = [] for endpoint in download_endpoints(): if not hf_endpoint_reachable(endpoint): failures.append((endpoint, "endpoint unreachable")) continue 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) except Exception as error: failures.append((endpoint, error)) continue write_status( args.status, start + span, f"{label} ready", **stage_kwargs, ) return # Hugging Face 官方源与已配置镜像均失败时，回退 ModelScope（仓库名与 Hugging Face 一致） try: write_status( args.status, start, f"Downloading {label} | ModelScope mirror", **stage_kwargs, ) download_from_modelscope( args, model_dir, start=start, span=span, label=label, stage_kwargs=stage_kwargs, ) _verify_model_complete(model_dir) except Exception as error: failures.append((MODELSCOPE_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(): """Hugging Face download candidates: configured mirror first, official endpoint otherwise.""" endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT return [endpoint.rstrip("/")] def hf_endpoint_reachable(endpoint, timeout=10.0): """Quick connectivity probe so slow or blocked Hugging Face endpoints fail fast.""" try: request = Request( f"{endpoint.rstrip('/')}/api/models", method="HEAD", headers={"User-Agent": "SnapForm"}, ) with urlopen(request, timeout=timeout): return True except Exception: return False def _verify_model_complete(model_dir): """Refuse to report success when the required model files are missing.""" model_dir = Path(model_dir) missing = [] if not (model_dir / "config.json").is_file(): missing.append("config.json") if not any(model_dir.glob("*.safetensors")): missing.append("*.safetensors") if missing: raise RuntimeError( "Model download reported success but files are missing: " + ", ".join(missing) ) def _hidden_repo_path(path): return any(part.startswith(".") for part in Path(path).parts) def _model_file_allowed(path, patterns): if not patterns: return True return any(fnmatch.fnmatch(path, pattern) for pattern in patterns) def modelscope_file_list(repository, revision=MODELSCOPE_REVISION): """List (path, size) file entries of a ModelScope repository.""" url = ( f"{MODELSCOPE_ENDPOINT}/api/v1/models/{repository}/repo/files?" + urlencode({"Revision": revision, "Recursive": "true"}) ) request = Request(url, headers={"User-Agent": "SnapForm"}) with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response: payload = json.loads(response.read().decode("utf-8")) entries = [] for entry in payload.get("Data", {}).get("Files", []): path = entry.get("Path") if path and entry.get("Type") == "blob" and not _hidden_repo_path(path): entries.append((path, int(entry.get("Size") or 0))) return entries def _download_modelscope_file( repository, path, destination, progress=None, retries=2, ): """Download one ModelScope file atomically with a per-chunk progress callback.""" url = ( f"{MODELSCOPE_ENDPOINT}/models/{repository}/resolve/" f"{MODELSCOPE_REVISION}/{quote(path)}" ) request = Request(url, headers={"User-Agent": "SnapForm"}) last_error = None for _attempt in range(retries): temporary = destination.with_name(f"{destination.name}.part") try: with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response: with open(temporary, "wb") as output: while True: chunk = response.read(1024 * 1024) if not chunk: break output.write(chunk) if progress is not None: progress(len(chunk)) temporary.replace(destination) return except Exception as error: last_error = error try: temporary.unlink() except OSError: pass raise last_error def download_from_modelscope( args, model_dir, start=0.05, span=0.90, label=None, stage_kwargs=None, ): """Download a model from ModelScope using only standard library requests.""" repository = args.repository files = modelscope_file_list(repository) patterns = args.allow_pattern or None selected = [ (path, size) for path, size in files if _model_file_allowed(path, patterns) ] if not selected: raise RuntimeError( f"No matching files in ModelScope repository: {repository}" ) reporter = DownloadProgressReporter( args.status, repository, start=start, span=span, label=label or repository, **(stage_kwargs or {}), ) total_size = sum(size for _path, size in selected) downloaded = 0 def on_chunk(count): nonlocal downloaded downloaded += count reporter.report(downloaded, total_size) for path, _size in selected: destination = Path(model_dir) / path destination.parent.mkdir(parents=True, exist_ok=True) _download_modelscope_file( repository, path, destination, 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, ) 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 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 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=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.parent, 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.parent, 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 == "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 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=True) write_status(args.status, 0.08, "Preparing input") images, temporary_directory = prepare_input( args.input, args.video_frames, output_directory.parent, 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", "", ) 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") 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, ) 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" ) if args.export_format == "mini_npz": nested_path = ( output_directory / "exports" / "mini_npz" / "results.npz" ) result_path = output_directory / "results.npz" if nested_path.is_file(): nested_path.replace(result_path) result_paths = [result_path] for directory in (nested_path.parent, nested_path.parent.parent): try: directory.rmdir() except OSError: pass 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("--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()