Rewrite of https://github.com/valgur/surface-normal using taichi

surface_normal_taichi.py

import numpy as np
import taichi as ti
from skimage.io import imread, imsave

ti.init(arch=ti.cuda)


@ti.kernel
def normals_from_depth_kernel(f: ti.f32, cx: ti.f32, cy: ti.f32, window_size: ti.u16, max_rel_depth_diff: ti.f32):
    height, width = ti_depth.shape()
    f_inv = 1. / f
    r = window_size // 2
    for row, col in ti_depth:
        center_depth = ti_depth[row, col]
        if center_depth == 0:
            continue

        mid = ti.Vector([(col - cx) * center_depth * f_inv, (row - cy) * center_depth * f_inv, center_depth])
        n = 0
        centroid = ti.Vector([0, 0, 0])
        outer_prod_sum = ti.Matrix([[0, 0, 0]] * 3, dt=ti.f32)
        for i in range(-r, r + 1):
            for j in range(-r, r + 1):
                x = col + j
                y = row + i
                if not (0 <= x < width and 0 <= y < height):
                    continue

                z = ti_depth[y, x]
                if z == 0 or ti.abs(z - center_depth) > max_rel_depth_diff * center_depth:
                    continue

                p = ti.Vector([(x - cx) * z * f_inv, (y - cy) * z * f_inv, z])
                p -= mid  # subtract midpoint for improved numeric stability in outer product
                centroid += p
                outer_prod_sum += ti.outer_product(p, p)
                n += 1

        if n >= 3:
            # centroid /= n
            cov = (outer_prod_sum - ti.outer_product(centroid, centroid / n)) / (n - 1)
            U, S, V = ti.svd(cov, ti.f32)
            normal = ti.Vector([V[0, 2], V[1, 2], V[2, 2]])
            normal /= normal.norm()

            if mid.dot(normal) < 0:
                normal *= -1

            ti_normals[row, col][0] = 127.5 * (1 - normal[0])
            ti_normals[row, col][1] = 127.5 * (1 - normal[2])
            ti_normals[row, col][2] = 127.5 * (1 - normal[1])


def normals_from_depth(depth, intrinsics, window_size=15, max_rel_depth_diff=0.1):
    global ti_depth, ti_normals
    if not ti.get_runtime().materialized:
        ti_depth = ti.var(dt=ti.f32, shape=depth.shape)
        ti_normals = ti.Vector(3, dt=ti.u8, shape=depth.shape)
    else:
        assert depth.shape == ti_depth.shape()
    ti_depth.from_numpy(depth)
    normals_from_depth_kernel(*intrinsics, window_size, max_rel_depth_diff)
    return np.squeeze(ti_normals.to_numpy())


f = 721.5377
cx = 609.5593
cy = 172.8540
intrinsics = (f, cx, cy)
depth = imread("depth.png")
normals = normals_from_depth(depth, intrinsics)
imsave("normals_taichi.png", normals)