Implement Omoo DA3 functionality with property group, runtime, worker, and status bar
- Added properties for Omoo DA3 settings in properties.py, including input path, video frames, model selection, and export format. - Created runtime management in runtime.py, handling model downloads, job paths, and command execution for reconstruction. - Developed status bar functionality in status_bar.py to display job progress and status. - Implemented worker.py for image and video processing, including frame extraction and model inference. - Added comprehensive unit tests covering addon registration, model handling, export processes, and UI interactions. - Included a lock file for dependency management with uv.
This commit is contained in:
304
src/omoo_da3/worker.py
Normal file
304
src/omoo_da3/worker.py
Normal file
@@ -0,0 +1,304 @@
|
||||
import argparse
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
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="omoo-da3-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
|
||||
|
||||
|
||||
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}")
|
||||
snapshot_download(
|
||||
repo_id=args.repository,
|
||||
local_dir=model_dir,
|
||||
)
|
||||
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):
|
||||
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.03, "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.10, 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.18, f"Processing {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()
|
||||
Reference in New Issue
Block a user