1105 lines
34 KiB
Python
1105 lines
34 KiB
Python
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()
|