feat: adaptation of Canny processor

Adapted from controlnet processors package

fix: do final resize in canny processor

canny
This commit is contained in:
psychedelicious 2024-03-21 18:58:05 +11:00 committed by Kent Keirsey
parent ca496f0380
commit 611006b692
2 changed files with 45 additions and 5 deletions

View File

@ -7,7 +7,6 @@ from typing import Dict, List, Literal, Union
import cv2 import cv2
import numpy as np import numpy as np
from controlnet_aux import ( from controlnet_aux import (
CannyDetector,
ContentShuffleDetector, ContentShuffleDetector,
HEDdetector, HEDdetector,
LeresDetector, LeresDetector,
@ -39,6 +38,7 @@ from invokeai.app.invocations.model import ModelIdentifierField
from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.invocations.primitives import ImageOutput
from invokeai.app.invocations.util import validate_begin_end_step, validate_weights from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
from invokeai.app.services.shared.invocation_context import InvocationContext 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.depth_anything import DepthAnythingDetector
from invokeai.backend.image_util.dw_openpose import DWOpenposeDetector 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 # Keep alpha channel for Canny processing to detect edges of transparent areas
return context.images.get_pil(self.image.image_name, "RGBA") return context.images.get_pil(self.image.image_name, "RGBA")
def run_processor(self, image): def run_processor(self, image: Image.Image) -> Image.Image:
canny_processor = CannyDetector() processed_image = get_canny_edges(
processed_image = canny_processor(
image, image,
self.low_threshold, self.low_threshold,
self.high_threshold, self.high_threshold,
image_resolution=self.image_resolution,
detect_resolution=self.detect_resolution, detect_resolution=self.detect_resolution,
image_resolution=self.image_resolution,
) )
return processed_image return processed_image

View File

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