380 lines
20 KiB
Python
380 lines
20 KiB
Python
# Core_Nodes/overlay_node.py
|
|
import torch
|
|
import torch.nn.functional as F # For affine_grid and grid_sample
|
|
import numpy as np
|
|
import math # Needed for ceil, sin, cos, radians
|
|
import traceback # For detailed error printing
|
|
import comfy.model_management as mm
|
|
|
|
class OverlayQRCodeNode:
|
|
"""
|
|
Overlays a QR Code image onto a base image.
|
|
Supports single placement with alignment/margin/rotation options OR
|
|
repeating the QR code as a tiled pattern (rotation applied to the overall pattern).
|
|
Handles RGBA inputs and outputs an RGBA image.
|
|
"""
|
|
|
|
ALIGN_VERT = ["top", "bottom", "center"]
|
|
ALIGN_HORIZ = ["left", "right", "center"]
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_image": ("IMAGE",),
|
|
"qr_code_image": ("IMAGE",),
|
|
"size_percent": ("FLOAT", {"default": 20.0, "min": 1.0, "max": 100.0, "step": 0.5, "display": "slider"}),
|
|
"opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),
|
|
"rotation_angle": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.1, "display": "number"}),
|
|
"repeat_pattern": ("BOOLEAN", {"default": False}), # Renamed
|
|
},
|
|
"optional": {
|
|
# These are primarily used if repeat_pattern is False, but margin could maybe apply to pattern?
|
|
# Let's keep them optional and clarify usage. Vertical/Horizontal align definitely only for single.
|
|
"vertical_align": (cls.ALIGN_VERT, {"default": "bottom"}),
|
|
"horizontal_align": (cls.ALIGN_HORIZ, {"default": "right"}),
|
|
"margin_percent": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 50.0, "step": 0.1, "display": "slider"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "overlay_qr"
|
|
CATEGORY = "image/postprocessing"
|
|
|
|
def overlay_qr(self, base_image, qr_code_image,
|
|
size_percent, opacity, rotation_angle, repeat_pattern,
|
|
vertical_align="bottom", horizontal_align="right", margin_percent=2.0
|
|
):
|
|
mm.soft_empty_cache()
|
|
dummy_tensor = torch.zeros((1, 64, 64, 4), dtype=torch.float32) # RGBA dummy
|
|
|
|
if base_image is None or qr_code_image is None:
|
|
print("ERROR: OverlayQRCodeNode - Missing base_image or qr_code_image input.")
|
|
return (dummy_tensor,)
|
|
if base_image.dim() != 4 or qr_code_image.dim() != 4:
|
|
print("ERROR: OverlayQRCodeNode - Inputs must be 4D tensors (B, H, W, C).")
|
|
return (dummy_tensor,)
|
|
|
|
base_img_tensor = base_image[0]
|
|
qr_img_tensor = qr_code_image[0]
|
|
base_h, base_w, base_c = base_img_tensor.shape
|
|
qr_h, qr_w, qr_c = qr_img_tensor.shape
|
|
|
|
# --- Ensure RGBA format internally (float32, 0-1 range) ---
|
|
# Base Image
|
|
base_img_np = base_img_tensor.cpu().numpy()
|
|
if base_c == 3:
|
|
base_img_np_rgba = np.concatenate((base_img_np, np.ones((base_h, base_w, 1), dtype=np.float32)), axis=2)
|
|
elif base_c == 4:
|
|
base_img_np_rgba = base_img_np
|
|
else:
|
|
print(f"ERROR: OverlayQRCodeNode - Base image unsupported channel count: {base_c}")
|
|
return (dummy_tensor,)
|
|
|
|
# QR Code Image
|
|
qr_img_np = qr_img_tensor.cpu().numpy()
|
|
if qr_c == 3:
|
|
qr_img_np_rgba = np.concatenate((qr_img_np, np.ones((qr_h, qr_w, 1), dtype=np.float32)), axis=2)
|
|
elif qr_c == 4:
|
|
qr_img_np_rgba = qr_img_np
|
|
else:
|
|
print(f"ERROR: OverlayQRCodeNode - QR Code image unsupported channel count: {qr_c}")
|
|
return (dummy_tensor,)
|
|
|
|
# --- Calculate Target QR Tile Size (Always based on UNROTATED size_percent) ---
|
|
base_ref_dim = min(base_h, base_w)
|
|
target_qr_dim_unrot = max(1, int(base_ref_dim * (size_percent / 100.0)))
|
|
qr_aspect_ratio = qr_w / qr_h if qr_h > 0 else 1
|
|
if qr_w >= qr_h:
|
|
target_qr_w_unrot = target_qr_dim_unrot
|
|
target_qr_h_unrot = max(1, int(target_qr_w_unrot / qr_aspect_ratio))
|
|
else:
|
|
target_qr_h_unrot = target_qr_dim_unrot
|
|
target_qr_w_unrot = max(1, int(target_qr_h_unrot * qr_aspect_ratio))
|
|
|
|
if target_qr_h_unrot <= 0 or target_qr_w_unrot <= 0:
|
|
print(f"ERROR: OverlayQRCodeNode - Invalid target QR size ({target_qr_h_unrot}x{target_qr_w_unrot}).")
|
|
return (torch.from_numpy(base_img_np_rgba).unsqueeze(0).cpu(),)
|
|
|
|
# --- Resize QR Code (single tile) ---
|
|
try:
|
|
qr_tensor_nchw = torch.from_numpy(qr_img_np_rgba).permute(2, 0, 1).unsqueeze(0).cpu()
|
|
# Resize to the UNROTATED tile size
|
|
resized_qr_tensor_nchw = F.interpolate(
|
|
qr_tensor_nchw, size=(target_qr_h_unrot, target_qr_w_unrot), mode='bicubic', align_corners=False
|
|
)
|
|
# Keep resized tensor as NCHW float for further processing
|
|
except Exception as e:
|
|
print(f"ERROR: OverlayQRCodeNode - Failed to resize QR Code: {e}\n{traceback.format_exc()}")
|
|
return (torch.from_numpy(base_img_np_rgba).unsqueeze(0).cpu(),)
|
|
|
|
# --- Calculate clamped_opacity ONCE ---
|
|
clamped_opacity = max(0.0, min(1.0, opacity))
|
|
|
|
|
|
# --- Prepare Output Image ---
|
|
output_img_np = base_img_np_rgba.copy()
|
|
|
|
# --- Perform Overlay ---
|
|
try:
|
|
if repeat_pattern:
|
|
# --- Repeat Pattern Mode ---
|
|
print("INFO: OverlayQRCodeNode - Repeat pattern mode enabled.")
|
|
|
|
# Calculate size of canvas needed to rotate base image (this is the target size for the tiled pattern)
|
|
angle_rad = math.radians(rotation_angle)
|
|
cos_a = abs(math.cos(angle_rad))
|
|
sin_a = abs(math.sin(angle_rad))
|
|
# Size of the base image's bounding box after rotation
|
|
rotated_canvas_w = int(base_w * cos_a + base_h * sin_a)
|
|
rotated_canvas_h = int(base_w * sin_a + base_h * cos_a)
|
|
|
|
# Ensure rotated canvas is at least base size (can happen near 0/90/180/270)
|
|
rotated_canvas_w = max(base_w, rotated_canvas_w)
|
|
rotated_canvas_h = max(base_h, rotated_canvas_h)
|
|
|
|
# Calculate how many UNROTATED tiles fit into the ROTATED canvas size
|
|
tile_h, tile_w = target_qr_h_unrot, target_qr_w_unrot
|
|
if tile_h <= 0 or tile_w <= 0: # Safety check
|
|
print("Warning: OverlayQRCodeNode - Resized QR tile has zero dimension for tiling. Skipping overlay.")
|
|
else:
|
|
# Calculate repetitions needed for the rotated canvas dimensions
|
|
reps_h = math.ceil(rotated_canvas_h / tile_h)
|
|
reps_w = math.ceil(rotated_canvas_w / tile_w)
|
|
reps_h = max(1, reps_h)
|
|
reps_w = max(1, reps_w)
|
|
|
|
# Tile the UNROTATED resized QR code image
|
|
resized_qr_np_unrot = resized_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() # H, W, 4 float32
|
|
tiled_layer_unrot = np.tile(resized_qr_np_unrot, (reps_h, reps_w, 1)) # Create a large unrotated tiled image
|
|
|
|
# Crop the unrotated tiled image to the size of the rotated canvas
|
|
tiled_layer_cropped_unrot = tiled_layer_unrot[:rotated_canvas_h, :rotated_canvas_w, :]
|
|
|
|
# --- Rotate the Tiled Layer ---
|
|
tiled_tensor_nchw = torch.from_numpy(tiled_layer_cropped_unrot).permute(2, 0, 1).unsqueeze(0).cpu()
|
|
|
|
if abs(rotation_angle % 360) > 1e-3: # Rotate only if angle is significant
|
|
# Create affine matrix for rotation (inverse rotation for grid_sample)
|
|
cos_a = math.cos(angle_rad) # Use original sin/cos for matrix
|
|
sin_a = math.sin(angle_rad)
|
|
rotation_matrix = torch.tensor([
|
|
[cos_a, sin_a, 0],
|
|
[-sin_a, cos_a, 0]
|
|
], dtype=torch.float32, device=tiled_tensor_nchw.device).unsqueeze(0) # Add batch dim
|
|
|
|
# Target size is the rotated canvas size
|
|
target_size = torch.Size([1, 4, rotated_canvas_h, rotated_canvas_w])
|
|
grid = F.affine_grid(rotation_matrix, target_size, align_corners=False)
|
|
|
|
# Perform grid sampling
|
|
rotated_tiled_tensor_nchw = F.grid_sample(
|
|
tiled_tensor_nchw, grid, mode='bicubic', padding_mode='zeros', align_corners=False
|
|
)
|
|
else:
|
|
# No rotation needed for the tiled layer
|
|
rotated_tiled_tensor_nchw = tiled_tensor_nchw
|
|
|
|
# Convert rotated tiled tensor back to HWC numpy float32
|
|
rotated_tiled_np = rotated_tiled_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() # H, W, 4 float32
|
|
|
|
# --- Crop the CENTER of the rotated tiled layer to base image size ---
|
|
center_y = rotated_canvas_h // 2
|
|
center_x = rotated_canvas_w // 2
|
|
|
|
base_h_half = base_h // 2
|
|
base_w_half = base_w // 2
|
|
|
|
crop_y_start = center_y - base_h_half
|
|
crop_y_end = center_y + base_h_half + (base_h % 2) # Handle odd dimensions
|
|
crop_x_start = center_x - base_w_half
|
|
crop_x_end = center_x + base_w_half + (base_w % 2) # Handle odd dimensions
|
|
|
|
# Ensure crop is within bounds (should be if canvas was calculated correctly)
|
|
crop_y_start = max(0, crop_y_start)
|
|
crop_y_end = min(rotated_canvas_h, crop_y_end)
|
|
crop_x_start = max(0, crop_x_start)
|
|
crop_x_end = min(rotated_canvas_w, crop_x_end)
|
|
|
|
# The final foreground layer is this center crop
|
|
final_fg_layer_np = rotated_tiled_np[crop_y_start:crop_y_end, crop_x_start:crop_x_end, :]
|
|
|
|
# Ensure cropped size matches base image size (can be slightly off due to rounding/interpolation)
|
|
if final_fg_layer_np.shape[:2] != (base_h, base_w):
|
|
print(f"Warning: OverlayQRCodeNode - Final tiled layer crop size ({final_fg_layer_np.shape[:2]}) mismatch with base image size ({base_h},{base_w}). Resizing.")
|
|
fg_tensor_nchw = torch.from_numpy(final_fg_layer_np).permute(2, 0, 1).unsqueeze(0).cpu()
|
|
resized_fg_tensor_nchw = F.interpolate(
|
|
fg_tensor_nchw, size=(base_h, base_w), mode='bilinear', align_corners=False # Bilinear for speed/simplicity here
|
|
)
|
|
final_fg_layer_np = resized_fg_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy()
|
|
|
|
|
|
# --- Apply Opacity to the FINAL tiled layer alpha ---
|
|
# Apply the pre-calculated clamped_opacity
|
|
final_fg_layer_np[:, :, 3] *= clamped_opacity
|
|
|
|
|
|
# --- Blend Full Layers (Tiled FG over Base BG) ---
|
|
bg_area = base_img_np_rgba # The whole base image
|
|
fg_area = final_fg_layer_np # The whole final tiled layer
|
|
|
|
fg_rgb, fg_alpha = fg_area[:, :, :3], fg_area[:, :, 3:4]
|
|
bg_rgb, bg_alpha = bg_area[:, :, :3], bg_area[:, :, 3:4]
|
|
|
|
# Alpha Compositing (Fg over Bg)
|
|
out_alpha = fg_alpha + bg_alpha * (1.0 - fg_alpha)
|
|
numerator = (fg_rgb * fg_alpha) + (bg_rgb * bg_alpha * (1.0 - fg_alpha))
|
|
epsilon = 1e-7
|
|
safe_alpha_mask = out_alpha > epsilon # Shape (H, W, 1)
|
|
divided_rgb = numerator / (out_alpha + epsilon)
|
|
out_rgb = np.where(safe_alpha_mask, divided_rgb, bg_rgb)
|
|
|
|
blended_area = np.concatenate((out_rgb, out_alpha), axis=2)
|
|
output_img_np = np.clip(blended_area, 0.0, 1.0) # Assign blended result to output
|
|
|
|
else:
|
|
# --- Single Placement Mode ---
|
|
print("INFO: OverlayQRCodeNode - Single placement mode enabled.")
|
|
|
|
# In single mode, rotation is applied to the *individual* tile
|
|
final_qr_h, final_qr_w = target_qr_h_unrot, target_qr_w_unrot
|
|
rotated_qr_np = None # Initialize
|
|
|
|
if abs(rotation_angle % 360) > 1e-3: # Rotate only if angle is significant
|
|
print(f"INFO: OverlayQRCodeNode - Rotating single QR code by {rotation_angle} degrees.")
|
|
try:
|
|
angle_rad = math.radians(rotation_angle)
|
|
cos_a = math.cos(angle_rad)
|
|
sin_a = math.sin(angle_rad)
|
|
|
|
# Calculate new bounding box size for the single rotated tile
|
|
final_qr_w = int(abs(target_qr_w_unrot * cos_a) + abs(target_qr_h_unrot * sin_a))
|
|
final_qr_h = int(abs(target_qr_w_unrot * sin_a) + abs(target_qr_h_unrot * cos_a))
|
|
final_qr_w = max(1, final_qr_w)
|
|
final_qr_h = max(1, final_qr_h)
|
|
|
|
# Create affine matrix for inverse rotation
|
|
rotation_matrix = torch.tensor([
|
|
[cos_a, sin_a, 0],
|
|
[-sin_a, cos_a, 0]
|
|
], dtype=torch.float32, device=resized_qr_tensor_nchw.device).unsqueeze(0)
|
|
|
|
# Target size is the new rotated bounding box size
|
|
target_size = torch.Size([1, 4, final_qr_h, final_qr_w])
|
|
grid = F.affine_grid(rotation_matrix, target_size, align_corners=False)
|
|
|
|
# Perform grid sampling
|
|
rotated_qr_tensor_nchw = F.grid_sample(
|
|
resized_qr_tensor_nchw, grid, mode='bicubic', padding_mode='zeros', align_corners=False
|
|
)
|
|
rotated_qr_np = rotated_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy()
|
|
|
|
except Exception as e:
|
|
print(f"ERROR: OverlayQRCodeNode - Failed during single tile rotation: {e}\n{traceback.format_exc()}")
|
|
# Fallback: use unrotated resized image
|
|
rotated_qr_np = resized_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy()
|
|
final_qr_h = target_qr_h_unrot
|
|
final_qr_w = target_qr_w_unrot
|
|
else:
|
|
# No rotation needed, just convert the resized tensor to numpy
|
|
rotated_qr_np = resized_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy()
|
|
# final_qr_h and final_qr_w remain the unrotated dimensions
|
|
|
|
|
|
# Apply Opacity to the SINGLE rotated tile alpha
|
|
# Apply the pre-calculated clamped_opacity
|
|
rotated_qr_np[:, :, 3] *= clamped_opacity
|
|
|
|
|
|
# Calculate Position and Slices using FINAL (potentially rotated) dimensions
|
|
target_qr_h_final, target_qr_w_final = final_qr_h, final_qr_w
|
|
|
|
margin_px_h = int(base_h * (margin_percent / 100.0))
|
|
margin_px_w = int(base_w * (margin_percent / 100.0))
|
|
|
|
if vertical_align == "top": y_start = margin_px_h
|
|
elif vertical_align == "bottom": y_start = base_h - target_qr_h_final - margin_px_h
|
|
else: y_start = (base_h - target_qr_h_final) // 2
|
|
|
|
if horizontal_align == "left": x_start = margin_px_w
|
|
elif horizontal_align == "right": x_start = base_w - target_qr_w_final - margin_px_w
|
|
else: x_start = (base_w - target_qr_w_final) // 2
|
|
|
|
y_end = y_start + target_qr_h_final
|
|
x_end = x_start + target_qr_w_final
|
|
|
|
base_slice_y_start = max(0, y_start)
|
|
base_slice_y_end = min(base_h, y_end)
|
|
base_slice_x_start = max(0, x_start)
|
|
base_slice_x_end = min(base_w, x_end)
|
|
|
|
# Recalculate QR slices relative to the part that overlaps the base image
|
|
qr_slice_y_start = base_slice_y_start - y_start
|
|
qr_slice_x_start = base_slice_x_start - x_start
|
|
qr_slice_y_end = base_slice_y_end - y_start
|
|
qr_slice_x_end = base_slice_x_end - x_start
|
|
|
|
# Check for valid overlap dimensions by comparing slice sizes
|
|
overlap_h = base_slice_y_end - base_slice_y_start
|
|
overlap_w = base_slice_x_end - base_slice_x_start
|
|
qr_overlap_h = qr_slice_y_end - qr_slice_y_start
|
|
qr_overlap_w = qr_slice_x_end - qr_slice_x_start
|
|
|
|
|
|
if overlap_h <= 0 or overlap_w <= 0 or qr_overlap_h <= 0 or qr_overlap_w <= 0 or \
|
|
overlap_h != qr_overlap_h or overlap_w != qr_overlap_w:
|
|
print("Warning: OverlayQRCodeNode - Single placement results in no valid overlap area.")
|
|
else:
|
|
# Extract the overlapping background area
|
|
bg_area = base_img_np_rgba[base_slice_y_start:base_slice_y_end, base_slice_x_start:base_slice_x_end, :]
|
|
|
|
# Extract the corresponding foreground QR area (single rotated tile)
|
|
# Use the calculated QR slices
|
|
fg_area = rotated_qr_np[qr_slice_y_start:qr_slice_y_end, qr_slice_x_start:qr_slice_x_end, :]
|
|
|
|
# Ensure shapes match after slicing (should be guaranteed now)
|
|
if bg_area.shape[:2] != fg_area.shape[:2]:
|
|
# Fallback/debug check - shouldn't be needed with fixed slicing
|
|
print(f"ERROR: OverlayQRCodeNode - Slice shape mismatch! BG: {bg_area.shape[:2]}, FG: {fg_area.shape[:2]}. Skipping blend.")
|
|
else:
|
|
# --- Blend Single Placement Area ---
|
|
fg_rgb, fg_alpha = fg_area[:, :, :3], fg_area[:, :, 3:4]
|
|
bg_rgb, bg_alpha = bg_area[:, :, :3], bg_area[:, :, 3:4]
|
|
out_alpha = fg_alpha + bg_alpha * (1.0 - fg_alpha)
|
|
numerator = (fg_rgb * fg_alpha) + (bg_rgb * bg_alpha * (1.0 - fg_alpha))
|
|
epsilon = 1e-7
|
|
safe_alpha_mask = out_alpha > epsilon
|
|
divided_rgb = numerator / (out_alpha + epsilon)
|
|
out_rgb = np.where(safe_alpha_mask, divided_rgb, bg_rgb)
|
|
blended_area = np.concatenate((out_rgb, out_alpha), axis=2)
|
|
blended_area = np.clip(blended_area, 0.0, 1.0)
|
|
|
|
# Place the blended area back into the output image copy
|
|
output_img_np[base_slice_y_start:base_slice_y_end, base_slice_x_start:base_slice_x_end, :] = blended_area
|
|
|
|
|
|
except Exception as e:
|
|
print(f"ERROR: OverlayQRCodeNode - Error during overlay process: {e}\n{traceback.format_exc()}")
|
|
# Return original base image copy on error
|
|
|
|
# --- Convert Final Image to Torch Tensor ---
|
|
try:
|
|
output_tensor = torch.from_numpy(output_img_np).unsqueeze(0).cpu()
|
|
except Exception as e:
|
|
print(f"ERROR: OverlayQRCodeNode - Failed to convert final numpy array to tensor: {e}")
|
|
return (dummy_tensor,)
|
|
|
|
return (output_tensor,)
|
|
|
|
|
|
# --- ComfyUI Mappings ---
|
|
NODE_CLASS_MAPPINGS = {
|
|
"⧉ IsulionOverlay": OverlayQRCodeNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"⧉ IsulionOverlay": "Overlay QR Code on Image"
|
|
} |