import array
import numpy as np
import matplotlib.pyplot as plt
import OpenEXR, Imath, imageio

'''
    Utilitary code to load and display images
'''

def imshow(axe, img, gamma=False):
    if len(img.shape) == 2:
        img   = img / np.max(np.abs(img))
        img_r = img.reshape(img.shape[0], img.shape[1], 1)
        col_R = np.array([1,0,0]).reshape(1, 1, 3)
        col_B = np.array([0,0,1]).reshape(1, 1, 3)
        img = np.where(img_r>0, img_r*col_R, -img_r*col_B)
    if gamma:
        img = np.power(img, 1.0/2.1)
    h = axe.imshow(img.clip(0, 1))
    axe.set_xticks([])
    axe.set_yticks([])
    return h

def load_EXR_image(filename):
    pt = Imath.PixelType(Imath.PixelType.FLOAT)
    exr_image = OpenEXR.InputFile(filename)
    dw = exr_image.header()['dataWindow']
    size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1, 1)
    (r, g, b, a) = exr_image.channels('RGBA', pt)
    (r, g, b, a) = (np.array(array.array('f', r)), np.array(array.array('f', g)), np.array(array.array('f', b)), np.array(array.array('f', a)))
    (r, g, b, a) = (r.reshape(size), g.reshape(size), b.reshape(size), a.reshape(size))
    img = np.concatenate((r, g, b, a), axis=-1)
    return img

def load_image(k):
    direct = './data/{}/direct_{:04}.exr'.format(SCENE, k)
    indirect = './data/{}/lightmap_{:04}.exr'.format(SCENE, k)
    img_D = load_EXR_image(direct)
    img_I = load_EXR_image(indirect)
    return img_D, img_I

Load Database

Here, we load the viking database and reshape the elements to produce two matrices: the matrix of direct illumination $x$ and the matrix of indirect illumination $y$. Note that we scale every indirect illumination vector by the norm of the associated direct illumination vector to be agnostic of light intensity.

K = 128
NB_VECTORS = 4*16
W, H = 512, 512
DIRECT_W, DIRECT_H = 16, 16
TIKHONOV_REGUL = 0.001
SCENE = 'viking'

if True:
    img_D, img_I = load_image(0)
    DIRECT_W, DIRECT_H = img_D.shape[0], img_D.shape[1]
    W, H = img_I.shape[0], img_I.shape[1]

x = np.zeros((K, DIRECT_W, DIRECT_H, 3))
y = np.zeros((K, W, H, 3))

for k in range(K):
    img_D, img_I = load_image(k)
    x[k] = img_D[:,:,0:3]
    y[k] = img_I[:,:,0:3]

x  = x.reshape((x.shape[0], -1))
y  = np.transpose(y, (0, 3, 1, 2))
y  = y.reshape((y.shape[0], 3, -1))

# Scale by the inverse direct illumination norm
x_norm = np.linalg.norm(x.reshape(x.shape[0], -1), ord=2, axis=-1, keepdims=True)
x = x / x_norm
for k in range(3):
    y[:,k,:] = y[:,k,:] / x_norm

yb = y.reshape(y.shape[0]*3, -1)

Indirect Illumination Basis

def compute_pca(x, nb_elem=-1):
    if(nb_elem <= 0):
        nb_elem = x.shape[0]
    
    # Compute mean
    m  = np.mean(x, axis=0, keepdims=True)
    dx = x - m
    
    # Get covariance matrix
    C = (1.0/(K-1)) * np.matmul(dx, dx.transpose())
    v,W = np.linalg.eig(C)
    
    cv = np.cumsum(v)
    cv = cv / cv[-1]
    print("Nb of eigenvectors to reach 90% = ", np.argmin(np.square(cv - 0.90)))
    print("Nb of eigenvectors to reach 95% = ", np.argmin(np.square(cv - 0.95)))
    print("Nb of eigenvectors to reach 99% = ", np.argmin(np.square(cv - 0.99)))
    
    B = np.matmul(W[:,0:nb_elem].transpose(), dx)
    
    # Orthogonalize the vector space
    B,r = np.linalg.qr(B.transpose(), mode='reduced')
    B = B.transpose()

    em = np.sqrt(v[nb_elem-1])
    return m, B, em

m, B, em = compute_pca(yb, NB_VECTORS)
m = m.reshape((1, -1))
B_exp = np.concatenate((m, B), axis=0)

# Plot the basis elements
f, a = plt.subplots(1, 4, figsize=(10, 2.5))
for k in range(4):
    imshow(a[k], B_exp[k].reshape(W, H), gamma=True)
plt.tight_layout()

Nb of eigenvectors to reach 90% =  4
Nb of eigenvectors to reach 95% =  9
Nb of eigenvectors to reach 99% =  34

Compute the Transfer Matrix

We decompose the indirect illumination $y$ in the basis $B$ to obtain coefficients $c_I$: $$ c_I = (y - m) \times B $$ where $m = E[y]$ is the mean of the dataset.

We regress the transfer matrix $T$ using the matrix of direct illumination $x$ and the matrix of coefficients $c_I$. We regress a transfer matrix per color channel.

Note that we substract the mean to project the indirect illumination.

# Coefficients for indirect illumination in the basis
cI = np.matmul(y-m, B.transpose())

# Build linear regression
C    = np.matmul(x.transpose(), x)
C    = C
C    = C + TIKHONOV_REGUL*np.eye(C.shape[0]) # Tikhonov regularization
Cinv = np.linalg.pinv(C)
DtI  = np.einsum('ji, ikl -> jkl', x.transpose(), cI)

# Transfer Matrix
T    = np.einsum('ij, jkc -> kic', Cinv, DtI)

Reconstruction Routine

The reconstruction code takes a vector of direction illumination and multiply it with the transfer matrix to get the coefficients to reconstruct the indirect illumination: $$ c_I = {x \over || x ||} \times T $$

Use the basis elements and adding back the mean we get the indirect illumination: $$ y = ||x|| \left[ c_I \times B + m \right] $$

def reconstruct_indirect(p):
    cI = np.zeros((3, NB_VECTORS))
    cI[0] = np.matmul(p, T[0])
    cI[1] = np.matmul(p, T[1])
    cI[2] = np.matmul(p, T[2])

    z  = np.matmul(B.transpose(), cI.transpose()) + m.transpose()
    z  = z.clip(0.0, 1.0)

    return z.reshape(W, H, 3)

In this late cell, we compare our reconstruction method with projection of the indirect illumination in the basis and later recomposed with it: $$ \tilde{y} = ||x|| \left[ B^T \times \left( \left({y \over ||x||}-m\right) \times B \right) + m \right] $$ We show that our method reconstruct correctly the indirect illumination.

def reconstruct_indirect_PCA(y, printVectors=False):
    cI = np.matmul(y.reshape((3,-1))-m, B.transpose())
    z  = np.matmul(cI, B)+m
    z  = np.transpose(z, (1,0))
    z  = z.clip(0, 1.0)
    return z.reshape(W, H, 3)

NB_EXEMPLES = 5
f, a = plt.subplots(3, NB_EXEMPLES, figsize=(12, 7))
h = [[None for k in range(NB_EXEMPLES)] for i in range(3)]

id = np.arange(0, y.shape[0])
a[0][0].set_ylabel('reference')
a[1][0].set_ylabel('Transfer matrix')
a[2][0].set_ylabel('Basis projection')
for k in range(NB_EXEMPLES):
    h[0][k] = imshow(a[0][k], x_norm[k]*np.transpose(y[id[k]], (1,0)).reshape(W, H, 3))
    h[1][k] = imshow(a[1][k], x_norm[k]*reconstruct_indirect(x[id[k]]))
    h[2][k] = imshow(a[2][k], x_norm[k]*reconstruct_indirect_PCA(y[id[k]]))
    a[0][k].set_title('{:d}'.format(id[k]))