| |
| 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 |
|
|
