From a35f4c5f11de4f784c228d6dfd0eedbede08b392 Mon Sep 17 00:00:00 2001 From: Alan Burns Date: Tue, 5 May 2026 17:23:16 -0500 Subject: [PATCH] feat: add AutoStraighten node for scanned photos --- __init__.py | 9 +++ auto_straighten.py | 164 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+) create mode 100644 auto_straighten.py diff --git a/__init__.py b/__init__.py index 88f6851..d497763 100644 --- a/__init__.py +++ b/__init__.py @@ -74,6 +74,11 @@ from .memory_reporter import ( NODE_CLASS_MAPPINGS as MEMORY_REPORTER_CLASS_MAP, NODE_DISPLAY_NAME_MAPPINGS as MEMORY_REPORTER_DISPLAY_MAP, ) +from .auto_straighten import ( + NODE_CLASS_MAPPINGS as AUTO_STRAIGHTEN_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as AUTO_STRAIGHTEN_DISPLAY_MAP, +) + NODE_CLASS_MAPPINGS = {} @@ -97,6 +102,8 @@ NODE_CLASS_MAPPINGS.update(PAIRED_LOADER_CLASS_MAP) NODE_CLASS_MAPPINGS.update(COMPARISON_GATE_CLASS_MAP) NODE_CLASS_MAPPINGS.update(SIMPLE_IMAGE_SAVER_CLASS_MAP) NODE_CLASS_MAPPINGS.update(MEMORY_REPORTER_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(AUTO_STRAIGHTEN_CLASS_MAP) + NODE_DISPLAY_NAME_MAPPINGS.update(AUTO_ADJUST_DISPLAY_MAP) NODE_DISPLAY_NAME_MAPPINGS.update(OUTPAINT_DISPLAY_MAP) @@ -115,6 +122,8 @@ NODE_DISPLAY_NAME_MAPPINGS.update(PAIRED_LOADER_DISPLAY_MAP) NODE_DISPLAY_NAME_MAPPINGS.update(COMPARISON_GATE_DISPLAY_MAP) NODE_DISPLAY_NAME_MAPPINGS.update(SIMPLE_IMAGE_SAVER_DISPLAY_MAP) NODE_DISPLAY_NAME_MAPPINGS.update(MEMORY_REPORTER_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(AUTO_STRAIGHTEN_DISPLAY_MAP) + # Register the load-image node which does not expose pre-built mappings. NODE_CLASS_MAPPINGS["LoadImageCombined"] = LoadImageCombined diff --git a/auto_straighten.py b/auto_straighten.py new file mode 100644 index 0000000..613e1cb --- /dev/null +++ b/auto_straighten.py @@ -0,0 +1,164 @@ +import torch +import numpy as np +import cv2 + +class ProcessScannedPhoto: + """Automatically straightens and crops scanned photos in a single pass.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "straighten": ("BOOLEAN", {"default": True}), + "crop_mode": (["Inner Photo Frame", "Scanner Bed Only", "None"], {"default": "Inner Photo Frame"}), + "padding": ("INT", {"default": 0, "min": -500, "max": 500, "step": 1}), + "threshold": ("FLOAT", {"default": 0.80, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "process" + CATEGORY = "PortraitUtils/Transform" + + def process(self, image, straighten, crop_mode, padding, threshold): + out_images = [] + for img_tensor in image: + # Convert to numpy + img_np = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) + + if img_np.shape[-1] >= 3: + gray = cv2.cvtColor(img_np[..., :3], cv2.COLOR_RGB2GRAY) + else: + gray = img_np.squeeze() + + thresh_val = int(threshold * 255) + + # --- STEP 1: STRAIGHTEN --- + if straighten: + # Use Hybrid Edge Detection to find outer bounds + gray_blur = cv2.GaussianBlur(gray, (5, 5), 0) + edges = cv2.Canny(gray_blur, 30, 100) + edges = cv2.dilate(edges, np.ones((3,3), np.uint8), iterations=1) + + _, thresh_img = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY_INV) + combined_mask = cv2.bitwise_or(edges, thresh_img) + + contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + valid_points = [c for c in contours if cv2.contourArea(c) > 100 or len(c) > 50] + + if valid_points: + all_points = np.vstack(valid_points) + hull = cv2.convexHull(all_points) + rect = cv2.minAreaRect(hull) + center, size, angle = rect + box = cv2.boxPoints(rect) + + edges_info = [] + for i in range(4): + p1 = box[i] + p2 = box[(i+1)%4] + dx = p2[0] - p1[0] + dy = p2[1] - p1[1] + length = np.hypot(dx, dy) + angle_deg = np.degrees(np.arctan2(dy, dx)) + edges_info.append((length, angle_deg)) + + longest_edge = max(edges_info, key=lambda x: x[0]) + longest_angle = longest_edge[1] + + closest_90 = round(longest_angle / 90.0) * 90.0 + angle_to_rotate = longest_angle - closest_90 + + h, w = img_np.shape[:2] + M = cv2.getRotationMatrix2D(center, angle_to_rotate, 1.0) + + abs_cos = abs(M[0, 0]) + abs_sin = abs(M[0, 1]) + bound_w = int(h * abs_sin + w * abs_cos) + bound_h = int(h * abs_cos + w * abs_sin) + M[0, 2] += bound_w / 2 - center[0] + M[1, 2] += bound_h / 2 - center[1] + + border_val = (255, 255, 255, 255) if img_np.shape[-1] == 4 else ((255, 255, 255) if img_np.shape[-1] == 3 else 255) + img_np = cv2.warpAffine(img_np, M, (bound_w, bound_h), borderMode=cv2.BORDER_CONSTANT, borderValue=border_val) + + # Update gray image for crop step + if img_np.shape[-1] >= 3: + gray = cv2.cvtColor(img_np[..., :3], cv2.COLOR_RGB2GRAY) + else: + gray = img_np.squeeze() + + # --- STEP 2: CROP --- + if crop_mode != "None": + h_img, w_img = img_np.shape[:2] + + if crop_mode == "Scanner Bed Only": + # Use Hybrid logic on rotated image + gray_blur = cv2.GaussianBlur(gray, (5, 5), 0) + edges = cv2.Canny(gray_blur, 30, 100) + edges = cv2.dilate(edges, np.ones((3,3), np.uint8), iterations=1) + + _, thresh_img = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY_INV) + combined = cv2.bitwise_or(edges, thresh_img) + contours, _ = cv2.findContours(combined, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + valid_points = [c for c in contours if cv2.contourArea(c) > 100 or len(c) > 50] + if valid_points: + x, y, cw, ch = cv2.boundingRect(np.vstack(valid_points)) + + x_min = max(0, x - padding) + y_min = max(0, y - padding) + x_max = min(w_img, x + cw + padding) + y_max = min(h_img, y + ch + padding) + img_np = img_np[y_min:y_max, x_min:x_max] + + elif crop_mode == "Inner Photo Frame": + # Use Morphological Contour logic to bypass borders + _, thresh_img = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY_INV) + kernel = np.ones((15, 15), np.uint8) + closed = cv2.morphologyEx(thresh_img, cv2.MORPH_CLOSE, kernel) + contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if contours: + largest_contour = max(contours, key=cv2.contourArea) + x, y, cw, ch = cv2.boundingRect(largest_contour) + + x_min = max(0, x - padding) + y_min = max(0, y - padding) + x_max = min(w_img, x + cw + padding) + y_max = min(h_img, y + ch + padding) + img_np = img_np[y_min:y_max, x_min:x_max] + + # Convert back to tensor + out_tensor = torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0) + out_images.append(out_tensor) + + # Batching logic + if len(out_images) == 1: + return (out_images[0],) + + max_h = max(img.shape[1] for img in out_images) + max_w = max(img.shape[2] for img in out_images) + padded_images = [] + for img in out_images: + h_i, w_i = img.shape[1:3] + pad_h = max_h - h_i + pad_w = max_w - w_i + if pad_h > 0 or pad_w > 0: + c = img.shape[-1] + padded = torch.ones((1, max_h, max_w, c), dtype=img.dtype, device=img.device) + padded[0, :h_i, :w_i, :] = img[0] + padded_images.append(padded) + else: + padded_images.append(img) + + return (torch.cat(padded_images, dim=0),) + +NODE_CLASS_MAPPINGS = { + "ProcessScannedPhoto": ProcessScannedPhoto +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ProcessScannedPhoto": "Process Scanned Photo" +}