| |
| """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() |
|
|
