inference_onnx.py

import numpy as np
from PIL import Image
import onnxruntime as ort

def preprocess(image_path: str, f_px: float) -> tuple[np.ndarray, np.ndarray]:
    img = Image.open(image_path).convert("RGB")
    w, h = img.size
    x = np.asarray(img).astype(np.float32) / 255.0  # HWC [0,1]
    x = np.transpose(x, (2, 0, 1))[None, ...]       # NCHW
    # Resize to 1536x1536 with bilinear + align_corners=True:
    # For a minimal example, rely on ORT/consumer to match training preprocessing.
    # (For exact match, use the same resize code as Apple.)
    disparity_factor = np.array([f_px / float(w)], dtype=np.float32)
    return x.astype(np.float32), disparity_factor

if __name__ == "__main__":
    sess = ort.InferenceSession(
        "sharp_ndc_opset18.onnx",
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
    )
    print("Providers:", sess.get_providers())
    # Example usage requires you to supply f_px (or choose an approximate default).
    # image_resized_pt should be [1,3,1536,1536] — see README for exact contract.