GrowMask now works with mask batches.
This commit is contained in:
parent
d76d71de3f
commit
1d36dfb9fe
|
@ -327,7 +327,10 @@ 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]))
|
||||||
|
out = []
|
||||||
|
for m in mask:
|
||||||
|
output = m.numpy()
|
||||||
while expand < 0:
|
while expand < 0:
|
||||||
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
||||||
expand += 1
|
expand += 1
|
||||||
|
@ -335,7 +338,8 @@ class GrowMask:
|
||||||
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
||||||
expand -= 1
|
expand -= 1
|
||||||
output = torch.from_numpy(output)
|
output = torch.from_numpy(output)
|
||||||
return (output,)
|
out.append(output)
|
||||||
|
return (torch.cat(out, dim=0),)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue