fix: 🐛 support batch masks (colored image node)

Also replace ANTIALIAS with LANCZOS.

closes #156
This commit is contained in:
Mel Massadian
2024-03-13 16:57:31 +01:00
parent 48f91b74e2
commit 2465ffb0d3
+5 -5
View File
@@ -3,7 +3,6 @@ import json
import math
import os
import cv2
import folder_paths
import numpy as np
import torch
@@ -476,7 +475,7 @@ class ColoredImage:
# Resize the image based on calculated scale
new_size = (int(img.width * scale), int(img.height * scale))
img = img.resize(new_size, Image.ANTIALIAS)
img = img.resize(new_size, Image.LANCZOS)
# Calculate cropping coordinates
left = (img.width - target_size[0]) / 2
@@ -488,7 +487,7 @@ class ColoredImage:
return img.crop((left, top, right, bottom))
def resize_and_crop_thumbnails(self, img, target_size):
img.thumbnail(target_size, Image.ANTIALIAS)
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
@@ -516,7 +515,7 @@ class ColoredImage:
raise ValueError(
"Foreground image and mask must have same batch size"
)
fg_masks = tensor2pil(foreground_mask)
fg_masks = tensor2pil(foreground_mask.unsqueeze(-1))
for fg_image, fg_mask in zip(fg_images, fg_masks):
# Resize and crop if dimensions mismatch
@@ -536,7 +535,8 @@ class ColoredImage:
else:
if fg_image.mode != "RGBA":
raise ValueError(
"Foreground image must be in 'RGBA' mode when no mask is provided, got {fg_image.mode}"
"Foreground image must be in 'RGBA' mode "
f"when no mask is provided, got {fg_image.mode}"
)
output.append(
Image.alpha_composite(image, fg_image).convert("RGB")