fix: 🐛 colored image mask input
- mask input was completely broken, now it works. - size can now be different from the input images, it will crop or expand accordingly. - the image and mask tensors still must be of the same size
This commit is contained in:
melMass
+102
-37
@@ -22,7 +22,9 @@ from ..utils import pil2tensor, tensor2np, tensor2pil
|
||||
# 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):
|
||||
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),
|
||||
@@ -115,7 +117,9 @@ class ColorCorrect:
|
||||
return pil2tensor(out)
|
||||
|
||||
@staticmethod
|
||||
def hsv_adjustment_tensor_not_working(image: torch.Tensor, hue, saturation, value):
|
||||
def hsv_adjustment_tensor_not_working(
|
||||
image: torch.Tensor, hue, saturation, value
|
||||
):
|
||||
"""Abandonning for now"""
|
||||
image = image.squeeze(0).permute(2, 0, 1)
|
||||
|
||||
@@ -344,11 +348,16 @@ class Sharpen_:
|
||||
(sharpen_radius, sharpen_radius, sharpen_radius, sharpen_radius),
|
||||
"reflect",
|
||||
)
|
||||
sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)
|
||||
sharpened = F.conv2d(
|
||||
tensor_image, kernel, padding=center, groups=channels
|
||||
)
|
||||
|
||||
# Remove padding
|
||||
sharpened = sharpened[
|
||||
:, :, sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius
|
||||
:,
|
||||
:,
|
||||
sharpen_radius:-sharpen_radius,
|
||||
sharpen_radius:-sharpen_radius,
|
||||
]
|
||||
|
||||
sharpened = sharpened.permute(0, 2, 3, 1)
|
||||
@@ -408,7 +417,9 @@ class MaskToImage:
|
||||
for m in masks:
|
||||
_mask = Image.fromarray(m).convert("L")
|
||||
|
||||
log.debug(f"Converted mask to PIL Image format, size: {_mask.size}")
|
||||
log.debug(
|
||||
f"Converted mask to PIL Image format, size: {_mask.size}"
|
||||
)
|
||||
|
||||
image = Image.new("RGBA", _mask.size, color=color)
|
||||
# apply the mask
|
||||
@@ -425,8 +436,11 @@ class MaskToImage:
|
||||
return (pil2tensor(images),)
|
||||
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class ColoredImage:
|
||||
"""Constant color image of given size"""
|
||||
"""Constant color image of given size."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
@@ -451,40 +465,82 @@ class ColoredImage:
|
||||
|
||||
FUNCTION = "render_img"
|
||||
|
||||
def resize_and_crop(self, img, target_size):
|
||||
# Calculate scaling factors for both dimensions
|
||||
scale_x = target_size[0] / img.width
|
||||
scale_y = target_size[1] / img.height
|
||||
|
||||
# Use the smaller scaling factor to maintain aspect ratio
|
||||
scale = max(scale_x, scale_y)
|
||||
|
||||
# Resize the image based on calculated scale
|
||||
new_size = (int(img.width * scale), int(img.height * scale))
|
||||
img = img.resize(new_size, Image.ANTIALIAS)
|
||||
|
||||
# Calculate cropping coordinates
|
||||
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
|
||||
|
||||
# Crop and return the image
|
||||
return img.crop((left, top, right, bottom))
|
||||
|
||||
def resize_and_crop_thumbnails(self, img, target_size):
|
||||
img.thumbnail(target_size, Image.ANTIALIAS)
|
||||
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))
|
||||
|
||||
def render_img(
|
||||
self, color, width, height, foreground_image=None, foreground_mask=None
|
||||
self,
|
||||
color,
|
||||
width,
|
||||
height,
|
||||
foreground_image: Optional[torch.Tensor] = None,
|
||||
foreground_mask: Optional[torch.Tensor] = None,
|
||||
):
|
||||
image = Image.new("RGBA", (width, height), color=color)
|
||||
output = []
|
||||
if foreground_image is not None:
|
||||
if foreground_mask is None:
|
||||
fg_images = tensor2pil(foreground_image)
|
||||
for img in fg_images:
|
||||
if image.size != img.size:
|
||||
raise ValueError(
|
||||
f"Dimension mismatch: image {image.size}, img {img.size}"
|
||||
)
|
||||
fg_images = tensor2pil(foreground_image)
|
||||
fg_masks = [None] * len(
|
||||
fg_images
|
||||
) # Default to None for each foreground image
|
||||
|
||||
if img.mode != "RGBA":
|
||||
raise ValueError(
|
||||
f"Foreground image must be in 'RGBA' mode when no mask is provided, got {img.mode}"
|
||||
)
|
||||
|
||||
output.append(Image.alpha_composite(image, img).convert("RGB"))
|
||||
|
||||
elif foreground_image.size[0] != foreground_mask.size[0]:
|
||||
raise ValueError("Foreground image and mask must have same batch size")
|
||||
else:
|
||||
fg_images = tensor2pil(foreground_image)
|
||||
if foreground_mask is not None:
|
||||
if foreground_image.size()[0] != foreground_mask.size()[0]:
|
||||
raise ValueError(
|
||||
"Foreground image and mask must have same batch size"
|
||||
)
|
||||
fg_masks = tensor2pil(foreground_mask)
|
||||
output.extend(
|
||||
Image.composite(
|
||||
fg_image.convert("RGBA"),
|
||||
image,
|
||||
fg_mask,
|
||||
).convert("RGB")
|
||||
for fg_image, fg_mask in zip(fg_images, fg_masks)
|
||||
)
|
||||
|
||||
for fg_image, fg_mask in zip(fg_images, fg_masks):
|
||||
# Resize and crop if dimensions mismatch
|
||||
if fg_image.size != image.size:
|
||||
fg_image = self.resize_and_crop(fg_image, image.size)
|
||||
if fg_mask:
|
||||
fg_mask = self.resize_and_crop(fg_mask, image.size)
|
||||
|
||||
if fg_mask:
|
||||
output.append(
|
||||
Image.composite(
|
||||
fg_image.convert("RGBA"),
|
||||
image,
|
||||
fg_mask,
|
||||
).convert("RGB")
|
||||
)
|
||||
else:
|
||||
if fg_image.mode != "RGBA":
|
||||
raise ValueError(
|
||||
"Foreground image must be in 'RGBA' mode when no mask is provided, got {fg_image.mode}"
|
||||
)
|
||||
output.append(
|
||||
Image.alpha_composite(image, fg_image).convert("RGB")
|
||||
)
|
||||
|
||||
elif foreground_mask is not None:
|
||||
log.warn("Mask ignored because no foreground image is given")
|
||||
|
||||
@@ -577,7 +633,9 @@ class ImageResizeFactor:
|
||||
):
|
||||
# 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)")
|
||||
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:
|
||||
@@ -689,7 +747,10 @@ class SaveImageGrid_:
|
||||
subfolder,
|
||||
filename_prefix,
|
||||
) = folder_paths.get_save_image_path(
|
||||
filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
|
||||
filename_prefix,
|
||||
self.output_dir,
|
||||
images[0].shape[1],
|
||||
images[0].shape[0],
|
||||
)
|
||||
image_list = []
|
||||
batch_counter = counter
|
||||
@@ -719,10 +780,14 @@ class SaveImageGrid_:
|
||||
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
|
||||
os.path.join(full_output_folder, file),
|
||||
pnginfo=metadata,
|
||||
compress_level=4,
|
||||
)
|
||||
|
||||
results = [{"filename": file, "subfolder": subfolder, "type": self.type}]
|
||||
results = [
|
||||
{"filename": file, "subfolder": subfolder, "type": self.type}
|
||||
]
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user