Add ONNX model inference implementations for MoGe-2, BEN2, and DA3
- 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.
This commit is contained in:
306
src/snapform/server/media.py
Normal file
306
src/snapform/server/media.py
Normal file
@@ -0,0 +1,306 @@
|
||||
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
|
||||
Reference in New Issue
Block a user