#!/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()