Function integrate_vector_field

scripts/normal_to_height_map.py:113–164

(
    vector_field: np.ndarray,
    mask: np.ndarray,
    target_iteration_count: int,
    thread_count: int,
)

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(