diff --git a/.gitignore b/.gitignore index d56e2d6..d89672b 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,4 @@ logs/ # Cache .cache/ .temp/ +nodes/__pycache__/nodes.cpython-311.pyc diff --git a/README.md b/README.md index 05d6653..1451918 100644 --- a/README.md +++ b/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 diff --git a/nodes/nodes.py b/nodes/nodes.py index ab930b4..10147b6 100644 --- a/nodes/nodes.py +++ b/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. diff --git a/tests/__pycache__/test_mask_contour_processor.cpython-311.pyc b/tests/__pycache__/test_mask_contour_processor.cpython-311.pyc new file mode 100644 index 0000000..da7095f Binary files /dev/null and b/tests/__pycache__/test_mask_contour_processor.cpython-311.pyc differ diff --git a/tests/test_mask_contour_processor.py b/tests/test_mask_contour_processor.py new file mode 100644 index 0000000..fa0daa0 --- /dev/null +++ b/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()