Unique3D / scripts /normal_to_height_map.py
Wuvin's picture
init
37aeb5b
raw
history blame
6.39 kB
# code modified from https://github.com/YertleTurtleGit/depth-from-normals
import numpy as np
import cv2 as cv
from multiprocessing.pool import ThreadPool as Pool
from multiprocessing import cpu_count
from typing import Tuple, List, Union
import numba
def calculate_gradients(
normals: np.ndarray, mask: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
horizontal_angle_map = np.arccos(np.clip(normals[:, :, 0], -1, 1))
left_gradients = np.zeros(normals.shape[:2])
left_gradients[mask != 0] = (1 - np.sin(horizontal_angle_map[mask != 0])) * np.sign(
horizontal_angle_map[mask != 0] - np.pi / 2
)
vertical_angle_map = np.arccos(np.clip(normals[:, :, 1], -1, 1))
top_gradients = np.zeros(normals.shape[:2])
top_gradients[mask != 0] = -(1 - np.sin(vertical_angle_map[mask != 0])) * np.sign(
vertical_angle_map[mask != 0] - np.pi / 2
)
return left_gradients, top_gradients
@numba.jit(nopython=True)
def integrate_gradient_field(
gradient_field: np.ndarray, axis: int, mask: np.ndarray
) -> np.ndarray:
heights = np.zeros(gradient_field.shape)
for d1 in numba.prange(heights.shape[1 - axis]):
sum_value = 0
for d2 in range(heights.shape[axis]):
coordinates = (d1, d2) if axis == 1 else (d2, d1)
if mask[coordinates] != 0:
sum_value = sum_value + gradient_field[coordinates]
heights[coordinates] = sum_value
else:
sum_value = 0
return heights
def calculate_heights(
left_gradients: np.ndarray, top_gradients, mask: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
left_heights = integrate_gradient_field(left_gradients, 1, mask)
right_heights = np.fliplr(
integrate_gradient_field(np.fliplr(-left_gradients), 1, np.fliplr(mask))
)
top_heights = integrate_gradient_field(top_gradients, 0, mask)
bottom_heights = np.flipud(
integrate_gradient_field(np.flipud(-top_gradients), 0, np.flipud(mask))
)
return left_heights, right_heights, top_heights, bottom_heights
def combine_heights(*heights: np.ndarray) -> np.ndarray:
return np.mean(np.stack(heights, axis=0), axis=0)
def rotate(matrix: np.ndarray, angle: float) -> np.ndarray:
h, w = matrix.shape[:2]
center = (w / 2, h / 2)
rotation_matrix = cv.getRotationMatrix2D(center, angle, 1.0)
corners = cv.transform(
np.array([[[0, 0], [w, 0], [w, h], [0, h]]]), rotation_matrix
)[0]
_, _, w, h = cv.boundingRect(corners)
rotation_matrix[0, 2] += w / 2 - center[0]
rotation_matrix[1, 2] += h / 2 - center[1]
result = cv.warpAffine(matrix, rotation_matrix, (w, h), flags=cv.INTER_LINEAR)
return result
def rotate_vector_field_normals(normals: np.ndarray, angle: float) -> np.ndarray:
angle = np.radians(angle)
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
rotated_normals = np.empty_like(normals)
rotated_normals[:, :, 0] = (
normals[:, :, 0] * cos_angle - normals[:, :, 1] * sin_angle
)
rotated_normals[:, :, 1] = (
normals[:, :, 0] * sin_angle + normals[:, :, 1] * cos_angle
)
return rotated_normals
def centered_crop(image: np.ndarray, target_resolution: Tuple[int, int]) -> np.ndarray:
return image[
(image.shape[0] - target_resolution[0])
// 2 : (image.shape[0] - target_resolution[0])
// 2
+ target_resolution[0],
(image.shape[1] - target_resolution[1])
// 2 : (image.shape[1] - target_resolution[1])
// 2
+ target_resolution[1],
]
def integrate_vector_field(
vector_field: np.ndarray,
mask: np.ndarray,
target_iteration_count: int,
thread_count: int,
) -> np.ndarray:
shape = vector_field.shape[:2]
angles = np.linspace(0, 90, target_iteration_count, endpoint=False)
def integrate_vector_field_angles(angles: List[float]) -> np.ndarray:
all_combined_heights = np.zeros(shape)
for angle in angles:
rotated_vector_field = rotate_vector_field_normals(
rotate(vector_field, angle), angle
)
rotated_mask = rotate(mask, angle)
left_gradients, top_gradients = calculate_gradients(
rotated_vector_field, rotated_mask
)
(
left_heights,
right_heights,
top_heights,
bottom_heights,
) = calculate_heights(left_gradients, top_gradients, rotated_mask)
combined_heights = combine_heights(
left_heights, right_heights, top_heights, bottom_heights
)
combined_heights = centered_crop(rotate(combined_heights, -angle), shape)
all_combined_heights += combined_heights / len(angles)
return all_combined_heights
with Pool(processes=thread_count) as pool:
heights = pool.map(
integrate_vector_field_angles,
np.array(
np.array_split(angles, thread_count),
dtype=object,
),
)
pool.close()
pool.join()
isotropic_height = np.zeros(shape)
for height in heights:
isotropic_height += height / thread_count
return isotropic_height
def estimate_height_map(
normal_map: np.ndarray,
mask: Union[np.ndarray, None] = None,
height_divisor: float = 1,
target_iteration_count: int = 250,
thread_count: int = cpu_count(),
raw_values: bool = False,
) -> np.ndarray:
if mask is None:
if normal_map.shape[-1] == 4:
mask = normal_map[:, :, 3] / 255
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
else:
mask = np.ones(normal_map.shape[:2], dtype=np.uint8)
normals = ((normal_map[:, :, :3].astype(np.float64) / 255) - 0.5) * 2
heights = integrate_vector_field(
normals, mask, target_iteration_count, thread_count
)
if raw_values:
return heights
heights /= height_divisor
heights[mask > 0] += 1 / 2
heights[mask == 0] = 1 / 2
heights *= 2**16 - 1
if np.min(heights) < 0 or np.max(heights) > 2**16 - 1:
raise OverflowError("Height values are clipping.")
heights = np.clip(heights, 0, 2**16 - 1)
heights = heights.astype(np.uint16)
return heights