nan
b02372c457
- Implemented RuntimeModelsPanel and DebugPanel for UI integration in Blender. - Created SnapFormPreferences to manage user settings for storage and runtime configuration. - Defined SnapFormSettings to hold various properties related to input, model selection, and export options. - Developed runtime management functions for installation, model downloading, and reconstruction commands. - Added status bar functionality to display job progress and status in the Blender UI. - Implemented worker script for handling model downloads and depth reconstruction tasks. - Introduced unit tests for download progress reporting and image depth extraction functionalities.
382 lines
12 KiB
Python
382 lines
12 KiB
Python
import argparse
|
|
import json
|
|
import shutil
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
|
VIDEO_SUFFIXES = {
|
|
".mp4",
|
|
".mov",
|
|
".avi",
|
|
".mkv",
|
|
".webm",
|
|
".m4v",
|
|
".mts",
|
|
".mpg",
|
|
".mpeg",
|
|
}
|
|
SCENE_FILENAME = "scene.glb"
|
|
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
|
|
|
|
|
|
def write_status(path, progress, message, error=None):
|
|
payload = {"progress": progress, "message": message}
|
|
if error:
|
|
payload["error"] = error
|
|
Path(path).write_text(
|
|
json.dumps(payload, ensure_ascii=False),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def collect_images(directory):
|
|
directory = Path(directory)
|
|
return sorted(
|
|
path
|
|
for path in directory.iterdir()
|
|
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
|
|
)
|
|
|
|
|
|
def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
|
|
capture = cv2.VideoCapture(str(video_path))
|
|
if not capture.isOpened():
|
|
raise RuntimeError(f"Unable to open video: {video_path}")
|
|
|
|
output_directory = Path(output_directory)
|
|
output_directory.mkdir(parents=True, exist_ok=True)
|
|
total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
if total_frames <= 0:
|
|
capture.release()
|
|
raise RuntimeError("The video does not report a readable frame count")
|
|
|
|
import numpy as np
|
|
|
|
sample_count = min(maximum_frames, total_frames)
|
|
frame_indices = np.linspace(
|
|
0,
|
|
total_frames - 1,
|
|
sample_count,
|
|
dtype=np.int64,
|
|
)
|
|
image_paths = []
|
|
try:
|
|
for output_index, frame_index in enumerate(frame_indices):
|
|
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
|
|
success, frame = capture.read()
|
|
if not success:
|
|
raise RuntimeError(
|
|
f"Unable to read video frame {int(frame_index)}"
|
|
)
|
|
image_path = output_directory / f"frame-{output_index:04d}.jpg"
|
|
if not cv2.imwrite(str(image_path), frame):
|
|
raise RuntimeError(f"Unable to write video frame: {image_path}")
|
|
image_paths.append(image_path)
|
|
finally:
|
|
capture.release()
|
|
|
|
return image_paths
|
|
|
|
|
|
def prepare_input(input_path, video_frames, temporary_root):
|
|
input_path = Path(input_path)
|
|
suffix = input_path.suffix.lower()
|
|
if suffix in IMAGE_SUFFIXES:
|
|
return [input_path], None
|
|
if suffix not in VIDEO_SUFFIXES:
|
|
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
|
|
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
|
|
|
|
import cv2
|
|
|
|
temporary_directory = Path(
|
|
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
|
|
)
|
|
try:
|
|
images = extract_video_frames(
|
|
input_path,
|
|
temporary_directory,
|
|
video_frames,
|
|
cv2,
|
|
)
|
|
except Exception:
|
|
shutil.rmtree(temporary_directory, ignore_errors=True)
|
|
raise
|
|
return images, temporary_directory
|
|
|
|
|
|
class DownloadProgressReporter:
|
|
def __init__(self, status_path, repository, clock=time.monotonic):
|
|
self.status_path = status_path
|
|
self.repository = repository
|
|
self.clock = clock
|
|
self.lock = threading.Lock()
|
|
self.progress = 0.05
|
|
self.last_write_time = 0.0
|
|
|
|
def report(self, current, total):
|
|
if not total or current <= 0:
|
|
return
|
|
|
|
factor = min(max(current / total, 0.0), 1.0)
|
|
progress = 0.05 + factor * 0.90
|
|
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
|
|
write_status(
|
|
self.status_path,
|
|
progress,
|
|
f"Downloading {self.repository}",
|
|
)
|
|
|
|
|
|
def create_download_progress_class(
|
|
status_path,
|
|
repository,
|
|
tqdm_type=None,
|
|
clock=time.monotonic,
|
|
):
|
|
if tqdm_type is None:
|
|
from tqdm.auto import tqdm as tqdm_type
|
|
|
|
reporter = DownloadProgressReporter(status_path, repository, clock)
|
|
|
|
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)
|
|
return result
|
|
|
|
def refresh(self, *args, **kwargs):
|
|
result = super().refresh(*args, **kwargs)
|
|
if self.snapform_reports_progress:
|
|
reporter.report(self.n, self.total)
|
|
return result
|
|
|
|
return DownloadProgress
|
|
|
|
|
|
def download_model(args):
|
|
from huggingface_hub import snapshot_download
|
|
|
|
model_dir = Path(args.model_dir)
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
write_status(args.status, 0.05, f"Downloading {args.repository}")
|
|
progress_class = create_download_progress_class(
|
|
args.status,
|
|
args.repository,
|
|
)
|
|
snapshot_download(
|
|
repo_id=args.repository,
|
|
local_dir=model_dir,
|
|
tqdm_class=progress_class,
|
|
)
|
|
write_status(args.status, 1.0, "Model is ready")
|
|
|
|
|
|
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")
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def exported_paths(output_directory, export_format):
|
|
output_directory = Path(output_directory)
|
|
if export_format == "glb":
|
|
scene_path = output_directory / SCENE_FILENAME
|
|
return [scene_path] if scene_path.is_file() else []
|
|
if export_format == "gs_ply":
|
|
return sorted((output_directory / "gs_ply").glob("*.ply"))
|
|
if export_format == "depth_vis":
|
|
return sorted((output_directory / "depth_vis").glob("*.jpg"))
|
|
if export_format == "mini_npz":
|
|
result_path = output_directory / "exports" / "mini_npz" / "results.npz"
|
|
return [result_path] if result_path.is_file() else []
|
|
return []
|
|
|
|
|
|
def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
|
|
# 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
result_paths = exported_paths(output_directory, export_format)
|
|
if result_paths:
|
|
return result_paths
|
|
time.sleep(0.05)
|
|
return []
|
|
|
|
|
|
def reconstruct(args):
|
|
write_status(args.status, 0.02, "Starting DA3 runtime")
|
|
|
|
import torch
|
|
from depth_anything_3.api import DepthAnything3
|
|
from depth_anything_3.utils.export import export
|
|
|
|
output_directory = Path(args.output_dir)
|
|
output_directory.mkdir(parents=True, exist_ok=False)
|
|
write_status(args.status, 0.08, "Preparing input")
|
|
images, temporary_directory = prepare_input(
|
|
args.input,
|
|
args.video_frames,
|
|
output_directory.parent,
|
|
)
|
|
|
|
try:
|
|
device = select_device(args.device, torch)
|
|
write_status(args.status, 0.15, f"Loading model on {device.type}")
|
|
model = DepthAnything3.from_pretrained(args.model_dir)
|
|
model = model.to(device=device)
|
|
model.eval()
|
|
|
|
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,
|
|
)
|
|
write_status(args.status, 0.88, f"Exporting {args.export_format}")
|
|
export(
|
|
prediction,
|
|
args.export_format,
|
|
str(output_directory),
|
|
**export_kwargs,
|
|
)
|
|
|
|
result_paths = wait_for_exported_paths(
|
|
output_directory,
|
|
args.export_format,
|
|
)
|
|
if not result_paths:
|
|
raise RuntimeError(
|
|
f"DA3 did not export any {args.export_format} files"
|
|
)
|
|
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)
|
|
|
|
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("--process-res", type=int, required=True)
|
|
reconstruction.add_argument(
|
|
"--process-res-method",
|
|
choices=("upper_bound_resize", "lower_bound_resize"),
|
|
required=True,
|
|
)
|
|
reconstruction.add_argument(
|
|
"--conf-thresh-percentile",
|
|
type=float,
|
|
required=True,
|
|
)
|
|
reconstruction.add_argument("--num-max-points", type=int, required=True)
|
|
reconstruction.add_argument(
|
|
"--show-cameras",
|
|
type=int,
|
|
choices=(0, 1),
|
|
required=True,
|
|
)
|
|
reconstruction.add_argument(
|
|
"--gs-views-interval",
|
|
type=int,
|
|
required=True,
|
|
)
|
|
reconstruction.add_argument("--video-frames", type=int, required=True)
|
|
reconstruction.add_argument(
|
|
"--device",
|
|
choices=("auto", "cuda", "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)
|
|
else:
|
|
reconstruct(args)
|
|
except Exception as error:
|
|
write_status(args.status, 1.0, "DA3 task failed", str(error))
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|