mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
|
import cv2
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from invokeai.backend.image_util.util import (
|
||
|
|
cv2_to_pil,
|
||
|
|
fit_image_to_resolution,
|
||
|
|
normalize_image_channel_count,
|
||
|
|
pil_to_cv2,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_canny_edges(
|
||
|
|
image: Image.Image, low_threshold: int, high_threshold: int, detect_resolution: int, image_resolution: int
|
||
|
|
) -> Image.Image:
|
||
|
|
"""Returns the edges of an image using the Canny edge detection algorithm.
|
||
|
|
|
||
|
|
This function is adapted from https://github.com/lllyasviel/ControlNet.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image: The input image.
|
||
|
|
low_threshold: The lower threshold for the hysteresis procedure.
|
||
|
|
high_threshold: The upper threshold for the hysteresis procedure.
|
||
|
|
input_resolution: The resolution of the input image. The image will be resized to this resolution before edge detection.
|
||
|
|
output_resolution: The resolution of the output image. The edges will be resized to this resolution before returning.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The Canny edges of the input image.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if image.mode != "RGB":
|
||
|
|
image = image.convert("RGB")
|
||
|
|
|
||
|
|
np_image = pil_to_cv2(image)
|
||
|
|
np_image = normalize_image_channel_count(np_image)
|
||
|
|
np_image = fit_image_to_resolution(np_image, detect_resolution)
|
||
|
|
|
||
|
|
edge_map = cv2.Canny(np_image, low_threshold, high_threshold)
|
||
|
|
edge_map = normalize_image_channel_count(edge_map)
|
||
|
|
edge_map = fit_image_to_resolution(edge_map, image_resolution)
|
||
|
|
|
||
|
|
return cv2_to_pil(edge_map)
|