# Core_Nodes/overlay_node.py import math import logging import traceback import numpy as np import torch import torch.nn.functional as F # For affine_grid and grid_sample logger = logging.getLogger(__name__) 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; the output preserves the BASE image channel count (RGB in -> RGB out, RGBA in -> RGBA out). Full batches are processed. """ 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}), }, "optional": { "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 ): if base_image is None or qr_code_image is None: logger.error("OverlayQRCodeNode - Missing base_image or qr_code_image input.") return (torch.zeros((1, 64, 64, 3), dtype=torch.float32),) if base_image.dim() != 4 or qr_code_image.dim() != 4: logger.error("OverlayQRCodeNode - Inputs must be 4D tensors (B, H, W, C).") return (torch.zeros((1, 64, 64, 3), dtype=torch.float32),) batch_size = base_image.shape[0] qr_batch_size = max(1, qr_code_image.shape[0]) outputs = [] for i in range(batch_size): base_item = base_image[i:i + 1] qr_index = min(i, qr_batch_size - 1) # broadcast single QR over batch qr_item = qr_code_image[qr_index:qr_index + 1] try: outputs.append(self._overlay_single( base_item, qr_item, size_percent, opacity, rotation_angle, repeat_pattern, vertical_align, horizontal_align, margin_percent, )) except Exception: logger.error( "OverlayQRCodeNode - Failed on batch item %d, passing base image through.\n%s", i, traceback.format_exc(), ) outputs.append(base_item) return (torch.cat(outputs, dim=0),) def _overlay_single(self, base_image, qr_code_image, size_percent, opacity, rotation_angle, repeat_pattern, vertical_align, horizontal_align, margin_percent): """Overlay one QR image onto one base image. Raises on internal errors.""" 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_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: raise ValueError(f"Base image unsupported channel count: {base_c}") # QR Code 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: raise ValueError(f"QR Code image unsupported channel count: {qr_c}") # --- 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)) # --- Resize QR Code (single tile) --- try: qr_tensor_nchw = torch.from_numpy(qr_img_np_rgba).permute(2, 0, 1).unsqueeze(0).cpu() resized_qr_tensor_nchw = F.interpolate( qr_tensor_nchw, size=(target_qr_h_unrot, target_qr_w_unrot), mode='bicubic', align_corners=False ) except Exception: logger.error("OverlayQRCodeNode - Failed to resize QR Code:\n%s", traceback.format_exc()) raise 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: self._overlay_repeat_pattern( output_img_np, resized_qr_tensor_nchw, base_h, base_w, target_qr_h_unrot, target_qr_w_unrot, rotation_angle, clamped_opacity, ) else: self._overlay_single_placement( output_img_np, base_img_np_rgba, resized_qr_tensor_nchw, base_h, base_w, target_qr_h_unrot, target_qr_w_unrot, rotation_angle, clamped_opacity, vertical_align, horizontal_align, margin_percent, ) except Exception: logger.error("OverlayQRCodeNode - Error during overlay process:\n%s", traceback.format_exc()) raise # --- Preserve base image channel count (RGB in -> RGB out) --- final_np = output_img_np[:, :, :base_c] return torch.from_numpy(final_np).unsqueeze(0).cpu() @staticmethod def _alpha_composite(fg_area, bg_area): """Standard 'over' alpha compositing of two HWC float arrays (RGB+A).""" 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) return np.concatenate((out_rgb, out_alpha), axis=2) def _overlay_repeat_pattern(self, output_img_np, resized_qr_tensor_nchw, base_h, base_w, target_qr_h_unrot, target_qr_w_unrot, rotation_angle, clamped_opacity): """Tile the QR across a rotated canvas, then center-crop onto the base.""" logger.debug("OverlayQRCodeNode - Repeat pattern mode enabled.") angle_rad = math.radians(rotation_angle) cos_a = abs(math.cos(angle_rad)) sin_a = abs(math.sin(angle_rad)) rotated_canvas_w = max(base_w, int(base_w * cos_a + base_h * sin_a)) rotated_canvas_h = max(base_h, int(base_w * sin_a + base_h * cos_a)) tile_h, tile_w = target_qr_h_unrot, target_qr_w_unrot if tile_h <= 0 or tile_w <= 0: logger.warning("OverlayQRCodeNode - Resized QR tile has zero dimension for tiling. Skipping overlay.") return reps_h = max(1, math.ceil(rotated_canvas_h / tile_h)) reps_w = max(1, math.ceil(rotated_canvas_w / tile_w)) resized_qr_np_unrot = resized_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() tiled_layer_unrot = np.tile(resized_qr_np_unrot, (reps_h, reps_w, 1)) tiled_layer_cropped_unrot = tiled_layer_unrot[:rotated_canvas_h, :rotated_canvas_w, :] tiled_tensor_nchw = torch.from_numpy(tiled_layer_cropped_unrot).permute(2, 0, 1).unsqueeze(0).cpu() if abs(rotation_angle % 360) > 1e-3: cos_a = math.cos(angle_rad) 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) target_size = torch.Size([1, 4, rotated_canvas_h, rotated_canvas_w]) grid = F.affine_grid(rotation_matrix, target_size, align_corners=False) rotated_tiled_tensor_nchw = F.grid_sample( tiled_tensor_nchw, grid, mode='bicubic', padding_mode='zeros', align_corners=False ) else: rotated_tiled_tensor_nchw = tiled_tensor_nchw rotated_tiled_np = rotated_tiled_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() # --- Crop the CENTER of the rotated tiled layer to base image size --- center_y = rotated_canvas_h // 2 center_x = rotated_canvas_w // 2 crop_y_start = max(0, center_y - base_h // 2) crop_y_end = min(rotated_canvas_h, center_y + base_h // 2 + (base_h % 2)) crop_x_start = max(0, center_x - base_w // 2) crop_x_end = min(rotated_canvas_w, center_x + base_w // 2 + (base_w % 2)) final_fg_layer_np = rotated_tiled_np[crop_y_start:crop_y_end, crop_x_start:crop_x_end, :] if final_fg_layer_np.shape[:2] != (base_h, base_w): logger.warning( "OverlayQRCodeNode - Final tiled layer crop size (%s) mismatch with base image size (%s,%s). Resizing.", final_fg_layer_np.shape[:2], base_h, base_w, ) 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 ) final_fg_layer_np = resized_fg_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() final_fg_layer_np[:, :, 3] *= clamped_opacity blended_area = self._alpha_composite(final_fg_layer_np, output_img_np) output_img_np[:] = np.clip(blended_area, 0.0, 1.0) def _overlay_single_placement(self, output_img_np, base_img_np_rgba, resized_qr_tensor_nchw, base_h, base_w, target_qr_h_unrot, target_qr_w_unrot, rotation_angle, clamped_opacity, vertical_align, horizontal_align, margin_percent): """Rotate (optional) a single QR tile and place it with alignment/margin.""" logger.debug("OverlayQRCodeNode - Single placement mode enabled.") final_qr_h, final_qr_w = target_qr_h_unrot, target_qr_w_unrot if abs(rotation_angle % 360) > 1e-3: logger.debug("OverlayQRCodeNode - Rotating single QR code by %s degrees.", rotation_angle) angle_rad = math.radians(rotation_angle) cos_a = math.cos(angle_rad) sin_a = math.sin(angle_rad) final_qr_w = max(1, int(abs(target_qr_w_unrot * cos_a) + abs(target_qr_h_unrot * sin_a))) final_qr_h = max(1, int(abs(target_qr_w_unrot * sin_a) + abs(target_qr_h_unrot * cos_a))) 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 = torch.Size([1, 4, final_qr_h, final_qr_w]) grid = F.affine_grid(rotation_matrix, target_size, align_corners=False) 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() else: rotated_qr_np = resized_qr_tensor_nchw.squeeze(0).permute(1, 2, 0).numpy() rotated_qr_np[:, :, 3] *= clamped_opacity 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) 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 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: logger.warning("OverlayQRCodeNode - Single placement results in no valid overlap area.") return bg_area = base_img_np_rgba[base_slice_y_start:base_slice_y_end, base_slice_x_start:base_slice_x_end, :] fg_area = rotated_qr_np[qr_slice_y_start:qr_slice_y_end, qr_slice_x_start:qr_slice_x_end, :] if bg_area.shape[:2] != fg_area.shape[:2]: logger.error("OverlayQRCodeNode - Slice shape mismatch! BG: %s, FG: %s. Skipping blend.", bg_area.shape[:2], fg_area.shape[:2]) return blended_area = np.clip(self._alpha_composite(fg_area, bg_area), 0.0, 1.0) output_img_np[base_slice_y_start:base_slice_y_end, base_slice_x_start:base_slice_x_end, :] = blended_area