Merge pull request #5 from RodrigoSKohl:fix/mask-pano

fix: enhance image mask processing functions

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
 

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)
 

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]