Compute surface normals from 3-plane "xyz" image data

normals_from_xyz.py

import numpy as np

def normals_from_xyz(xyz, smooth_sigma=None):
    """
    Given a 3-plane float image containing x,y,z coordinates,
    compute the surface normals at each pixel.
    Returns a 3-plane floating point image with nx,ny,nz per pixel.
    """
    if smooth_sigma is not None:
        xyz = np.array(xyz)  # make copy of input
        for p in range(3):
            xyz[:,:,p] = gaussian_filter(xyz[:,:,p], sigma=smooth_sigma)

    dx_du, dx_dv = np.gradient(xyz[:,:,0])
    dy_du, dy_dv = np.gradient(xyz[:,:,1])
    dz_du, dz_dv = np.gradient(xyz[:,:,2])

    grad_u = np.stack((dx_du, dy_du, dz_du), axis=2)
    grad_v = np.stack((dx_dv, dy_dv, dz_dv), axis=2)

    # compute surface normals as cross product of dx/du and dx/dv
    xyz_norm = np.cross(grad_u, grad_v, axis=2)
    xyz_norm /= np.expand_dims(np.linalg.norm(xyz_norm, axis=2), axis=2)
    return xyz_norm