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()