From f7e080f7ee231459e420bdabdf4dd17ec854143c Mon Sep 17 00:00:00 2001 From: rodrigoskohl Date: Mon, 26 Aug 2024 17:23:22 -0300 Subject: [PATCH] update node --- .gitignore | 1 + README.md | 8 + __init__.py | 2 +- nodes/nodes.py | 141 ++++++++++++- requeriments.txt => requirements.txt | 0 worflow/example.json | 293 +++++++++++++++++++++++++++ 6 files changed, 433 insertions(+), 12 deletions(-) create mode 100644 .gitignore create mode 100644 README.md rename requeriments.txt => requirements.txt (100%) create mode 100644 worflow/example.json diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ba0430d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..ce44b07 --- /dev/null +++ b/README.md @@ -0,0 +1,8 @@ +
+ +# Panoramic Image Stitcher for ComfyUI + +
+ +## Simple Node to make panoramic images using [OpenCV](github.com/opencv) stitch function + diff --git a/__init__.py b/__init__.py index d1e7e7c..87cd9ed 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from .nodes import ImageStitchingNode +from .nodes.nodes import ImageStitchingNode NODE_CLASS_MAPPINGS = { diff --git a/nodes/nodes.py b/nodes/nodes.py index 082a20e..7c425af 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -1,28 +1,130 @@ import cv2 import torch +import torch.nn.functional as F import numpy as np from PIL import Image -def pil2tensor(image, device): - return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0).to(device) +def pil2tensor(image, device, rgb=True): + if rgb: + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + cv2pill = Image.fromarray(image) + return torch.from_numpy(np.array(cv2pill).astype(np.float32) / 255.0).unsqueeze(0).to(device) + + + +def image_mask(image, device): + # Converte a imagem para escala de cinza + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + # Define um limiar baixo para detectar a borda preta + _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY) + + # Remove pequenos pontos internos que não fazem parte da borda externa + kernel = np.ones((5, 5), np.uint8) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) + + # Opcional: Use uma máscara de área mínima para garantir que apenas a borda externa seja capturada + # Encontra contornos e filtra os pequenos contornos + cnts, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + mask = np.zeros_like(mask) + for c in cnts: + if cv2.contourArea(c) > 100: # Ajuste o valor conforme o tamanho da borda externa + cv2.drawContours(mask, [c], -1, 255, thickness=cv2.FILLED) + + # Converte a máscara para tensor + mask_tensor = pil2tensor(mask, device, rgb=False) + + return mask_tensor + + + +def apply_mask(image_tensor, mask_tensor): + image_tensor = image_tensor.float() + + # Redimensiona a máscara para combinar com a imagem, se necessário + if image_tensor.shape[2:] != mask_tensor.shape[2:]: + mask_tensor = F.interpolate(mask_tensor, size=image_tensor.shape[2:], mode='bilinear', align_corners=False) + + # Expande a máscara para ter o mesmo número de canais da imagem + if mask_tensor.shape[1] != image_tensor.shape[1]: + mask_tensor = mask_tensor.repeat(1, image_tensor.shape[1], 1, 1) + + # Aplica a máscara + masked_image = image_tensor * mask_tensor + print("Image tensor shape:", image_tensor.shape) + print("Mask tensor shape before resize:", mask_tensor.shape) + print("Mask tensor shape after resize:", mask_tensor.shape) + + return masked_image + + +def remove_black_border(image): + image = cv2.copyMakeBorder(image, 10, 10, 10, 10, cv2.BORDER_CONSTANT, (0, 0, 0)) # Adiciona uma borda preta + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)[1] + kernel = np.ones((5, 5), np.uint8) + thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2) + + cnts,_ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + thresh = np.zeros_like(thresh) + for c in cnts: + if cv2.contourArea(c) > 100: # Ajuste o valor conforme o tamanho da borda externa + cv2.drawContours(thresh, [c], -1, 255, thickness=cv2.FILLED) + + #cnts_h = imutils.grab_contours(cnts) + c = max(cnts, key=cv2.contourArea) + mask = np.zeros(thresh.shape, dtype="uint8") + (x, y, w, h) = cv2.boundingRect(c) + cv2.rectangle(mask, (x, y), (x + w, y + h), 255, -1) + minRect = mask.copy() + sub = mask.copy() + + while cv2.countNonZero(sub) > 0: + minRect = cv2.erode(minRect, None) + sub = cv2.subtract(minRect, thresh) + + cnts,_ = cv2.findContours(minRect.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + #cnts = imutils.grab_contours(cnts) + # Se não houver contornos, retorna a imagem original + if not cnts: + return image + + c = max(cnts, key=cv2.contourArea) + (x, y, w, h) = cv2.boundingRect(c) + image = image[y:y + h, x:x + w] + + return image + class ImageStitchingNode: - CATEGORY = "🧩 Custom Nodes" + @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE", {"list": True}), # Especifica que espera uma lista de imagens - "device": ("STRING", {"default": "cuda:0"}), # Permite escolher a GPU + "device": (["cuda", "cpu"],), # Permite escolher a GPU + "crop": (["enable", "disable"],), # Permite escolher se deseja cortar a imagem + "mode": (["panoramic", "scans"],), # Permite escolher o modo de stitching + "threshold": ("FLOAT",{ + "min": 0.0, + "max": 1.0, + "default": 1.0, + "step": 0.01, + "round": 0.001, + "display": "number", + }), # Permite escolher o limiar para a má + } } - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("Panoramic Image",) + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("IMAGE", "MASK") FUNCTION = "stitch_images" + CATEGORY = "🧩 Custom Nodes" - def stitch_images(self, images, device): + def stitch_images(self, images, device, crop, mode, threshold): # Verifica se recebeu pelo menos duas imagens if len(images) < 2: raise ValueError("At least two images are required for stitching.") @@ -37,15 +139,32 @@ class ImageStitchingNode: np_images = [cv2.cvtColor(img, cv2.COLOR_RGB2BGR) for img in np_images] # Cria o objeto Stitcher e realiza o stitching - stitcher = cv2.Stitcher_create() + if mode == 'panoramic': + stitcher = cv2.Stitcher_create(cv2.Stitcher_PANORAMA) + elif mode == 'scans': + stitcher = cv2.Stitcher_create(cv2.Stitcher_SCANS) + else: + raise ValueError("Invalid mode. Use 'PANORAMA' or 'SCANS'.") + + stitcher.setPanoConfidenceThresh(threshold) (status, pano) = stitcher.stitch(np_images) # Verifica se o stitching foi bem-sucedido if status != cv2.Stitcher_OK: raise RuntimeError(f"Error when stitching: {status}") + # Corta a imagem para remover as bordas pretas usando a técnica de bounding box + if crop == "enable": + pano = remove_black_border(pano) + else: + pano_mask = np.ones(pano.shape[:2], dtype=np.uint8) * 255 # Máscara branca se não for cortar + + pano_mask = image_mask(pano, device) + # Converte a imagem resultante para um tensor que o ComfyUI pode usar - pano_pil = Image.fromarray(cv2.cvtColor(pano, cv2.COLOR_BGR2RGB)) - pano_tensor = pil2tensor(pano_pil, device) + pano_tensor = pil2tensor(pano, device) + + + + return (pano_tensor, pano_mask) - return (pano_tensor,) diff --git a/requeriments.txt b/requirements.txt similarity index 100% rename from requeriments.txt rename to requirements.txt diff --git a/worflow/example.json b/worflow/example.json new file mode 100644 index 0000000..1b80402 --- /dev/null +++ b/worflow/example.json @@ -0,0 +1,293 @@ +{ + "last_node_id": 53, + "last_link_id": 79, + "nodes": [ + { + "id": 43, + "type": "MaskToImage", + "pos": [ + 224, + 192 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 79, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 53 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MaskToImage" + } + }, + { + "id": 39, + "type": "PreviewImage", + "pos": [ + 512, + 288 + ], + "size": { + "0": 224, + "1": 256 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 53 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 30, + "type": "PreviewImage", + "pos": [ + -64, + 160 + ], + "size": { + "0": 224, + "1": 256 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 28 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 34, + "type": "PreviewImage", + "pos": [ + 480, + -224 + ], + "size": { + "0": 576, + "1": 448 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 78 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 13, + "type": "VHS_LoadImagesPath", + "pos": [ + -512, + 96 + ], + "size": [ + 352, + 192 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 28, + 77 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + }, + { + "name": "frame_count", + "type": "INT", + "links": [], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadImagesPath" + }, + "widgets_values": { + "directory": "C:\\ComfyUI_windows_portable\\ComfyUI\\input\\test", + "image_load_cap": 0, + "skip_first_images": 0, + "select_every_nth": 1, + "choose folder to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_images": 0, + "filename": "C:\\ComfyUI_windows_portable\\ComfyUI\\input\\test", + "type": "path", + "format": "folder", + "select_every_nth": 1 + } + } + } + }, + { + "id": 49, + "type": "Image Stitching Node", + "pos": [ + -128, + -32 + ], + "size": { + "0": 315, + "1": 150 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 77, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "Panoramic Image", + "type": "IMAGE", + "links": [ + 78 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "Mask", + "type": "MASK", + "links": [ + 79 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "Image Stitching Node" + }, + "widgets_values": [ + "cuda", + "enable", + "panoramic", + 1 + ] + } + ], + "links": [ + [ + 28, + 13, + 0, + 30, + 0, + "IMAGE" + ], + [ + 53, + 43, + 0, + 39, + 0, + "IMAGE" + ], + [ + 77, + 13, + 0, + 49, + 0, + "IMAGE" + ], + [ + 78, + 49, + 0, + 34, + 0, + "IMAGE" + ], + [ + 79, + 49, + 1, + 43, + 0, + "MASK" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.7627768444385521, + "offset": [ + 532.7374172065774, + 249.47383484481796 + ] + } + }, + "version": 0.4 +} \ No newline at end of file