Ray Tracing with Neural Networks

Introduction

Accurate simulation of imaging systems is crucial for enabling cost-effective and rapid design processes and training machine learning algorithms. Such simulation tools typically incorporate the full imaging pipeline, including the scene, camera optics, and image sensor. However, the proprietary nature of lens designs often leads manufacturers to withhold specific details, posing a challenge for accurate simulation. To address this challenge, Goossens et al. (2022) propose a solution that leverages Zemax black-box lens models. At the core of this method is the Ray Transfer Function (RTF), which characterizes how the position and angle of a ray entering a lens determine its exit position and angle, providing a practical approach to lens modeling without exposing proprietary design information. Goossens et al., modeled the RTF using a polynomial fitting and showed promising results as in [1]. The idea of this project is to explore the potential of neural networks to approximate the ray transfer function and enable efficient, high-fidelity ray tracing without the need for detailed lens specifications. As neural networks have demonstrated remarkable capabilities in learning complex nonlinear functions, making them a promising candidate for modeling the ray transfer function [2].

Background

1. Dataset Generation using Zemax

2. Polynomial Estimation of Ray-Transfer Functions

In a general form, the position of each input ray can be represented in six variables, $x,y,z,dx,dy,dz$. Similarly, the position of each output ray is represented as $x^{\prime },y^{\prime },z^{\prime },d{x}^{\prime },d{y}^{\prime },d{z}^{\prime }$. The polynomial-based approach aims to model this mapping between input and output ray positions using a set of polynomial equations:

$\hat{\beta }=\underset{{\beta }_0,{\beta }_1,{\beta }_2,\dots ,{\beta }_k}{\arg \min }\Vert [1,x,x^2,x^3,\dots ,x^k{]}^T[{\beta }_0,{\beta }_1,{\beta }_2,\dots ,{\beta }_k]-x^{\prime }{\Vert }_2^2$ where $x$ here is a column vector storing the "x" positions of all rays. The overall polynomial is multivariate but the logic is the same for $x,y,z,dx,dy,dz$. The user must specify the order of polynomial (denoted by $k$), as reported in [1], the fitting error varies as the polynomial degree changes (see the figure below).

Methods

1. Network Architecture

As a starting point, we are provided with the following GitHub repo: https://github.com/ISET/mlp-tiny, where a multi-layer perceptron (MLP) was implemented but the complete training/testing code was missing. We implemented the complete training/testing code and put all free parameters in a configuration file: https://github.com/ISET/mlp-tiny/blob/ray_tracing/mlp%20scripts/train_prediction.sh

The provided MLP has 16 fully connected layers (each is followed by batch normalization and ReLU activation function), as shown in the following figure.

However, one issue we found (see next section), was that the network struggled with predicting values in varying ranges, i.e., the coordinates can have significantly different maximums and minimums. To address the issue, we used positional encoding [3] that remapped the coordinates into [-1, 1] with a series of sinusoid functions at different frequencies, e.g., $x,sin(x),sin(2x),\dots ,sin(2^{p}x)$ (also see the following figure).

We integrated the positional embedding technique into the MLP and the resulting architecture is shown below:

where we found using separate branch to pre-process spatial and directional information and later merge them together to be the most effective. The modified architecture even has fewer trainable parameters than the MLP provided.

Results

1. Ray Tracing

2. Ray Classification

3. Remaining Issue: Imbalanced Dataset

References

[1] Thomas Goossens, Zheng Lyu, Jamyuen Ko, Gordon C. Wan, Joyce Farrell, and Brian Wandell, "Ray-transfer functions for camera simulation of 3D scenes with hidden lens design," Opt. Express 30, 24031-24047 (2022).

[2] Ian Goodfellow and Yoshua Bengio and Aaron Courville, "Deep Learning," MIT Press (2016), http://www.deeplearningbook.org.

[3] Position Embeddings. https://paperswithcode.com/methods/category/position-embeddings.