GrowMask now works with mask batches.

This commit is contained in:
comfyanonymous 2023-09-26 02:53:57 -04:00
parent d76d71de3f
commit 1d36dfb9fe
1 changed files with 13 additions and 9 deletions

View File

@ -327,15 +327,19 @@ class GrowMask:
kernel = np.array([[c, 1, c], kernel = np.array([[c, 1, c],
[1, 1, 1], [1, 1, 1],
[c, 1, c]]) [c, 1, c]])
output = mask.numpy().copy() mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
while expand < 0: out = []
output = scipy.ndimage.grey_erosion(output, footprint=kernel) for m in mask:
expand += 1 output = m.numpy()
while expand > 0: while expand < 0:
output = scipy.ndimage.grey_dilation(output, footprint=kernel) output = scipy.ndimage.grey_erosion(output, footprint=kernel)
expand -= 1 expand += 1
output = torch.from_numpy(output) while expand > 0:
return (output,) output = scipy.ndimage.grey_dilation(output, footprint=kernel)
expand -= 1
output = torch.from_numpy(output)
out.append(output)
return (torch.cat(out, dim=0),)