nan
c5b77c5c85
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models.
307 lines
9.1 KiB
Python
307 lines
9.1 KiB
Python
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"
|
||
|
||
|
||
class TaskCancelled(RuntimeError):
|
||
pass
|
||
|
||
|
||
def check_cancelled(cancel_check):
|
||
if cancel_check is not None and cancel_check():
|
||
raise TaskCancelled("Task cancelled")
|
||
|
||
|
||
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,
|
||
av,
|
||
start_frame=1,
|
||
cancel_check=None,
|
||
):
|
||
output_directory = Path(output_directory)
|
||
output_directory.mkdir(parents=True, exist_ok=True)
|
||
try:
|
||
container = av.open(str(video_path))
|
||
except Exception as error:
|
||
raise RuntimeError(f"Unable to open video: {video_path}") from error
|
||
image_paths = []
|
||
skipped = 0
|
||
try:
|
||
stream = container.streams.video[0]
|
||
for frame in container.decode(stream):
|
||
check_cancelled(cancel_check)
|
||
if skipped < start_frame - 1:
|
||
skipped += 1
|
||
continue
|
||
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
|
||
frame.to_image().save(image_path, format="JPEG")
|
||
image_paths.append(image_path)
|
||
if len(image_paths) >= maximum_frames:
|
||
break
|
||
finally:
|
||
container.close()
|
||
|
||
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 av
|
||
|
||
temporary_directory = Path(
|
||
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
|
||
)
|
||
try:
|
||
images = extract_video_frames(
|
||
input_path,
|
||
temporary_directory,
|
||
video_frames,
|
||
av,
|
||
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 extract_frames(context, parameters):
|
||
"""Copy selected source frames into the current job directory."""
|
||
context.directory.mkdir(parents=True, exist_ok=True)
|
||
images, temporary_directory = prepare_input(
|
||
parameters["input"],
|
||
parameters.get("video_frames", 16),
|
||
context.directory,
|
||
start_frame=parameters.get("start_frame", 1),
|
||
cancel_check=context.check_cancelled,
|
||
)
|
||
try:
|
||
copy_frames(images, context.directory)
|
||
finally:
|
||
if temporary_directory is not None:
|
||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||
message = f"Extracted {len(images)} frame(s)"
|
||
context.progress(1.0, message)
|
||
return {"frames": len(images)}
|
||
|
||
|
||
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
|