245 lines
10 KiB
Python
245 lines
10 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def parse_hex_color(hex_str: str) -> tuple:
|
|
hex_str = hex_str.strip().lstrip('#')
|
|
if len(hex_str) == 3:
|
|
hex_str = ''.join([c*2 for c in hex_str])
|
|
try:
|
|
if len(hex_str) == 6:
|
|
return (int(hex_str[0:2], 16), int(hex_str[2:4], 16), int(hex_str[4:6], 16))
|
|
elif len(hex_str) == 8:
|
|
return (int(hex_str[0:2], 16), int(hex_str[2:4], 16), int(hex_str[4:6], 16))
|
|
except ValueError:
|
|
pass
|
|
return (0, 0, 0)
|
|
|
|
if hasattr(Image, "Resampling"):
|
|
PIL_RESAMPLING = {
|
|
"lanczos": Image.Resampling.LANCZOS,
|
|
"bicubic": Image.Resampling.BICUBIC,
|
|
"hamming": Image.Resampling.HAMMING,
|
|
"bilinear": Image.Resampling.BILINEAR,
|
|
"nearest": Image.Resampling.NEAREST,
|
|
}
|
|
else:
|
|
PIL_RESAMPLING = {
|
|
"lanczos": Image.LANCZOS,
|
|
"bicubic": Image.BICUBIC,
|
|
"hamming": Image.HAMMING,
|
|
"bilinear": Image.BILINEAR,
|
|
"nearest": Image.NEAREST,
|
|
}
|
|
|
|
class JetImageScale:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"aspect_ratio": (["original", "custom", "1:1", "3:2", "4:3", "16:9", "2:3", "3:4", "9:16"], {"default": "original"}),
|
|
"proportional_width": ("INT", {"default": 1, "min": 1, "max": 99999, "step": 1}),
|
|
"proportional_height": ("INT", {"default": 1, "min": 1, "max": 99999, "step": 1}),
|
|
"fit": (["crop", "letterbox", "fill"], {"default": "crop"}),
|
|
"method": (["lanczos", "bicubic", "hamming", "bilinear", "nearest"], {"default": "lanczos"}),
|
|
"scale_to_side": (["none", "longest", "shortest", "width", "height"], {"default": "longest"}),
|
|
"scale_to_length": ("INT", {"default": 1280, "min": 1, "max": 99999, "step": 1}),
|
|
"background_color": ("STRING", {"default": "#000000"}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "TUPLE", "INT", "INT")
|
|
RETURN_NAMES = ("image", "mask", "original_size", "width", "height")
|
|
FUNCTION = "scale_image"
|
|
CATEGORY = "Image/Scale"
|
|
|
|
def scale_image(self, image, aspect_ratio, proportional_width, proportional_height,
|
|
fit, method, scale_to_side, scale_to_length, background_color, mask=None):
|
|
B, H, W, C = image.shape
|
|
orig_w = W
|
|
orig_h = H
|
|
|
|
# 1. Determine target aspect ratio
|
|
if aspect_ratio == "original":
|
|
target_ar = orig_w / orig_h
|
|
elif aspect_ratio == "custom":
|
|
target_ar = proportional_width / proportional_height
|
|
elif aspect_ratio == "1:1":
|
|
target_ar = 1.0
|
|
elif aspect_ratio == "3:2":
|
|
target_ar = 1.5
|
|
elif aspect_ratio == "4:3":
|
|
target_ar = 4.0 / 3.0
|
|
elif aspect_ratio == "16:9":
|
|
target_ar = 16.0 / 9.0
|
|
elif aspect_ratio == "2:3":
|
|
target_ar = 2.0 / 3.0
|
|
elif aspect_ratio == "3:4":
|
|
target_ar = 3.0 / 4.0
|
|
elif aspect_ratio == "9:16":
|
|
target_ar = 9.0 / 16.0
|
|
else:
|
|
target_ar = orig_w / orig_h
|
|
|
|
# 2. Determine target dimensions
|
|
if scale_to_side == "none":
|
|
orig_ar = orig_w / orig_h
|
|
if target_ar >= orig_ar:
|
|
target_h = orig_h
|
|
target_w = max(1, int(target_h * target_ar))
|
|
else:
|
|
target_w = orig_w
|
|
target_h = max(1, int(target_w / target_ar))
|
|
elif scale_to_side == "width":
|
|
target_w = scale_to_length
|
|
target_h = max(1, int(target_w / target_ar))
|
|
elif scale_to_side == "height":
|
|
target_h = scale_to_length
|
|
target_w = max(1, int(target_h * target_ar))
|
|
elif scale_to_side == "longest":
|
|
if target_ar >= 1.0:
|
|
target_w = scale_to_length
|
|
target_h = max(1, int(target_w / target_ar))
|
|
else:
|
|
target_h = scale_to_length
|
|
target_w = max(1, int(target_h * target_ar))
|
|
elif scale_to_side == "shortest":
|
|
if target_ar >= 1.0:
|
|
target_h = scale_to_length
|
|
target_w = max(1, int(target_h * target_ar))
|
|
else:
|
|
target_w = scale_to_length
|
|
target_h = max(1, int(target_w / target_ar))
|
|
else:
|
|
target_w = scale_to_length
|
|
target_h = max(1, int(target_w / target_ar))
|
|
|
|
# 3. Normalize mask input
|
|
if mask is not None:
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0).repeat(B, 1, 1)
|
|
elif len(mask.shape) == 3:
|
|
if mask.shape[0] != B:
|
|
if mask.shape[0] == 1:
|
|
mask = mask.repeat(B, 1, 1)
|
|
else:
|
|
mask = mask[:B]
|
|
else:
|
|
mask = None
|
|
|
|
resampling_method = PIL_RESAMPLING.get(method.lower(), Image.Resampling.LANCZOS if hasattr(Image, "Resampling") else Image.LANCZOS)
|
|
bg_color_tuple = parse_hex_color(background_color)
|
|
|
|
output_images = []
|
|
output_masks = []
|
|
|
|
for i in range(B):
|
|
img_tensor = image[i]
|
|
img_np = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
|
|
if C == 4:
|
|
pil_img = Image.fromarray(img_np, mode="RGBA")
|
|
elif C == 3:
|
|
pil_img = Image.fromarray(img_np, mode="RGB")
|
|
elif C == 1:
|
|
pil_img = Image.fromarray(img_np[:, :, 0], mode="L")
|
|
else:
|
|
pil_img = Image.fromarray(img_np, mode="RGB")
|
|
|
|
if mask is not None:
|
|
mask_tensor = mask[i]
|
|
mask_np = (mask_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
pil_mask = Image.fromarray(mask_np, mode="L")
|
|
else:
|
|
pil_mask = None
|
|
|
|
if fit == "fill":
|
|
resized_img = pil_img.resize((target_w, target_h), resample=resampling_method)
|
|
if pil_mask is not None:
|
|
resized_mask = pil_mask.resize((target_w, target_h), resample=resampling_method)
|
|
else:
|
|
resized_mask = Image.new("L", (target_w, target_h), color=0)
|
|
|
|
elif fit == "crop":
|
|
s = max(target_w / orig_w, target_h / orig_h)
|
|
new_w = max(1, int(orig_w * s))
|
|
new_h = max(1, int(orig_h * s))
|
|
|
|
temp_img = pil_img.resize((new_w, new_h), resample=resampling_method)
|
|
if pil_mask is not None:
|
|
temp_mask = pil_mask.resize((new_w, new_h), resample=resampling_method)
|
|
else:
|
|
temp_mask = Image.new("L", (new_w, new_h), color=0)
|
|
|
|
left = (new_w - target_w) // 2
|
|
top = (new_h - target_h) // 2
|
|
right = left + target_w
|
|
bottom = top + target_h
|
|
|
|
resized_img = temp_img.crop((left, top, right, bottom))
|
|
resized_mask = temp_mask.crop((left, top, right, bottom))
|
|
|
|
elif fit == "letterbox":
|
|
s = min(target_w / orig_w, target_h / orig_h)
|
|
new_w = max(1, int(orig_w * s))
|
|
new_h = max(1, int(orig_h * s))
|
|
|
|
temp_img = pil_img.resize((new_w, new_h), resample=resampling_method)
|
|
if pil_mask is not None:
|
|
temp_mask = pil_mask.resize((new_w, new_h), resample=resampling_method)
|
|
else:
|
|
temp_mask = Image.new("L", (new_w, new_h), color=0)
|
|
|
|
if C == 4:
|
|
alpha = bg_color_tuple[3] if len(bg_color_tuple) > 3 else 255
|
|
canvas_color = (bg_color_tuple[0], bg_color_tuple[1], bg_color_tuple[2], alpha)
|
|
resized_img = Image.new("RGBA", (target_w, target_h), color=canvas_color)
|
|
elif C == 3:
|
|
resized_img = Image.new("RGB", (target_w, target_h), color=bg_color_tuple)
|
|
elif C == 1:
|
|
gray_val = int(0.299 * bg_color_tuple[0] + 0.587 * bg_color_tuple[1] + 0.114 * bg_color_tuple[2])
|
|
resized_img = Image.new("L", (target_w, target_h), color=gray_val)
|
|
else:
|
|
resized_img = Image.new("RGB", (target_w, target_h), color=bg_color_tuple)
|
|
|
|
resized_mask = Image.new("L", (target_w, target_h), color=0)
|
|
|
|
left = (target_w - new_w) // 2
|
|
top = (target_h - new_h) // 2
|
|
|
|
resized_img.paste(temp_img, (left, top))
|
|
resized_mask.paste(temp_mask, (left, top))
|
|
else:
|
|
resized_img = pil_img.resize((target_w, target_h), resample=resampling_method)
|
|
if pil_mask is not None:
|
|
resized_mask = pil_mask.resize((target_w, target_h), resample=resampling_method)
|
|
else:
|
|
resized_mask = Image.new("L", (target_w, target_h), color=0)
|
|
|
|
out_img_np = np.array(resized_img).astype(np.float32) / 255.0
|
|
if C == 1 and len(out_img_np.shape) == 2:
|
|
out_img_np = np.expand_dims(out_img_np, axis=-1)
|
|
|
|
out_mask_np = np.array(resized_mask).astype(np.float32) / 255.0
|
|
|
|
output_images.append(torch.from_numpy(out_img_np))
|
|
output_masks.append(torch.from_numpy(out_mask_np))
|
|
|
|
stacked_images = torch.stack(output_images, dim=0)
|
|
stacked_masks = torch.stack(output_masks, dim=0)
|
|
original_size = (orig_w, orig_h)
|
|
|
|
return (stacked_images, stacked_masks, original_size, target_w, target_h)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"JetImageScale": JetImageScale
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"JetImageScale": "Jet Image Scale"
|
|
}
|