def visualize_disparity_with_blinn_phong()

in run_nerf_helpers.py [0:0]

• 58 lines of code
• 5 McCabe index (conditional complexity)

def visualize_disparity_with_blinn_phong(depth_map):
    # follows https://en.wikipedia.org/wiki/Blinn%E2%80%93Phong_reflection_model
    lightPos = np.array([1.0, 1.0, 1.0])
    lightColor = np.array([1.0, 1.0, 1.0])
    lightPower = 2.0  # 40.0
    ambientColor = np.array([0.1, 0.0, 0.0])
    diffuseColor = np.array([0.5, 0.0, 0.0])
    specColor = np.array([1.0, 1.0, 1.0])
    shininess = 2.0  # 16.0

    height, width = depth_map.shape

    # normals from depth map
    # https://stackoverflow.com/questions/53350391/surface-normal-calculation-from-depth-map-in-python
    spacing = 2.0 / (height - 1)
    zy, zx = np.gradient(depth_map, spacing)
    normal = np.dstack(
        (-zx, zy, np.ones_like(depth_map))
    )  # need to flip zy because OpenGL indexes bottom left as (0,0) (this is a guess, it simply turns out to work if zy is flipped)
    normal_length = np.linalg.norm(normal, axis=2, keepdims=True)
    normal /= normal_length  # height x width x 3

    i, j = np.meshgrid(
        np.arange(width, dtype=np.float32) / width,
        np.arange(height, dtype=np.float32) / width,
        indexing="xy",
    )  # note: if height != width then dividing the second argument by height would lead to anisotropic scaling
    vertPos = np.stack(
        [i, j, depth_map], axis=-1
    )  # height x width x 3. note that (x,y) and (depth) have different scaling factors and offsets because we don't do proper unprojection - might need to adjust them

    lightDir = -vertPos + lightPos.reshape(1, 1, 3)  # height x width x 3
    distance = np.linalg.norm(lightDir, axis=2, keepdims=True)  # height x width x 1
    lightDir /= distance
    # distance = distance ** 2
    distance = (distance + 1.0) ** 2

    def dot_product(A, B):
        return np.sum(A * B, axis=-1)

    lightDir_x_normal = dot_product(lightDir, normal)
    lambertian = np.clip(lightDir_x_normal, a_max=None, a_min=0.0).reshape(
        height, width, 1
    )  # height x width x 1

    invalid_mask = lambertian <= 0.0

    def normalize(image):
        return image / np.linalg.norm(image, axis=-1, keepdims=True)

    viewDir = normalize(-vertPos)  # height x width x 3
    halfDir = normalize(lightDir + viewDir)  # height x width x 3
    # specAngle = np.clip(dot_product(halfDir, normal), a_max=None, a_min=0.).reshape(height, width, 1) # height x width x 1
    specAngle = np.clip(dot_product(halfDir, -normal), a_max=None, a_min=0.0).reshape(
        height, width, 1
    )  # height x width x 1
    specular = specAngle ** shininess

    specular[invalid_mask] = 0.0

    colorLinear = (
        lambertian
        * diffuseColor.reshape(1, 1, 3)
        * lightColor.reshape(1, 1, 3)
        * lightPower
        / distance
        + specular
        * specColor.reshape(1, 1, 3)
        * lightColor.reshape(1, 1, 3)
        * lightPower
        / distance
        + ambientColor.reshape(1, 1, 3)
    )  # height x width x 3
    return colorLinear