Convert blender world normals EXR to camera normal

world_to_camera_normal.py

import numpy as np
import OpenEXR, Imath
from scipy.misc import imsave

INPUT_EXR_NORMALS_FILE = '1-world.exr'
OUTPUT_RGB_NORMAL_FILE = '1-world-to-camera.png'

INVERTED_CAMERA_QUARTERNION = np.array([0.8127532005310059, -0.456145316362381, 0.17738527059555054, 0.31606677174568176], dtype=np.float32)
X_CHANNEL_NAME = 'RenderLayer.Normal.X'
Y_CHANNEL_NAME = 'RenderLayer.Normal.Y'
Z_CHANNEL_NAME = 'RenderLayer.Normal.Z'

##
# q: quaternion
# v: 3-element array
# @see adapted from blender's math_rotation.c
#
# \note:
# Assumes a unit quaternion?
#
# in fact not, but you may want to use a unit quat, read on...
#
# Shortcut for 'q v q*' when \a v is actually a quaternion.
# This removes the need for converting a vector to a quaternion,
# calculating q's conjugate and converting back to a vector.
# It also happens to be faster (17+,24* vs * 24+,32*).
# If \a q is not a unit quaternion, then \a v will be both rotated by
# the same amount as if q was a unit quaternion, and scaled by the square of
# the length of q.
#
# For people used to python mathutils, its like:
# def mul_qt_v3(q, v): (q * Quaternion((0.0, v[0], v[1], v[2])) * q.conjugated())[1:]
#
# \note: multiplying by 3x3 matrix is ~25% faster.
##
def _multiply_quaternion_vec3(q, v):
    t0 = -q[1] * v[0] - q[2] * v[1] - q[3] * v[2]
    t1 = q[0] * v[0] + q[2] * v[2] - q[3] * v[1]
    t2 = q[0] * v[1] + q[3] * v[0] - q[1] * v[2]
    i = [t1, t2, q[0] * v[2] + q[1] * v[1] - q[2] * v[0]]
    t1 = t0 * -q[1] + i[0] * q[0] - i[1] * q[3] + i[2] * q[2]
    t2 = t0 * -q[2] + i[1] * q[0] - i[2] * q[1] + i[0] * q[3]
    i[2] = t0 * -q[3] + i[2] * q[0] - i[0] * q[2] + i[1] * q[1]
    i[0] = t1
    i[1] = t2
    return i


def world_to_camera_normals(inverted_camera_quaternation, exr_x, exr_y, exr_z):
    camera_normal = np.empty([exr_x.shape[0], exr_x.shape[1], 3], dtype=np.float32)
    for i in range(exr_x.shape[0]):
        for j in range(exr_x.shape[1]):
            pixel_camera_normal = _multiply_quaternion_vec3(
                inverted_camera_quaternation,
                [exr_x[i][j], exr_y[i][j], exr_z[i][j]]
            )
            camera_normal[i][j][0] = pixel_camera_normal[0]
            camera_normal[i][j][1] = pixel_camera_normal[1]
            camera_normal[i][j][2] = pixel_camera_normal[2]
    return camera_normal

# checks exr_x, exr_y, exr_z for original zero values
def normal_to_rgb(normals_to_convert, exr_x, exr_y, exr_z):
    camera_normal_rgb = normals_to_convert + 1
    camera_normal_rgb *= 127.5
    camera_normal_rgb = camera_normal_rgb.astype(np.uint8)
    # Correct case when we had an empty normal (0,0,0), and turn it to black instead of gray
    for i in range(exr_x.shape[0]):
        for j in range(exr_x.shape[1]):
            if exr_x[i][j] == 0 and exr_y[i][j] == 0 and exr_z[i][j] == 0:
                camera_normal_rgb[i][j][0] = 0
                camera_normal_rgb[i][j][1] = 0
                camera_normal_rgb[i][j][2] = 0
    return camera_normal_rgb

# Return X, Y, Z normals as numpy arrays
def read_exr_normal_file(exr_path):
    exr_file = OpenEXR.InputFile(exr_path)
    cm_dw = exr_file.header()['dataWindow']
    exr_x = np.fromstring(
        exr_file.channel(X_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_x.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols
    exr_y = np.fromstring(
        exr_file.channel(Y_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_y.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols

    exr_z = np.fromstring(
        exr_file.channel(Z_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_z.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols
    return exr_x, exr_y, exr_z




exr_x, exr_y, exr_z = read_exr_normal_file(INPUT_EXR_NORMALS_FILE)

assert(exr_x.shape == exr_y.shape)
assert(exr_y.shape == exr_z.shape)
print("exr_z.shape =", exr_z.shape)

camera_normal = world_to_camera_normals(INVERTED_CAMERA_QUARTERNION, exr_x, exr_y, exr_z)
camera_normal_rgb = normal_to_rgb(camera_normal, exr_x, exr_y, exr_z)

imsave(OUTPUT_RGB_NORMAL_FILE, camera_normal_rgb)
print("RGB normals saved to", OUTPUT_RGB_NORMAL_FILE)