mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-21 12:42:11 -03:00
76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
class CombineBackgroundOverlay:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"background": ("IMAGE",),
|
|
"overlay": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"horizontal_position": ("FLOAT", {"default": 50, "min": -50, "max": 150, "step": 0.1}),
|
|
"vertical_position": ("FLOAT", {"default": 50, "min": -50, "max": 150, "step": 0.1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "combine_background_overlay"
|
|
CATEGORY = "Bjornulf"
|
|
|
|
def combine_background_overlay(self, background, overlay, mask, horizontal_position, vertical_position):
|
|
# Convert background from torch tensor to numpy array
|
|
bg = background[0].numpy()
|
|
bg = (bg * 255).astype(np.uint8)
|
|
bg_img = Image.fromarray(bg, 'RGB')
|
|
|
|
results = []
|
|
|
|
for ov, m in zip(overlay, mask):
|
|
# Convert overlay from torch tensor to numpy array
|
|
ov = ov.numpy()
|
|
ov = (ov * 255).astype(np.uint8)
|
|
|
|
# Convert mask from torch tensor to numpy array
|
|
m = m.numpy()
|
|
m = (m * 255).astype(np.uint8)
|
|
|
|
# Create PIL Image for overlay
|
|
ov_img = Image.fromarray(ov, 'RGB')
|
|
|
|
# Create alpha channel from mask
|
|
alpha = Image.fromarray(m, 'L')
|
|
|
|
# Combine RGB overlay with alpha mask
|
|
ov_img.putalpha(alpha)
|
|
|
|
# Calculate horizontal position
|
|
x = int((horizontal_position / 100) * (bg_img.width - ov_img.width))
|
|
|
|
# Calculate vertical position
|
|
y = int((vertical_position / 100) * (bg_img.height - ov_img.height))
|
|
|
|
# Create a new image for this overlay
|
|
result = Image.new('RGBA', bg_img.size, (0, 0, 0, 0))
|
|
|
|
# Paste the background
|
|
result.paste(bg_img, (0, 0))
|
|
|
|
# Paste the overlay in the calculated position
|
|
result.paste(ov_img, (x, y), ov_img)
|
|
|
|
# Convert back to numpy array and then to torch tensor
|
|
result_np = np.array(result)
|
|
|
|
# Convert RGBA to RGB
|
|
result_np = result_np[:,:,:3]
|
|
|
|
result_tensor = torch.from_numpy(result_np).float() / 255.0
|
|
|
|
results.append(result_tensor)
|
|
|
|
# Stack all results into a single tensor
|
|
final_result = torch.stack(results)
|
|
|
|
return (final_result,) |