Fix Detect Edges (Canny) on images with an alpha channel

kornia's canny only accepts 1 or 3 channel input, so a 4 channel image raised
a RuntimeError. Run edge detection on the colour channels.

CORE-393
This commit is contained in:
Glary-Bot
2026-08-15 00:20:02 +00:00
parent b963f4ad21
commit 79d16de082
2 changed files with 59 additions and 1 deletions

View File

@@ -0,0 +1,58 @@
import torch
from comfy.cli_args import args as cli_args
if not torch.cuda.is_available():
cli_args.cpu = True
from comfy_extras.nodes_canny import Canny # noqa: E402
def edged_image(channels, alpha=0.8, size=16):
"""Half black, half white, so there is a real edge down the middle."""
t = torch.zeros(1, size, size, channels)
t[:, :, size // 2:, :3] = 1.0
if channels == 4:
t[..., 3] = alpha
return t
def test_rgb_detects_the_edge():
out = Canny.execute(edged_image(3), 0.4, 0.8).result[0]
assert out.shape[-1] == 3
assert out.max() > 0.0
def test_rgba_does_not_raise():
out = Canny.execute(edged_image(4), 0.4, 0.8).result[0]
assert out.shape[-1] == 3
assert out.max() > 0.0
def test_rgba_and_rgb_give_the_same_edges():
"""Alpha must not influence edge detection."""
from_rgb = Canny.execute(edged_image(3), 0.4, 0.8).result[0]
from_rgba = Canny.execute(edged_image(4), 0.4, 0.8).result[0]
assert torch.equal(from_rgb, from_rgba)
def test_alpha_pattern_does_not_change_the_result():
opaque = edged_image(4, alpha=1.0)
transparent = edged_image(4, alpha=0.0)
from_opaque = Canny.execute(opaque, 0.4, 0.8).result[0]
from_transparent = Canny.execute(transparent, 0.4, 0.8).result[0]
assert torch.equal(from_opaque, from_transparent)
def test_does_not_mutate_input():
src = edged_image(4)
before = src.clone()
Canny.execute(src, 0.4, 0.8)
assert torch.equal(src, before)