icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update SnapForm to version 0.10.10 and add technical documentation

Browse Source 
- Bump SnapForm version in uv.lock from 0.9.30 to 0.10.10.
- Introduce comprehensive technical documentation for SnapForm, detailing features, usage instructions, and development guidelines.
This commit is contained in:
马南
2026-08-08 19:05:42 +08:00
parent d08d0e9c1a
commit aee460fe7f
22 changed files with 3905 additions and 1472 deletions
Download Patch File Download Diff File Expand all files Collapse all files

576
src/snapform/worker.py
View File

@@ -1,13 +1,14 @@
import argparse
import fnmatch
import json
import os
import shutil
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import quote, urlencode
from urllib.request import Request, urlopen
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
@@ -24,9 +25,10 @@ VIDEO_SUFFIXES = {
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
HF_ENDPOINT = "https://huggingface.co/"
HF_MIRRORS = ("https://hf-mirror.com",)
HF_PROBE_TIMEOUT = 5.0
HF_ENDPOINT = "https://huggingface.co"
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
MODELSCOPE_REVISION = "master"
MODELSCOPE_TIMEOUT = 60.0
class TaskCancelled(RuntimeError):
@@ -38,34 +40,6 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def reachable(url, timeout=HF_PROBE_TIMEOUT):
request = urllib.request.Request(
url,
method="HEAD",
headers={"User-Agent": "SnapForm-Runtime-Installer"},
)
try:
with urllib.request.urlopen(request, timeout=timeout):
return True
except urllib.error.HTTPError:
return True
except Exception:
return False
def resolve_hf_endpoint():
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
configured = os.environ.get("HF_ENDPOINT")
if configured:
return configured, False
if reachable(HF_ENDPOINT):
return HF_ENDPOINT, False
for mirror in HF_MIRRORS:
if reachable(mirror):
return mirror, True
return HF_ENDPOINT, False
def write_status(
path,
progress,
@@ -106,56 +80,66 @@ def collect_images(directory):
)
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}")
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
if total_frames <= 0:
capture.release()
raise RuntimeError("The video does not report a readable frame count")
import numpy as np
sample_count = min(maximum_frames, total_frames)
frame_indices = np.linspace(
0,
total_frames - 1,
sample_count,
dtype=np.int64,
)
image_paths = []
skipped = 0
try:
for output_index, frame_index in enumerate(frame_indices):
while len(image_paths) < maximum_frames:
check_cancelled(cancel_check)
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
success, frame = capture.read()
if not success:
raise RuntimeError(
f"Unable to read video frame {int(frame_index)}"
)
image_path = output_directory / f"frame-{output_index:04d}.jpg"
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, cancel_check=None):
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
@@ -174,6 +158,7 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
temporary_directory,
video_frames,
cv2,
start_frame=start_frame,
cancel_check=cancel_check,
)
except Exception:
@@ -182,6 +167,43 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
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,
@@ -318,21 +340,8 @@ def download_model(args):
"stages": getattr(args, "stages", None),
"stage_label": getattr(args, "stage_label", None),
}
write_status(
args.status,
start,
"checking download source",
**stage_kwargs,
)
endpoint, used_mirror = resolve_hf_endpoint()
if endpoint != os.environ.get("HF_ENDPOINT"):
os.environ["HF_ENDPOINT"] = endpoint
from huggingface_hub import snapshot_download
message = f"downloading {label}"
if used_mirror:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
progress_class = create_download_progress_class(
args.status,
args.repository,
@@ -341,18 +350,213 @@ def download_model(args):
label=label,
**stage_kwargs,
)
snapshot_download(
repo_id=args.repository,
local_dir=model_dir,
tqdm_class=progress_class,
allow_patterns=args.allow_pattern or None,
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,
"model ready",
**stage_kwargs,
)
return
# Hugging Face 官方源与已配置镜像均失败时，回退 ModelScope（仓库名与 Hugging Face 一致）
try:
write_status(
args.status,
start,
f"downloading {label} (mirror: modelscope.cn)",
**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,
"model ready",
**stage_kwargs,
)
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
write_status(
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,
start + span,
"model ready",
**stage_kwargs,
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):
@@ -435,6 +639,49 @@ def infer_foreground(
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")
@@ -446,41 +693,90 @@ def remove_background(args, model_cache=None, cancel_check=None):
torch = model_cache.torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=False)
output_directory.mkdir(parents=True, exist_ok=True)
device = select_device(args.device, torch)
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
write_status(args.status, 0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
video_frames = getattr(args, "video_frames", 16)
start_frame = getattr(args, "start_frame", 1)
images, temporary_directory = prepare_background_input(
args.input,
args.model_dir,
device,
torch,
model_cache=model_cache,
video_frames,
output_directory.parent,
start_frame=start_frame,
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")
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 in {inference_ms:.0f} ms "
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":
@@ -491,7 +787,12 @@ def exported_paths(output_directory, export_format):
if export_format == "depth_vis":
return sorted((output_directory / "depth_vis").glob("*.jpg"))
if export_format == "mini_npz":
result_path = output_directory / "exports" / "mini_npz" / "results.npz"
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 []
@@ -543,12 +844,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
from depth_anything_3.utils.export import export
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=False)
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,
)
@@ -562,22 +864,36 @@ def reconstruct(args, model_cache=None, cancel_check=None):
)
if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
remove_background_model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
check_cancelled(cancel_check)
foreground.convert("RGBA").save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
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"
@@ -636,7 +952,21 @@ def reconstruct(args, model_cache=None, cancel_check=None):
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,
@@ -673,6 +1003,15 @@ def parser():
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)
@@ -709,6 +1048,7 @@ def parser():
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"),
@@ -735,6 +1075,8 @@ def main():
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: