48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
"""
|
|
Script for generating sphere ground truth normal images.
|
|
"""
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
|
|
def generate_sphere_gt_normals(img_mask, center_x, center_y, radius):
|
|
"""
|
|
Generates sphere ground truth normal images for an image.
|
|
Args:
|
|
img_mask: a numpy array of shape [H, W, 3]
|
|
center_x: x coordinate of the center of the sphere
|
|
center_y: y coordinate of the center of the sphere
|
|
radius: the radius of the sphere
|
|
Returns:
|
|
img_normal: a numpy array of shape [H, W, 3]
|
|
"""
|
|
img_normal = np.zeros(img_mask.shape, dtype="float64")
|
|
|
|
for y in range(img_mask.shape[0]):
|
|
for x in range(img_mask.shape[1]):
|
|
|
|
img_normal[y, x, 0] = 0.0
|
|
img_normal[y, x, 1] = 0.0
|
|
img_normal[y, x, 2] = 1.0
|
|
|
|
if np.sum(img_mask[y, x, :]) > 0:
|
|
dist = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2)
|
|
ang_xz = math.acos(dist / radius)
|
|
ang_xy = math.atan2(y - center_y, x - center_x)
|
|
|
|
nx = math.cos(ang_xz) * math.cos(ang_xy)
|
|
ny = math.cos(ang_xz) * math.sin(ang_xy)
|
|
nz = math.sin(ang_xz)
|
|
|
|
img_normal[y, x, 0] = nx
|
|
img_normal[y, x, 1] = -ny
|
|
img_normal[y, x, 2] = nz
|
|
|
|
norm_val = np.linalg.norm(img_normal[y, x, :])
|
|
img_normal[y, x, :] = img_normal[y, x, :] / norm_val
|
|
|
|
# img_normal between [-1., 1.], converting to [0., 1.]
|
|
img_normal = (img_normal + 1.0) * 0.5
|
|
return img_normal
|