update node
This commit is contained in:
rodrigoskohl
@@ -0,0 +1 @@
|
||||
__pycache__/
|
||||
@@ -0,0 +1,8 @@
|
||||
<div align="center">
|
||||
|
||||
# Panoramic Image Stitcher for ComfyUI
|
||||
|
||||
</div>
|
||||
|
||||
## Simple Node to make panoramic images using [OpenCV](github.com/opencv) stitch function
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
from .nodes import ImageStitchingNode
|
||||
from .nodes.nodes import ImageStitchingNode
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
+130
-11
@@ -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,)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user