import itertools
import json
import math
import os
import comfy.model_management as model_management
import folder_paths
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageOps
from PIL.PngImagePlugin import PngInfo
from skimage.filters import gaussian
from skimage.util import compare_images
from ..log import log
from ..utils import np2tensor, pil2tensor, tensor2pil
# try:
# from cv2.ximgproc import guidedFilter
# except ImportError:
# log.warning("cv2.ximgproc.guidedFilter not found, use opencv-contrib-python")
def gaussian_kernel(
kernel_size: int, sigma_x: float, sigma_y: float, device=None
):
x, y = torch.meshgrid(
torch.linspace(-1, 1, kernel_size, device=device),
torch.linspace(-1, 1, kernel_size, device=device),
indexing="ij",
)
d_x = x * x / (2.0 * sigma_x * sigma_x)
d_y = y * y / (2.0 * sigma_y * sigma_y)
g = torch.exp(-(d_x + d_y))
return g / g.sum()
class MTB_CoordinatesToString:
RETURN_TYPES = ("STRING",)
FUNCTION = "convert"
CATEGORY = "mtb/coordinates"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"coordinates": ("BATCH_COORDINATES",),
"frame": ("INT",),
}
}
def convert(
self, coordinates: list[list[tuple[int, int]]], frame: int
) -> tuple[str]:
frame = max(frame, len(coordinates) - 1)
coords = coordinates[frame]
output: list[dict[str, int]] = []
for x, y in coords:
output.append({"x": x, "y": y})
return (json.dumps(output),)
class MTB_ExtractCoordinatesFromImage:
"""Extract 2D points from a batch of images based on a threshold."""
RETURN_TYPES = ("BATCH_COORDINATES", "IMAGE")
FUNCTION = "extract"
CATEGORY = "mtb/coordinates"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"threshold": ("FLOAT",),
"max_points": ("INT", {"default": 50, "min": 0}),
},
"optional": {"image": ("IMAGE",), "mask": ("MASK",)},
}
def extract(
self,
threshold: float,
max_points: int,
image: torch.Tensor | None = None,
mask: torch.Tensor | None = None,
) -> tuple[list[list[tuple[int, int]]], torch.Tensor]:
if image is not None:
batch_count, height, width, channel_count = image.shape
imgs = image
else:
if mask is None:
raise ValueError("Must provide either image or mask")
batch_count, height, width = mask.shape
channel_count = 1
imgs = mask
if channel_count not in [1, 2, 3, 4]:
raise ValueError(f"Incorrect channel count: {channel_count}")
all_points: list[list[tuple[int, int]]] = []
debug_images = torch.zeros(
(batch_count, height, width, 3),
dtype=torch.uint8,
device=imgs.device,
)
for i, img in enumerate(imgs):
if channel_count == 1:
alpha_channel = img if len(img.shape) == 2 else img[:, :, 0]
elif channel_count == 2:
alpha_channel = img[:, :, 1]
elif channel_count == 4:
alpha_channel = img[:, :, 3]
else:
# get intensity
alpha_channel = img[:, :, :3].max(dim=2)[0]
points = (alpha_channel > threshold).nonzero(as_tuple=False)
if len(points) > max_points:
indices = torch.randperm(points.size(0), device=img.device)[
:max_points
]
points = points[indices]
points = [(int(y.item()), int(x.item())) for x, y in points]
all_points.append(points)
for x, y in points:
self._draw_circle(debug_images[i], (x, y), 5)
return (all_points, debug_images)
@staticmethod
def _draw_circle(
image: torch.Tensor, center: tuple[int, int], radius: int
):
"""Draw a 5px circle on the image."""
x0, y0 = center
for x in range(-radius, radius + 1):
for y in range(-radius, radius + 1):
in_radius = x**2 + y**2 <= radius**2
in_bounds = (
0 <= x0 + x < image.shape[1]
and 0 <= y0 + y < image.shape[0]
)
if in_radius and in_bounds:
image[y0 + y, x0 + x] = torch.tensor(
[255, 255, 255],
dtype=torch.uint8,
device=image.device,
)
class MTB_ColorCorrectGPU:
"""Various color correction methods using only Torch."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"force_gpu": ("BOOLEAN", {"default": True}),
"clamp": ([True, False], {"default": True}),
"gamma": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"contrast": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"exposure": (
"FLOAT",
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
),
"offset": (
"FLOAT",
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
),
"hue": (
"FLOAT",
{"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.01},
),
"saturation": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"value": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
},
"optional": {"mask": ("MASK",)},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "correct"
CATEGORY = "mtb/image processing"
@staticmethod
def get_device(tensor: torch.Tensor, force_gpu: bool):
if force_gpu:
if torch.cuda.is_available():
return torch.device("cuda")
elif (
hasattr(torch.backends, "mps")
and torch.backends.mps.is_available()
):
return torch.device("mps")
elif hasattr(torch, "hip") and torch.hip.is_available():
return torch.device("hip")
return (
tensor.device
) # model_management.get_torch_device() # torch.device("cpu")
@staticmethod
def rgb_to_hsv(image: torch.Tensor):
r, g, b = image.unbind(-1)
max_rgb, argmax_rgb = image.max(-1)
min_rgb, _ = image.min(-1)
diff = max_rgb - min_rgb
h = torch.empty_like(max_rgb)
s = diff / (max_rgb + 1e-7)
v = max_rgb
h[argmax_rgb == 0] = (g - b)[argmax_rgb == 0] / (diff + 1e-7)[
argmax_rgb == 0
]
h[argmax_rgb == 1] = (
2.0 + (b - r)[argmax_rgb == 1] / (diff + 1e-7)[argmax_rgb == 1]
)
h[argmax_rgb == 2] = (
4.0 + (r - g)[argmax_rgb == 2] / (diff + 1e-7)[argmax_rgb == 2]
)
h = (h / 6.0) % 1.0
h = h.unsqueeze(-1)
s = s.unsqueeze(-1)
v = v.unsqueeze(-1)
return torch.cat((h, s, v), dim=-1)
@staticmethod
def hsv_to_rgb(hsv: torch.Tensor):
h, s, v = hsv.unbind(-1)
h = h * 6.0
i = torch.floor(h)
f = h - i
p = v * (1.0 - s)
q = v * (1.0 - s * f)
t = v * (1.0 - s * (1.0 - f))
i = i.long() % 6
mask = torch.stack(
(i == 0, i == 1, i == 2, i == 3, i == 4, i == 5), -1
)
rgb = torch.stack(
(
torch.where(
mask[..., 0],
v,
torch.where(
mask[..., 1],
q,
torch.where(
mask[..., 2],
p,
torch.where(
mask[..., 3],
p,
torch.where(mask[..., 4], t, v),
),
),
),
),
torch.where(
mask[..., 0],
t,
torch.where(
mask[..., 1],
v,
torch.where(
mask[..., 2],
v,
torch.where(
mask[..., 3],
q,
torch.where(mask[..., 4], p, p),
),
),
),
),
torch.where(
mask[..., 0],
p,
torch.where(
mask[..., 1],
p,
torch.where(
mask[..., 2],
t,
torch.where(
mask[..., 3],
v,
torch.where(mask[..., 4], v, q),
),
),
),
),
),
dim=-1,
)
return rgb
def correct(
self,
image: torch.Tensor,
force_gpu: bool,
clamp: bool,
gamma: float = 1.0,
contrast: float = 1.0,
exposure: float = 0.0,
offset: float = 0.0,
hue: float = 0.0,
saturation: float = 1.0,
value: float = 1.0,
mask: torch.Tensor | None = None,
):
device = self.get_device(image, force_gpu)
image = image.to(device)
if mask is not None:
if mask.shape[0] != image.shape[0]:
mask = mask.expand(image.shape[0], -1, -1)
mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3)
mask = mask.to(device)
model_management.throw_exception_if_processing_interrupted()
adjusted = image.pow(1 / gamma) * (2.0**exposure) * contrast + offset
model_management.throw_exception_if_processing_interrupted()
hsv = self.rgb_to_hsv(adjusted)
hsv[..., 0] = (hsv[..., 0] + hue) % 1.0 # Hue
hsv[..., 1] = hsv[..., 1] * saturation # Saturation
hsv[..., 2] = hsv[..., 2] * value # Value
adjusted = self.hsv_to_rgb(hsv)
model_management.throw_exception_if_processing_interrupted()
if clamp:
adjusted = torch.clamp(adjusted, 0.0, 1.0)
# apply mask
result = (
adjusted
if mask is None
else torch.where(mask > 0, adjusted, image)
)
if not force_gpu:
result = result.cpu()
return (result,)
class MTB_ColorCorrect:
"""Various color correction methods"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"clamp": ([True, False], {"default": True}),
"gamma": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"contrast": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"exposure": (
"FLOAT",
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
),
"offset": (
"FLOAT",
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
),
"hue": (
"FLOAT",
{"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.01},
),
"saturation": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
"value": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
),
},
"optional": {"mask": ("MASK",)},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "correct"
CATEGORY = "mtb/image processing"
@staticmethod
def gamma_correction_tensor(image, gamma):
gamma_inv = 1.0 / gamma
return image.pow(gamma_inv)
@staticmethod
def contrast_adjustment_tensor(image, contrast):
r, g, b = image.unbind(-1)
# Using Adobe RGB luminance weights.
luminance_image = 0.33 * r + 0.71 * g + 0.06 * b
luminance_mean = torch.mean(luminance_image.unsqueeze(-1))
# Blend original with mean luminance using contrast factor as blend ratio.
contrasted = image * contrast + (1.0 - contrast) * luminance_mean
return torch.clamp(contrasted, 0.0, 1.0)
@staticmethod
def exposure_adjustment_tensor(image, exposure):
return image * (2.0**exposure)
@staticmethod
def offset_adjustment_tensor(image, offset):
return image + offset
@staticmethod
def hsv_adjustment(image: torch.Tensor, hue, saturation, value):
images = tensor2pil(image)
out = []
for img in images:
hsv_image = img.convert("HSV")
h, s, v = hsv_image.split()
h = h.point(lambda x: (x + hue * 255) % 256)
s = s.point(lambda x: int(x * saturation))
v = v.point(lambda x: int(x * value))
hsv_image = Image.merge("HSV", (h, s, v))
rgb_image = hsv_image.convert("RGB")
out.append(rgb_image)
return pil2tensor(out)
@staticmethod
def hsv_adjustment_tensor_not_working(
image: torch.Tensor, hue, saturation, value
):
"""Abandonning for now"""
image = image.squeeze(0).permute(2, 0, 1)
max_val, _ = image.max(dim=0, keepdim=True)
min_val, _ = image.min(dim=0, keepdim=True)
delta = max_val - min_val
hue_image = torch.zeros_like(max_val)
mask = delta != 0.0
r, g, b = image[0], image[1], image[2]
hue_image[mask & (max_val == r)] = ((g - b) / delta)[
mask & (max_val == r)
] % 6.0
hue_image[mask & (max_val == g)] = ((b - r) / delta)[
mask & (max_val == g)
] + 2.0
hue_image[mask & (max_val == b)] = ((r - g) / delta)[
mask & (max_val == b)
] + 4.0
saturation_image = delta / (max_val + 1e-7)
value_image = max_val
hue_image = (hue_image + hue) % 1.0
saturation_image = torch.where(
mask, saturation * saturation_image, saturation_image
)
value_image = value * value_image
c = value_image * saturation_image
x = c * (1 - torch.abs((hue_image % 2) - 1))
m = value_image - c
prime_image = torch.zeros_like(image)
prime_image[0] = torch.where(
max_val == r, c, torch.where(max_val == g, x, prime_image[0])
)
prime_image[1] = torch.where(
max_val == r, x, torch.where(max_val == g, c, prime_image[1])
)
prime_image[2] = torch.where(
max_val == g, x, torch.where(max_val == b, c, prime_image[2])
)
rgb_image = prime_image + m
rgb_image = rgb_image.permute(1, 2, 0).unsqueeze(0)
return rgb_image
def correct(
self,
image: torch.Tensor,
clamp: bool,
gamma: float = 1.0,
contrast: float = 1.0,
exposure: float = 0.0,
offset: float = 0.0,
hue: float = 0.0,
saturation: float = 1.0,
value: float = 1.0,
mask: torch.Tensor | None = None,
):
if mask is not None:
if mask.shape[0] != image.shape[0]:
mask = mask.expand(image.shape[0], -1, -1)
mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3)
# Apply color correction operations
adjusted = self.gamma_correction_tensor(image, gamma)
adjusted = self.contrast_adjustment_tensor(adjusted, contrast)
adjusted = self.exposure_adjustment_tensor(adjusted, exposure)
adjusted = self.offset_adjustment_tensor(adjusted, offset)
adjusted = self.hsv_adjustment(adjusted, hue, saturation, value)
if clamp:
adjusted = torch.clamp(image, 0.0, 1.0)
result = (
adjusted
if mask is None
else torch.where(mask > 0, adjusted, image)
)
return (result,)
class MTB_ImageCompare:
"""Compare two images and return a difference image"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"imageA": ("IMAGE",),
"imageB": ("IMAGE",),
"mode": (
["checkerboard", "diff", "blend"],
{"default": "checkerboard"},
),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "compare"
CATEGORY = "mtb/image"
def compare(self, imageA: torch.Tensor, imageB: torch.Tensor, mode):
if imageA.dim() == 4:
batch_count = imageA.size(0)
return (
torch.cat(
tuple(
self.compare(imageA[i], imageB[i], mode)[0]
for i in range(batch_count)
),
dim=0,
),
)
num_channels_A = imageA.size(2)
num_channels_B = imageB.size(2)
# handle RGBA/RGB mismatch
if num_channels_A == 3 and num_channels_B == 4:
imageA = torch.cat(
(imageA, torch.ones_like(imageA[:, :, 0:1])), dim=2
)
elif num_channels_B == 3 and num_channels_A == 4:
imageB = torch.cat(
(imageB, torch.ones_like(imageB[:, :, 0:1])), dim=2
)
match mode:
case "diff":
compare_image = torch.abs(imageA - imageB)
case "blend":
compare_image = 0.5 * (imageA + imageB)
case "checkerboard":
imageA = imageA.numpy()
imageB = imageB.numpy()
compared_channels = [
torch.from_numpy(
compare_images(
imageA[:, :, i], imageB[:, :, i], method=mode
)
)
for i in range(imageA.shape[2])
]
compare_image = torch.stack(compared_channels, dim=2)
case _:
compare_image = None
raise ValueError(f"Unknown mode {mode}")
compare_image = compare_image.unsqueeze(0)
return (compare_image,)
import requests
class MTB_LoadImageFromUrl:
"""Load an image from the given URL"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"url": (
"STRING",
{
"default": "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg"
},
),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "load"
CATEGORY = "mtb/IO"
def load(self, url):
# get the image from the url
image = Image.open(requests.get(url, stream=True).raw)
image = ImageOps.exif_transpose(image)
return (pil2tensor(image),)
class MTB_Blur:
"""Blur an image using a Gaussian filter."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"sigmaX": (
"FLOAT",
{"default": 3.0, "min": 0.0, "max": 200.0, "step": 0.01},
),
"sigmaY": (
"FLOAT",
{"default": 3.0, "min": 0.0, "max": 200.0, "step": 0.01},
),
},
"optional": {"sigmasX": ("FLOATS",), "sigmasY": ("FLOATS",)},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "blur"
CATEGORY = "mtb/image processing"
def blur(
self, image: torch.Tensor, sigmaX, sigmaY, sigmasX=None, sigmasY=None
):
image_np = image.numpy() * 255
blurred_images = []
if sigmasX is not None:
if sigmasY is None:
sigmasY = sigmasX
if len(sigmasX) != image.size(0):
raise ValueError(
f"SigmasX must have same length as image, sigmasX is {len(sigmasX)} but the batch size is {image.size(0)}"
)
for i in range(image.size(0)):
blurred = gaussian(
image_np[i],
sigma=(sigmasX[i], sigmasY[i], 0),
channel_axis=2,
)
blurred_images.append(blurred)
image_np = np.array(blurred_images)
else:
for i in range(image.size(0)):
blurred = gaussian(
image_np[i], sigma=(sigmaX, sigmaY, 0), channel_axis=2
)
blurred_images.append(blurred)
image_np = np.array(blurred_images)
return (np2tensor(image_np).squeeze(0),)
class MTB_Sharpen:
"""Sharpens an image using a Gaussian kernel."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"sharpen_radius": (
"INT",
{"default": 1, "min": 1, "max": 31, "step": 1},
),
"sigma_x": (
"FLOAT",
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
),
"sigma_y": (
"FLOAT",
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
),
"alpha": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1},
),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_sharp"
CATEGORY = "mtb/image processing"
def do_sharp(
self,
image: torch.Tensor,
sharpen_radius: int,
sigma_x: float,
sigma_y: float,
alpha: float,
):
if sharpen_radius == 0:
return (image,)
channels = image.shape[3]
kernel_size = 2 * sharpen_radius + 1
kernel = gaussian_kernel(kernel_size, sigma_x, sigma_y) * -(alpha * 10)
# Modify center of kernel to make it a sharpening kernel
center = kernel_size // 2
kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0
kernel = kernel.repeat(channels, 1, 1).unsqueeze(1)
tensor_image = image.permute(0, 3, 1, 2)
tensor_image = F.pad(
tensor_image,
(sharpen_radius, sharpen_radius, sharpen_radius, sharpen_radius),
"reflect",
)
sharpened = F.conv2d(
tensor_image, kernel, padding=center, groups=channels
)
# Remove padding
sharpened = sharpened[
:,
:,
sharpen_radius:-sharpen_radius,
sharpen_radius:-sharpen_radius,
]
sharpened = sharpened.permute(0, 2, 3, 1)
result = torch.clamp(sharpened, 0, 1)
return (result,)
# https://github.com/lllyasviel/AdverseCleaner/blob/main/clean.py
# def deglaze_np_img(np_img):
# y = np_img.copy()
# for _ in range(64):
# y = cv2.bilateralFilter(y, 5, 8, 8)
# for _ in range(4):
# y = guidedFilter(np_img, y, 4, 16)
# return y
# class DeglazeImage:
# """Remove adversarial noise from images"""
# @classmethod
# def INPUT_TYPES(cls):
# return {"required": {"image": ("IMAGE",)}}
# CATEGORY = "mtb/image processing"
# RETURN_TYPES = ("IMAGE",)
# FUNCTION = "deglaze_image"
# def deglaze_image(self, image):
# return (np2tensor(deglaze_np_img(tensor2np(image))),)
class MTB_MaskToImage:
"""Converts a mask (alpha) to an RGB image with a color and background"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mask": ("MASK",),
"color": ("COLOR",),
"background": ("COLOR", {"default": "#000000"}),
},
"optional": {
"invert": ("BOOLEAN", {"default": False}),
},
}
CATEGORY = "mtb/generate"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "render_mask"
def render_mask(self, mask, color, background, invert=False):
masks = tensor2pil(1.0 - mask) if invert else tensor2pil(mask)
images = []
for m in masks:
_mask = m.convert("L")
log.debug(
f"Converted mask to PIL Image format, size: {_mask.size}"
)
image = Image.new("RGBA", _mask.size, color=color)
# apply the mask
image = Image.composite(
image, Image.new("RGBA", _mask.size, color=background), _mask
)
# image = ImageChops.multiply(image, mask)
# apply over background
# image = Image.alpha_composite(Image.new("RGBA", image.size, color=background), image)
images.append(image.convert("RGB"))
return (pil2tensor(images),)
class MTB_ColoredImage:
"""Constant color image of given size."""
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"color": ("COLOR",),
"width": ("INT", {"default": 512, "min": 16, "max": 8160}),
"height": ("INT", {"default": 512, "min": 16, "max": 8160}),
},
"optional": {
"foreground_image": ("IMAGE",),
"foreground_mask": ("MASK",),
"invert": ("BOOLEAN", {"default": False}),
"mask_opacity": (
"FLOAT",
{"default": 1.0, "step": 0.1, "min": 0},
),
},
}
CATEGORY = "mtb/generate"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "render_img"
def resize_and_crop(self, img: Image.Image, target_size: tuple[int, int]):
scale = max(target_size[0] / img.width, target_size[1] / img.height)
new_size = (int(img.width * scale), int(img.height * scale))
img = img.resize(new_size, Image.LANCZOS)
left = (img.width - target_size[0]) // 2
top = (img.height - target_size[1]) // 2
return img.crop(
(left, top, left + target_size[0], top + target_size[1])
)
def resize_and_crop_thumbnails(
self, img: Image.Image, target_size: tuple[int, int]
):
img.thumbnail(target_size, Image.LANCZOS)
left = (img.width - target_size[0]) / 2
top = (img.height - target_size[1]) / 2
right = (img.width + target_size[0]) / 2
bottom = (img.height + target_size[1]) / 2
return img.crop((left, top, right, bottom))
@staticmethod
def process_mask(
mask: torch.Tensor | None,
invert: bool,
# opacity: float,
batch_size: int,
) -> list[Image.Image] | None:
if mask is None:
return [None] * batch_size
masks = tensor2pil(mask if not invert else 1.0 - mask)
if len(masks) == 1 and batch_size > 1:
masks = masks * batch_size
if len(masks) != batch_size:
raise ValueError(
"Foreground image and mask must have the same batch size"
)
return masks
def render_img(
self,
color: str,
width: int,
height: int,
foreground_image: torch.Tensor | None = None,
foreground_mask: torch.Tensor | None = None,
invert: bool = False,
mask_opacity: float = 1.0,
) -> tuple[torch.Tensor]:
background = Image.new("RGBA", (width, height), color=color)
if foreground_image is None:
return (pil2tensor([background.convert("RGB")]),)
fg_images = tensor2pil(foreground_image)
fg_masks = self.process_mask(foreground_mask, invert, len(fg_images))
output: list[Image.Image] = []
for fg_image, fg_mask in zip(fg_images, fg_masks, strict=False):
fg_image = self.resize_and_crop(fg_image, background.size)
if fg_mask:
fg_mask = self.resize_and_crop(fg_mask, background.size)
fg_mask_array = np.array(fg_mask)
fg_mask_array = (fg_mask_array * mask_opacity).astype(np.uint8)
fg_mask = Image.fromarray(fg_mask_array)
output.append(
Image.composite(
fg_image.convert("RGBA"), background, fg_mask
).convert("RGB")
)
else:
if fg_image.mode != "RGBA":
raise ValueError(
f"Foreground image must be in 'RGBA' mode when no mask is provided, got {fg_image.mode}"
)
output.append(
Image.alpha_composite(background, fg_image).convert("RGB")
)
return (pil2tensor(output),)
class MTB_ImagePremultiply:
"""Premultiply image with mask"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"invert": ("BOOLEAN", {"default": False}),
}
}
CATEGORY = "mtb/image"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("RGBA",)
FUNCTION = "premultiply"
def premultiply(self, image, mask, invert):
images = tensor2pil(image)
masks = tensor2pil(mask) if invert else tensor2pil(1.0 - mask)
single = len(mask) == 1
masks = [x.convert("L") for x in masks]
out = []
for i, img in enumerate(images):
cur_mask = masks[0] if single else masks[i]
img.putalpha(cur_mask)
out.append(img)
# if invert:
# image = Image.composite(image,Image.new("RGBA", image.size, color=(0,0,0,0)), mask)
# else:
# image = Image.composite(Image.new("RGBA", image.size, color=(0,0,0,0)), image, mask)
return (pil2tensor(out),)
class MTB_ImageResizeFactor:
"""Extracted mostly from WAS Node Suite, with a few edits (most notably multiple image support) and less features."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"factor": (
"FLOAT",
{"default": 2, "min": 0.01, "max": 16.0, "step": 0.01},
),
"supersample": ("BOOLEAN", {"default": True}),
"resampling": (
[
"nearest",
"linear",
"bilinear",
"bicubic",
"trilinear",
"area",
"nearest-exact",
],
{"default": "nearest"},
),
},
"optional": {
"mask": ("MASK",),
},
}
CATEGORY = "mtb/image"
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "resize"
def resize(
self,
image: torch.Tensor,
factor: float,
supersample: bool,
resampling: str,
mask=None,
):
# Check if the tensor has the correct dimension
if len(image.shape) not in [3, 4]: # HxWxC or BxHxWxC
raise ValueError(
"Expected image tensor of shape (H, W, C) or (B, H, W, C)"
)
# Transpose to CxHxW or BxCxHxW for PyTorch
if len(image.shape) == 3:
image = image.permute(2, 0, 1).unsqueeze(0) # CxHxW
else:
image = image.permute(0, 3, 1, 2) # BxCxHxW
# Compute new dimensions
B, C, H, W = image.shape
new_H, new_W = int(H * factor), int(W * factor)
align_corner_filters = ("linear", "bilinear", "bicubic", "trilinear")
# Resize the image
resized_image = F.interpolate(
image,
size=(new_H, new_W),
mode=resampling,
align_corners=resampling in align_corner_filters,
)
# Optionally supersample
if supersample:
resized_image = F.interpolate(
resized_image,
scale_factor=2,
mode=resampling,
align_corners=resampling in align_corner_filters,
)
# Transpose back to the original format: BxHxWxC or HxWxC
if len(image.shape) == 4:
resized_image = resized_image.permute(0, 2, 3, 1)
else:
resized_image = resized_image.squeeze(0).permute(1, 2, 0)
# Apply mask if provided
if mask is not None:
if len(mask.shape) != len(resized_image.shape):
raise ValueError(
"Mask tensor should have the same dimensions as the image tensor"
)
resized_image = resized_image * mask
return (resized_image,)
class MTB_SaveImageGrid:
"""Save all the images in the input batch as a grid of images."""
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
self.type = "output"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
"save_intermediate": ("BOOLEAN", {"default": False}),
},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
RETURN_TYPES = ()
FUNCTION = "save_images"
OUTPUT_NODE = True
CATEGORY = "mtb/IO"
def create_image_grid(self, image_list):
total_images = len(image_list)
# Calculate the grid size based on the square root of the total number of images
grid_size = (
int(math.sqrt(total_images)),
int(math.ceil(math.sqrt(total_images))),
)
# Get the size of the first image to determine the grid size
image_width, image_height = image_list[0].size
# Create a new blank image to hold the grid
grid_width = grid_size[0] * image_width
grid_height = grid_size[1] * image_height
grid_image = Image.new("RGB", (grid_width, grid_height))
# Iterate over the images and paste them onto the grid
for i, image in enumerate(image_list):
x = (i % grid_size[0]) * image_width
y = (i // grid_size[0]) * image_height
grid_image.paste(image, (x, y, x + image_width, y + image_height))
return grid_image
def save_images(
self,
images,
filename_prefix="Grid",
save_intermediate=False,
prompt=None,
extra_pnginfo=None,
):
(
full_output_folder,
filename,
counter,
subfolder,
filename_prefix,
) = folder_paths.get_save_image_path(
filename_prefix,
self.output_dir,
images[0].shape[1],
images[0].shape[0],
)
image_list = []
batch_counter = counter
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
for idx, image in enumerate(images):
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
image_list.append(img)
if save_intermediate:
file = f"{filename}_batch-{idx:03}_{batch_counter:05}_.png"
img.save(
os.path.join(full_output_folder, file),
pnginfo=metadata,
compress_level=4,
)
batch_counter += 1
file = f"{filename}_{counter:05}_.png"
grid = self.create_image_grid(image_list)
grid.save(
os.path.join(full_output_folder, file),
pnginfo=metadata,
compress_level=4,
)
results = [
{"filename": file, "subfolder": subfolder, "type": self.type}
]
return {"ui": {"images": results}}
class MTB_ImageTileOffset:
"""Mimics an old photoshop technique to check for seamless textures"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"tilesX": ("INT", {"default": 2, "min": 1}),
"tilesY": ("INT", {"default": 2, "min": 1}),
}
}
CATEGORY = "mtb/generate"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "tile_image"
def tile_image(
self, image: torch.Tensor, tilesX: int = 2, tilesY: int = 2
):
if tilesX < 1 or tilesY < 1:
raise ValueError("The number of tiles must be at least 1.")
batch_size, height, width, channels = image.shape
tile_height = height // tilesY
tile_width = width // tilesX
output_image = torch.zeros_like(image)
for i, j in itertools.product(range(tilesY), range(tilesX)):
start_h = i * tile_height
end_h = start_h + tile_height
start_w = j * tile_width
end_w = start_w + tile_width
tile = image[:, start_h:end_h, start_w:end_w, :]
output_start_h = (i + 1) % tilesY * tile_height
output_start_w = (j + 1) % tilesX * tile_width
output_end_h = output_start_h + tile_height
output_end_w = output_start_w + tile_width
output_image[
:, output_start_h:output_end_h, output_start_w:output_end_w, :
] = tile
return (output_image,)
__nodes__ = [
MTB_ColorCorrect,
MTB_ColorCorrectGPU,
MTB_ImageCompare,
MTB_ImageTileOffset,
MTB_Blur,
# DeglazeImage,
MTB_MaskToImage,
MTB_ColoredImage,
MTB_ImagePremultiply,
MTB_ImageResizeFactor,
MTB_SaveImageGrid,
MTB_LoadImageFromUrl,
MTB_Sharpen,
MTB_ExtractCoordinatesFromImage,
MTB_CoordinatesToString,
]