-
Notifications
You must be signed in to change notification settings - Fork 1
/
overlay_image_node.py
83 lines (67 loc) · 4.24 KB
/
overlay_image_node.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import torch
import numpy as np
from PIL import Image
class ScaleImageNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"background_image": ("IMAGE", {"default": None}),
"scale_x": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
"scale_y": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
"position_x": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"position_y": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"rotation_angle": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "scale_image"
CATEGORY = "Roblox"
def scale_image(self, image, background_image, scale_x, scale_y, position_x, position_y, rotation_angle):
# Convertir la imagen principal de tensor a matriz NumPy
if isinstance(image, torch.Tensor):
image = image.squeeze().cpu().numpy() # Asegurarse de que el tensor esté en CPU y comprimido
# Verificar los canales y convertir a RGBA si es necesario
if image.shape[-1] == 3:
# Crear un canal alfa basado en la intensidad de los píxeles, manteniendo opacos los píxeles de la imagen
alpha_channel = np.where(np.sum(image, axis=-1, keepdims=True) > 0, 1.0, 0.0)
image = np.concatenate((image, alpha_channel), axis=-1)
elif image.shape[-1] == 4:
# Mantener el canal alfa de la imagen si ya tiene 4 canales
image[..., 3] = np.where(image[..., 3] > 0, 1.0, 0.0) # Asegurar que el canal alfa sea 0 o 1
else:
raise ValueError(f"Se esperaba una imagen con 3 o 4 canales, pero tiene {image.shape[-1]}.")
# Convertir la imagen principal a PIL
pil_image = Image.fromarray((image * 255).astype(np.uint8), mode='RGBA')
# Aplicar el escalado a la imagen principal
width, height = pil_image.size
new_width = int(width * scale_x)
new_height = int(height * scale_y)
scaled_image = pil_image.resize((new_width, new_height), Image.LANCZOS) # Usar LANCZOS en lugar de ANTIALIAS
# Aplicar la rotación
rotated_image = scaled_image.rotate(rotation_angle, expand=True)
# Convertir la imagen de fondo a PIL sin cambiar su tamaño
if background_image is not None:
if isinstance(background_image, torch.Tensor):
background_image = background_image.squeeze().cpu().numpy() # Comprimir y enviar a CPU
if background_image.shape[-1] == 3:
background_image = np.concatenate((background_image, np.ones((background_image.shape[0], background_image.shape[1], 1), dtype=np.float32)), axis=-1)
elif background_image.shape[-1] != 4:
raise ValueError(f"Se esperaba una imagen de fondo con 3 o 4 canales, pero tiene {background_image.shape[-1]}.")
pil_background_image = Image.fromarray((background_image * 255).astype(np.uint8), mode='RGBA')
else:
raise ValueError("Se requiere una imagen de fondo válida.")
# Asegurarse de que la imagen de fondo sea más grande que la imagen escalada y rotada
bg_width, bg_height = pil_background_image.size
# Calcular las coordenadas de posición basadas en valores proporcionados (relativo a la imagen de fondo)
offset_x = int((bg_width - rotated_image.size[0]) * position_x)
offset_y = int((bg_height - rotated_image.size[1]) * position_y)
# Componer la imagen escalada y rotada sobre la imagen de fondo
combined_image = pil_background_image.copy()
combined_image.paste(rotated_image, (offset_x, offset_y), rotated_image)
# Convertir la imagen combinada de vuelta a matriz NumPy y normalizar a [0, 1]
combined_image_np = np.array(combined_image).astype(np.float32) / 255.0
# Añadir de nuevo la dimensión de lote y convertir a tensor
combined_image_tensor = torch.tensor(combined_image_np).unsqueeze(0)
return (combined_image_tensor,)