inference script

inference_onnx.py

@@ -0,0 +1,302 @@
+#!/usr/bin/env python3
+"""ONNX Inference Script for SHARP Model.
+
+Loads an ONNX model (fp32 or fp16), runs inference on an input image,
+and exports the result as a PLY file.
+
+Usage:
+    python inference_onnx.py -m sharp.onnx -i test.png -o output.ply
+    python inference_onnx.py -m sharp_inline_fp16.onnx -i test.png -o output.ply -d 0.5
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from plyfile import PlyData, PlyElement
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+LOGGER = logging.getLogger(__name__)
+
+DEFAULT_HEIGHT = 1536
+DEFAULT_WIDTH = 1536
+
+
+def linear_to_srgb(linear: float) -> float:
+    if linear <= 0.0031308:
+        return linear * 12.92
+    return 1.055 * pow(linear, 1.0 / 2.4) - 0.055
+
+
+def rgb_to_sh(rgb: float) -> float:
+    coeff_degree0 = 1.0 / np.sqrt(4.0 * np.pi)
+    return (rgb - 0.5) / coeff_degree0
+
+
+def inverse_sigmoid(x: float) -> float:
+    x = np.clip(x, 1e-6, 1.0 - 1e-6)
+    return np.log(x / (1.0 - x))
+
+
+def preprocess_image(image_path: str | Path, target_size: tuple[int, int] = (DEFAULT_HEIGHT, DEFAULT_WIDTH)):
+    """Load and preprocess an image for ONNX inference."""
+    image_path = Path(image_path)
+    target_h, target_w = target_size
+
+    img = Image.open(image_path)
+    original_size = img.size
+    focal_length_px = original_size[0]
+
+    if img.size != (target_w, target_h):
+        img = img.resize((target_w, target_h), Image.BILINEAR)
+
+    img_np = np.array(img, dtype=np.float32) / 255.0
+
+    if img_np.shape[2] == 4:
+        img_np = img_np[:, :, :3]
+
+    img_np = np.transpose(img_np, (2, 0, 1))
+    img_np = np.expand_dims(img_np, axis=0)
+
+    LOGGER.info(f"Loaded image: {image_path}, original size: {original_size}")
+    LOGGER.info(f"Preprocessed shape: {img_np.shape}, range: [{img_np.min():.4f}, {img_np.max():.4f}]")
+
+    return img_np, float(focal_length_px), original_size
+
+
+def run_inference(onnx_path: str | Path, image: np.ndarray, disparity_factor: float = 1.0) -> dict[str, np.ndarray]:
+    """Run ONNX inference on the preprocessed image."""
+    onnx_path = Path(onnx_path)
+
+    LOGGER.info(f"Loading ONNX model: {onnx_path}")
+    
+    # Try with default providers first, then fallback to CPU only
+    try:
+        session = ort.InferenceSession(str(onnx_path))
+    except Exception as e:
+        error_msg = str(e)
+        if "tensor(float16)" in error_msg and "tensor(float)" in error_msg:
+            LOGGER.error("FP16 model has mixed float16/float32 types. This model was converted incorrectly.")
+            LOGGER.error("For FP16 inference on Apple Silicon, use the Core ML model (sharp.mlpackage) instead.")
+            LOGGER.error("Or regenerate the ONNX model with proper FP16 conversion.")
+            raise RuntimeError(f"Invalid FP16 model: {error_msg}")
+        # Try CPU fallback
+        try:
+            session = ort.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider'])
+        except Exception as cpu_e:
+            raise RuntimeError(f"Failed to load ONNX model: {cpu_e}")
+
+    input_names = [inp.name for inp in session.get_inputs()]
+    output_names = [out.name for out in session.get_outputs()]
+
+    LOGGER.info(f"Input names: {input_names}")
+    LOGGER.info(f"Output names: {output_names}")
+
+    inputs = {
+        "image": image.astype(np.float32),
+        "disparity_factor": np.array([disparity_factor], dtype=np.float32)
+    }
+
+    LOGGER.info("Running inference...")
+    raw_outputs = session.run(None, inputs)
+
+    outputs = {}
+
+    if len(raw_outputs) == 1:
+        concat = raw_outputs[0]
+        sizes = [3, 3, 4, 3, 1]
+        names = [
+            "mean_vectors_3d_positions",
+            "singular_values_scales",
+            "quaternions_rotations",
+            "colors_rgb_linear",
+            "opacities_alpha_channel"
+        ]
+        start = 0
+        for name, size in zip(names, sizes):
+            outputs[name] = concat[:, :, start:start + size]
+            start += size
+    elif len(raw_outputs) == 5:
+        names = [
+            "mean_vectors_3d_positions",
+            "singular_values_scales",
+            "quaternions_rotations",
+            "colors_rgb_linear",
+            "opacities_alpha_channel"
+        ]
+        for name, out in zip(names, raw_outputs):
+            outputs[name] = out
+    else:
+        for name, out in zip(output_names, raw_outputs):
+            outputs[name] = out
+
+    for name, arr in outputs.items():
+        LOGGER.info(f"  {name}: shape {arr.shape}")
+
+    return outputs
+
+
+def export_ply(outputs: dict[str, np.ndarray], output_path: str | Path,
+               focal_length_px: float, image_shape: tuple[int, int],
+               decimation: float = 1.0) -> None:
+    """Export Gaussians to PLY file format."""
+    output_path = Path(output_path)
+
+    mean_vectors = outputs["mean_vectors_3d_positions"]
+    singular_values = outputs["singular_values_scales"]
+    quaternions = outputs["quaternions_rotations"]
+    colors = outputs["colors_rgb_linear"]
+    opacities = outputs["opacities_alpha_channel"]
+
+    mean_vectors = mean_vectors[0]
+    singular_values = singular_values[0]
+    quaternions = quaternions[0]
+    colors = colors[0]
+    opacities = opacities[0]
+
+    num_gaussians = mean_vectors.shape[0]
+    LOGGER.info(f"Exporting {num_gaussians} Gaussians to PLY")
+
+    if decimation < 1.0:
+        log_scales = np.log(np.maximum(singular_values, 1e-10))
+        scale_product = np.exp(np.sum(log_scales, axis=1))
+        importance = scale_product * opacities
+
+        indices = np.argsort(-importance)
+        keep_count = max(1, int(num_gaussians * decimation))
+        keep_indices = indices[:keep_count]
+        keep_indices.sort()
+
+        LOGGER.info(f"Decimating: keeping {keep_count} of {num_gaussians} ({decimation * 100:.1f}%)")
+
+        mean_vectors = mean_vectors[keep_indices]
+        singular_values = singular_values[keep_indices]
+        quaternions = quaternions[keep_indices]
+        colors = colors[keep_indices]
+        opacities = opacities[keep_indices]
+        num_gaussians = keep_count
+
+    vertex_data = np.zeros(num_gaussians, dtype=[
+        ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+        ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4'),
+        ('opacity', 'f4'),
+        ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
+        ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4')
+    ])
+
+    vertex_data['x'] = mean_vectors[:, 0]
+    vertex_data['y'] = mean_vectors[:, 1]
+    vertex_data['z'] = mean_vectors[:, 2]
+
+    for i in range(num_gaussians):
+        r, g, b = colors[i]
+        srgb_r = linear_to_srgb(float(r))
+        srgb_g = linear_to_srgb(float(g))
+        srgb_b = linear_to_srgb(float(b))
+
+        vertex_data['f_dc_0'][i] = rgb_to_sh(srgb_r)
+        vertex_data['f_dc_1'][i] = rgb_to_sh(srgb_g)
+        vertex_data['f_dc_2'][i] = rgb_to_sh(srgb_b)
+
+    vertex_data['opacity'] = inverse_sigmoid(opacities)
+
+    vertex_data['scale_0'] = np.log(np.maximum(singular_values[:, 0], 1e-10))
+    vertex_data['scale_1'] = np.log(np.maximum(singular_values[:, 1], 1e-10))
+    vertex_data['scale_2'] = np.log(np.maximum(singular_values[:, 2], 1e-10))
+
+    vertex_data['rot_0'] = quaternions[:, 0]
+    vertex_data['rot_1'] = quaternions[:, 1]
+    vertex_data['rot_2'] = quaternions[:, 2]
+    vertex_data['rot_3'] = quaternions[:, 3]
+
+    vertex_element = PlyElement.describe(vertex_data, 'vertex')
+
+    # Extrinsic: 4x4 identity matrix as 16 separate properties
+    extrinsic_data = np.zeros(1, dtype=[('extrinsic', 'f4', (16,))])
+    extrinsic_data['extrinsic'][0] = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]
+    extrinsic_element = PlyElement.describe(extrinsic_data, 'extrinsic')
+
+    img_h, img_w = image_shape
+    # Intrinsic: 3x3 matrix as 9 separate properties
+    intrinsic_data = np.zeros(1, dtype=[('intrinsic', 'f4', (9,))])
+    intrinsic_data['intrinsic'][0] = [focal_length_px, 0, img_w / 2, 0, focal_length_px, img_h / 2, 0, 0, 1]
+    intrinsic_element = PlyElement.describe(intrinsic_data, 'intrinsic')
+
+    # Image size: 2 separate uint32 properties
+    image_size_data = np.zeros(1, dtype=[('image_size', 'u4', (2,))])
+    image_size_data['image_size'][0] = [img_w, img_h]
+    image_size_element = PlyElement.describe(image_size_data, 'image_size')
+
+    # Frame: 2 separate int32 properties
+    frame_data = np.zeros(1, dtype=[('frame', 'i4', (2,))])
+    frame_data['frame'][0] = [1, num_gaussians]
+    frame_element = PlyElement.describe(frame_data, 'frame')
+
+    z_values = mean_vectors[:, 2]
+    z_safe = np.maximum(z_values, 1e-6)
+    disparities = 1.0 / z_safe
+    disparities.sort()
+    disparity_10 = disparities[int(len(disparities) * 0.1)] if len(disparities) > 0 else 0.0
+    disparity_90 = disparities[int(len(disparities) * 0.9)] if len(disparities) > 0 else 1.0
+    disparity_data = np.zeros(1, dtype=[('disparity', 'f4', (2,))])
+    disparity_data['disparity'][0] = [disparity_10, disparity_90]
+    disparity_element = PlyElement.describe(disparity_data, 'disparity')
+
+    # Color space: single uchar property
+    color_space_data = np.zeros(1, dtype=[('color_space', 'u1')])
+    color_space_data['color_space'][0] = 1
+    color_space_element = PlyElement.describe(color_space_data, 'color_space')
+
+    # Version: 3 uchar properties
+    version_data = np.zeros(1, dtype=[('version', 'u1', (3,))])
+    version_data['version'][0] = [1, 5, 0]
+    version_element = PlyElement.describe(version_data, 'version')
+
+    PlyData([
+        vertex_element,
+        extrinsic_element,
+        intrinsic_element,
+        image_size_element,
+        frame_element,
+        disparity_element,
+        color_space_element,
+        version_element
+    ], text=False).write(str(output_path))
+
+    LOGGER.info(f"Saved PLY with {num_gaussians} Gaussians to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="ONNX Inference for SHARP - Generate 3D Gaussians from an image"
+    )
+    parser.add_argument("-m", "--model", type=str, required=True,
+                        help="Path to ONNX model file")
+    parser.add_argument("-i", "--input", type=str, required=True,
+                        help="Path to input image")
+    parser.add_argument("-o", "--output", type=str, required=True,
+                        help="Path to output file (.ply)")
+    parser.add_argument("-d", "--decimate", type=float, default=1.0,
+                        help="Decimation ratio 0.0-1.0 (default: 1.0 = keep all)")
+    parser.add_argument("--disparity-factor", type=float, default=1.0,
+                        help="Disparity factor for depth conversion (default: 1.0)")
+
+    args = parser.parse_args()
+
+    # Preprocess image
+    image, focal_length_px, image_shape = preprocess_image(args.input)
+
+    # Run inference
+    outputs = run_inference(args.model, image, args.disparity_factor)
+
+    # Export to PLY
+    export_ply(outputs, args.output, focal_length_px, image_shape, args.decimate)
+
+
+if __name__ == "__main__":
+    main()