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