File size: 2,317 Bytes
d922f06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/usr/bin/env python3
"""Minimal ONNX Runtime example for EfficientViT-Depth-5M (float).
The model predicts affine-invariant inverse depth: higher values are closer.
For metric depth, fit a per-image scale and shift on a known reference.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
from PIL import Image
import onnxruntime as ort
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
def preprocess(image_path: Path, height: int = 336, width: int = 448) -> np.ndarray:
image = Image.open(image_path).convert("RGB").resize((width, height), Image.Resampling.BICUBIC)
array = np.asarray(image, dtype=np.float32) / 255.0
array = (array - MEAN) / STD
return np.transpose(array, (2, 0, 1))[None].astype(np.float32)
def save_visualization(depth: np.ndarray, output_path: Path) -> None:
depth = np.squeeze(depth).astype(np.float32)
lo, hi = float(np.nanmin(depth)), float(np.nanmax(depth))
normalized = (depth - lo) / max(hi - lo, 1e-6)
Image.fromarray((normalized * 255.0).astype(np.uint8), mode="L").save(output_path)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--image", type=Path, required=True, help="input RGB image")
parser.add_argument("--output", type=Path, default=Path("depth.png"), help="grayscale depth preview")
parser.add_argument(
"--model",
type=Path,
default=Path(__file__).with_name("efficientvit_depth_5m_fp32.onnx"),
help="ONNX model path",
)
parser.add_argument("--npy", type=Path, default=None, help="optional raw inverse-depth .npy output")
args = parser.parse_args()
session = ort.InferenceSession(str(args.model), providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
depth = session.run(None, {input_name: preprocess(args.image)})[0]
depth = np.squeeze(depth)
print(f"output shape: {depth.shape}")
print(f"inverse depth: min={depth.min():.4f} max={depth.max():.4f} mean={depth.mean():.4f}")
save_visualization(depth, args.output)
if args.npy is not None:
np.save(args.npy, depth)
print(f"saved preview: {args.output}")
if __name__ == "__main__":
main()
|