fix

.gitignore

@@ -31,3 +31,4 @@ logs/
 # Cache
 .cache/
 .temp/
+nodes/__pycache__/nodes.cpython-311.pyc

README.md

@@ -1,5 +1,7 @@
 # ComfyUI Mask Contour Processor
 
+_Version 0.1_
+
 A ComfyUI node that improves inpainting results by extending mask boundaries with geometric patterns, helping create smoother transitions and better context for AI-driven image completion.
 
 ## Installation

nodes/nodes.py

@@ -60,29 +60,24 @@ class MaskContourProcessor:
         self.stroke_color = stroke_color
         self.element_color = element_color
 
-    def calculate_mask_centroid(self, mask_array):
-        """Calculate the center of mass (centroid) of a binary mask.
-
-        Args:
-            mask_array (ndarray): Binary mask array with values between 0 and 1
-
-        Returns:
-            tuple: (x, y) coordinates of the mask's centroid
-        """
-        # Create binary mask by thresholding at 0.5
-        binary_mask = (mask_array >= 0.5).astype(np.float32)
+    def calculate_mask_centroid(self, mask):
+        """Calculate the centroid of the mask."""
+        binary_mask = (mask > 0.5).astype(np.float32)
         
-        # Find coordinates of mask pixels
-        y_coords, x_coords = np.nonzero(binary_mask)
+        # Get non-zero coordinates
+        coords = np.nonzero(binary_mask)
         
-        if len(x_coords) == 0:  # Handle empty mask case
-            return (mask_array.shape[1] // 2, mask_array.shape[0] // 2)
+        if len(coords[0]) == 0:  # If mask is empty
+            # Return center of image
+            height, width = mask.shape
+            return width // 2, height // 2
         
-        # Calculate average x and y coordinates
-        center_x = np.mean(x_coords)
+        # Calculate centroid from non-zero coordinates
+        y_coords, x_coords = coords
         center_y = np.mean(y_coords)
+        center_x = np.mean(x_coords)
         
-        return (center_x, center_y)
+        return center_x, center_y
 
     def detect_edge_points(self, mask_array, center):
         """Detect and sort edge points of a mask in clockwise order.

tests/test_mask_contour_processor.py

@@ -0,0 +1,89 @@
+import unittest
+import numpy as np
+import torch
+from nodes.nodes import MaskContourProcessor
+
+class TestMaskContourProcessor(unittest.TestCase):
+    def setUp(self):
+        """Set up test fixtures before each test method."""
+        self.processor = MaskContourProcessor()
+        
+        # Create some common test masks
+        self.empty_mask = np.zeros((64, 64), dtype=np.float32)
+        self.full_mask = np.ones((64, 64), dtype=np.float32)
+        self.circle_mask = np.zeros((64, 64), dtype=np.float32)
+        center = (32, 32)
+        radius = 20
+        y, x = np.ogrid[:64, :64]
+        self.circle_mask[(x - center[0])**2 + (y - center[1])**2 <= radius**2] = 1
+
+    def test_calculate_mask_centroid(self):
+        """Test centroid calculation for various masks."""
+        # Test empty mask
+        center_x, center_y = self.processor.calculate_mask_centroid(self.empty_mask)
+        self.assertEqual((center_x, center_y), (32, 32))  # Should return center of image
+
+        # Test full mask
+        center_x, center_y = self.processor.calculate_mask_centroid(self.full_mask)
+        self.assertEqual((center_x, center_y), (31.5, 31.5))  # Should be center of image
+
+        # Test circle mask
+        center_x, center_y = self.processor.calculate_mask_centroid(self.circle_mask)
+        self.assertAlmostEqual(center_x, 32, delta=1)
+        self.assertAlmostEqual(center_y, 32, delta=1)
+
+    def test_detect_edge_points(self):
+        """Test edge point detection."""
+        center = (32, 32)
+        edge_points = self.processor.detect_edge_points(self.circle_mask, center)
+        
+        # Check that we got some edge points
+        self.assertTrue(len(edge_points) > 0)
+        
+        # Check that all points are on the edge (should have at least one
+        # neighboring pixel with different value)
+        for x, y in edge_points:
+            has_different_neighbor = False
+            for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
+                nx, ny = x + dx, y + dy
+                if (0 <= nx < 64 and 0 <= ny < 64 and 
+                    self.circle_mask[ny, nx] != self.circle_mask[y, x]):
+                    has_different_neighbor = True
+                    break
+            self.assertTrue(has_different_neighbor)
+
+    def test_redistribute_points(self):
+        """Test point redistribution."""
+        # Create a simple square of points
+        square_points = [(0,0), (0,10), (10,10), (10,0)]
+        
+        # Test redistribution to different numbers of points
+        for target_count in [4, 8, 16]:
+            new_points = self.processor.redistribute_points(square_points, target_count)
+            self.assertEqual(len(new_points), target_count)
+
+    def test_process_mask(self):
+        """Test the main processing function."""
+        # Convert numpy mask to tensor
+        mask_tensor = torch.from_numpy(self.circle_mask)
+        
+        # Process the mask
+        result = self.processor.process_mask(
+            mask_tensor,
+            line_length=1.0,
+            line_count=20,
+            line_width=0.01
+        )
+        
+        # Check return type and shape
+        self.assertIsInstance(result, tuple)
+        self.assertEqual(len(result), 1)
+        self.assertIsInstance(result[0], torch.Tensor)
+        self.assertEqual(result[0].shape, mask_tensor.shape)
+        
+        # Check that the output contains at least the original mask
+        # (output should be >= input everywhere)
+        self.assertTrue(torch.all(result[0] >= mask_tensor))
+
+if __name__ == '__main__':
+    unittest.main()