From e3521809ec7de7b67b287636be466b29510fda37 Mon Sep 17 00:00:00 2001 From: Patrick Morgan Date: Sun, 15 Dec 2024 20:47:48 -0800 Subject: [PATCH] fix --- .gitignore | 1 + README.md | 2 + nodes/nodes.py | 31 +++--- ...est_mask_contour_processor.cpython-311.pyc | Bin 0 -> 5337 bytes tests/test_mask_contour_processor.py | 89 ++++++++++++++++++ 5 files changed, 105 insertions(+), 18 deletions(-) create mode 100644 tests/__pycache__/test_mask_contour_processor.cpython-311.pyc create mode 100644 tests/test_mask_contour_processor.py 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 0000000000000000000000000000000000000000..da7095f69392482702869096cf9b0e7f7d73e246 GIT binary patch literal 5337 zcmbVQZ*1Gf6+cp>e75B}aur#UEmz65#!BrtX_mxo-288vwu#fkPFEP+I9a4@*|Ovu zsnu4c!2?o2Edtan3Pb^1Kr6P)LDC@~QXs&-WWYZ3+eCqY1p)y!4BOy+QjrcA`el1} zlxQijv$iMl@p$judw0k8yZ27N>*#1hw0tg2{i6e+zp$G(=%$c5Z5q0)tDygN? zWb?H8dP2t;do^TE6QH41X6*WH{syejEw)OO37YKwO`7ag{tlX2>h_v6eV;Mmx+d9? z<76SnH_gWWuCgdQ-{UPB+0i2f2 zI2pP|yd^cu)_eq7WD&k1T7zE}zz(E~?}z!BhdKJfRfbQgVJblnY<{3E)5#PE;M9uGi%xOnD^6L()L z^&HiEjv75j^MZkW6}La{thk@fJ0EnquPziEg)X?q3X^bSqjSgmLY~X-t6aFqqzgL= zvEqf5`1Qfn9lHOy)ysDd-rI5ag`2s1m-RizbpNsYr@uV@$5+e#b2=U^vFjlWA4P4X z!V6tehVYM}^%|gWhKjI4G5>f02^*_7K;o7y#eAI0@mW-~5!Tp0%Vqf(>`u*i=f0gy zd%yTzHXqqzzNJljAKKEUy(cx#Q+CM0ZLv8TP@Bcg<@wF!t(`hEK;1G#@jy~>ITVc~ zqUl6Li>HzyfOY5>5fV?O)eyy)x`)u6D9rzpjc4iTF%p6r!!mJE4n?_=&?2MZ7V;D{ znB8@K77nGl9nwX)?GY7VroDVQ9Z3+n8IbL?*=;o+Pt2uMiy3y1Zpy(HZlwjgwy08P z0L+qJD$q+2DQ#gib6`sdKB*fWnEGcROQ`BZeLGivwiY~D4xYS!>QCqO;J6VS*Kr_^ zm-{Q&wT69V>?`ya&y@Uws|z|FHt=u>4?o&*R1c0B!LhA6W_3Jb;E@s@d8DIpYFrOq zGJ=;Hd-j!aUtwXj>(0Is?$hxB10N{i18g*#vw#a1>4P~np_Td%L+cVltIXxN?}pgy zBk0>2t(qOvXuN4gfXvS_JT|w2hx4b1rB)Hp(M|ttIq(xbyNN>#vf|jkFu><+!%G6ctY0E$eU$*6cJ=+ z4)VDqcXfQreM@|Fo>x`51>!Bv1Q<3Ic}fY*r{YO1BrBQ{rP#EP$O7i+umF+vzfjm~ zm}gL|*Xd?T!a$&%)hSL{$Q>UIixx<`EVU9fBr!v(1;jQBh~p%!n2v;!gs@CuNPIxd z_D1m@+8m(G-sy-c$?;fBA+V=NNhLltJ((gVFE7Bo2yq;q$wieW007%svC!MiVEFiG zGJBs4+{b(T0JN(dc<~bY3yXT*z$eKy@1e5y&^^HzKJ(?Xy7z+Ny^wDQfVFmf*!e-{ zr;ftwy072x^{?(L`v!D8XyCyT9;{-tJyb=w^~a37e=2`Gf4TCb-6hw~N?^E(#Et+X zA32x#zpItutG2DeoO08ebYd9#ms^ClKD`=0F1n{+t@n$@2JVKK=6LdiD_SH1(hH@K^c_|_3d z@@Fgl;7Z?`|8UuV_};lMAuD{#@V}LR4ZLK{)mwJ;emYhRT$fg))r9UkY`6}W*i{wl z8lK=T|6)O*@sKuSCufaYVe?03A%ty)FnUy^OY~TQ!w9&a?6@u50RAmohJ_D1O=|3x zQy<~R7ezzRP*shfthX$7*scrsKw_QaD@ zp;$U;g`p{H0M(;}4r%IL&+lLPYv#()u)}Om#FL5yL1RjrHrv=s4KS?!PF&WeHzL(& zxWg1R_yYoDVCsbw;=|( zDO_y37!6|+OHwj2r%007CP{NCISnUo$#h8)&Fw7?3BqcSU6kynYj|yWK^95vSk0Ne7iWtRZ#;CbchG5$R>SVB7|w1^kWe=Mw@uHi!W8tfl#oCkbZ-BaoQR}Y?OkLw-1kN*T2Ir`hKvj`9XfK*; zvk88KN37*JLVf~mbO%*G0#bEw99Kb|OYB!cU8VX}L9QkCtDu2W^Hmkz;5eAl=CrV2 uJhS&J#{9ZbMCOs;%4f{oRlApi_E5XgN1 literal 0 HcmV?d00001 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()