diff --git a/README.md b/README.md index 05ef8dd..0e825b3 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ This node is available on [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI- | `conf_thresh` | Decreases the confidence threshold for creating stitch points. Lowering this value makes the algorithm less strict, potentially causing stitching errors or misalignments. See [FAQ](#faq) | Float (0.0–1.0) | 0.0 | 1.0 | 1.0 | 0.01 | | `work_megapix` | Resolution (in megapixels) used for registration step. Higher values = better quality, slower process | Float (0.001–100.0) | 0.001 | 100.0 | 0.6 | 0.01 | | `seam_megapix` | Resolution (in megapixels) used for seam estimation. Lower values speed up the process with some quality loss | Float (0.001–100.0) | 0.001 | 100.0 | 0.1 | 0.01 | + [Default values encountered here](https://github.com/opencv/opencv/blob/master/samples/cpp/stitching_detailed.cpp) ## FAQ @@ -41,7 +42,7 @@ Controls the confidence threshold for creating stitch points. Lowering this valu Important: Always use the default value 1.0 unless your images are high quality, well aligned, and captured carefully; Images are provided in the correct sequence, ensuring smooth visual continuity. -Lowering the threshold offers no real advantage except to avoid false positives when there are very similar or duplicated images in the batch. Use caution, as lowering it improperly can degrade stitching quality. +Lowering the threshold offers no real advantage except to avoid false positives images. ## Next features diff --git a/nodes/nodes.py b/nodes/nodes.py index d031b07..3267258 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -5,16 +5,24 @@ import numpy as np from PIL import Image from comfy.model_management import get_torch_device -def pil2tensor(image, device, rgb=True): - if rgb: - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - else: - image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - _, image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY) - #kernel = np.ones((5, 5), np.uint8) - #image = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) - cv2pill = Image.fromarray(image) - return torch.from_numpy(np.array(cv2pill).astype(np.float32) / 255.0).unsqueeze(0).to(device) +def cv2_to_tensor(image, device, return_mask=False): + # Converte imagem BGR para RGB + rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + # Converte para float e normaliza + rgb_image = rgb_image.astype(np.float32) / 255.0 + # Cria tensor da imagem: [1, H, W, C] + image_tensor = torch.from_numpy(rgb_image).unsqueeze(0).to(device) + + if return_mask: + # Cria máscara em escala de cinza e aplica threshold + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + _, binary_mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY) + binary_mask = binary_mask.astype(np.float32) / 255.0 + binary_mask = np.expand_dims(binary_mask, axis=-1) # [H, W, 1] + mask_tensor = torch.from_numpy(binary_mask).unsqueeze(0).to(device) # [1, H, W, 1] + return image_tensor, mask_tensor + + return image_tensor def remove_black_border(image): @@ -23,14 +31,16 @@ def remove_black_border(image): thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)[1] 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) > 1: # Ajuste o valor conforme o tamanho da borda externa - cv2.drawContours(thresh, [c], -1, 255, thickness=cv2.FILLED) + if cnts: + # Se houver contornos, encontra o maior contorno e desenha ele(isso exclui pixeis pretos que não são parte do objeto) + c = max(cnts, key=cv2.contourArea) + (x, y, w, h) = cv2.boundingRect(c) + thresh = np.zeros_like(thresh) + cv2.drawContours(thresh, [c], -1, 255, thickness=cv2.FILLED) + else: + return image - 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() @@ -59,9 +69,9 @@ class ImageStitchingNode: def INPUT_TYPES(cls): return { "required": { - "images": ("IMAGE", {"list": True}), # Especifica que espera uma lista de imagens - "crop": (["enable", "disable"],), # Permite escolher se deseja cortar a imagem - "mode": (["panoramic", "scans"],), # Permite escolher o modo de stitching + "images": ("IMAGE", ), + "crop": (["enable", "disable"],), + "mode": (["panoramic", "scans"],), "conf_thresh": ("FLOAT",{ "min": 0.0, "max": 1.0, @@ -91,18 +101,17 @@ class ImageStitchingNode: } } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("IMAGE", "MASK") + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) FUNCTION = "stitch_images" CATEGORY = "🧩 Custom Nodes" + def stitch_images(self, images, crop, mode, conf_thresh, work_megapix, seam_megapix): # Verifica se recebeu pelo menos duas imagens if len(images) < 2: raise ValueError("At least two images are required for stitching.") - # Verifica se o dispositivo especificado está disponível - # Converter os tensores para arrays de numpy compatíveis com OpenCV np_images = [np.array(image.squeeze(0).cpu().numpy() * 255, dtype=np.uint8) for image in images] @@ -115,7 +124,7 @@ class ImageStitchingNode: elif mode == 'scans': stitcher = cv2.Stitcher_create(cv2.Stitcher_SCANS) else: - raise ValueError("Invalid mode. Use 'PANORAMA' or 'SCANS'.") + raise ValueError("Invalid mode. Use 'panoramic' or 'scans'.") stitcher.setPanoConfidenceThresh(conf_thresh) stitcher.setRegistrationResol(work_megapix) @@ -126,18 +135,19 @@ class ImageStitchingNode: if status != cv2.Stitcher_OK: raise RuntimeError(f"Error when stitching: {status}") - # Retorna a mascara original da imagem panorâmica mesmo se crop estiver habilitado - pano_mask = pil2tensor(pano, device=self.device, rgb=False) - # Corta a imagem para remover as bordas pretas usando a técnica de bounding box + + # Aplica crop na imagem se solicitado if crop == "enable": pano = remove_black_border(pano) + pano_tensor = cv2_to_tensor(pano, device=self.device) + # Se crop não for solicitado, converte a imagem e aplica máscara + else: + pano_tensor, pano_mask = cv2_to_tensor(pano, device=self.device, return_mask=True) + pano_mask = pano_mask.clamp(0, 1) + pano_tensor = torch.cat([pano_tensor, pano_mask], dim=-1) - - - # Converte a imagem resultante para um tensor que o ComfyUI pode usar - pano_tensor = pil2tensor(pano, device=self.device, rgb=True) + return (pano_tensor,) - return (pano_tensor, pano_mask) diff --git a/pyproject.toml b/pyproject.toml index d047252..b1aec4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-panoramic-imgstitcher" description = "Simple Node to make panoramic images" -version = "1.0.2" +version = "1.0.3" license = {file = "LICENSE"} [project.urls]