diff --git a/invokeai/app/invocations/controlnet_image_processors.py b/invokeai/app/invocations/controlnet_image_processors.py index ed56344847..547b335a70 100644 --- a/invokeai/app/invocations/controlnet_image_processors.py +++ b/invokeai/app/invocations/controlnet_image_processors.py @@ -7,7 +7,6 @@ from typing import Dict, List, Literal, Union import cv2 import numpy as np from controlnet_aux import ( - CannyDetector, ContentShuffleDetector, HEDdetector, LeresDetector, @@ -39,6 +38,7 @@ from invokeai.app.invocations.model import ModelIdentifierField from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.invocations.util import validate_begin_end_step, validate_weights from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.image_util.canny import get_canny_edges from invokeai.backend.image_util.depth_anything import DepthAnythingDetector from invokeai.backend.image_util.dw_openpose import DWOpenposeDetector @@ -189,14 +189,13 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation): # Keep alpha channel for Canny processing to detect edges of transparent areas return context.images.get_pil(self.image.image_name, "RGBA") - def run_processor(self, image): - canny_processor = CannyDetector() - processed_image = canny_processor( + def run_processor(self, image: Image.Image) -> Image.Image: + processed_image = get_canny_edges( image, self.low_threshold, self.high_threshold, - image_resolution=self.image_resolution, detect_resolution=self.detect_resolution, + image_resolution=self.image_resolution, ) return processed_image diff --git a/invokeai/backend/image_util/canny.py b/invokeai/backend/image_util/canny.py new file mode 100644 index 0000000000..e0067bd7fd --- /dev/null +++ b/invokeai/backend/image_util/canny.py @@ -0,0 +1,41 @@ +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)