def test_3d(self):
a = torch.tensor(
[[list(range(7)), list(range(7)),
list(range(7))],
[list(range(7)), list(range(7)),
list(range(7))]],
dtype=torch.float,
)
a = a[None][None]
a = a.expand(2, 3, -1, -1, -1)
expected = np.array([
[
[2.0, 6.0, 12.0, 18.0, 24.0, 30.0, 22.0],
[3.0, 9.0, 18.0, 27.0, 36.0, 45.0, 33.0],
[2.0, 6.0, 12.0, 18.0, 24.0, 30.0, 22.0],
],
[
[2.0, 6.0, 12.0, 18.0, 24.0, 30.0, 22.0],
[3.0, 9.0, 18.0, 27.0, 36.0, 45.0, 33.0],
[2.0, 6.0, 12.0, 18.0, 24.0, 30.0, 22.0],
],
])
expected = expected
# testing shapes
k = torch.tensor([[[1, 1, 1], [1, 1, 1], [1, 1, 1]]])
for kernel in (k, k[None], k[None][None]):
out = apply_filter(a, kernel)
np.testing.assert_allclose(out.cpu().numpy()[1][2],
expected,
rtol=1e-4)
if torch.cuda.is_available():
out = apply_filter(a.cuda(), kernel.cuda())
np.testing.assert_allclose(out.cpu().numpy()[0][1],
expected,
rtol=1e-4)
def test_1d(self):
a = torch.tensor([[list(range(10))]], dtype=torch.float)
out = apply_filter(a, torch.tensor([-1, 0, 1]), stride=1)
expected = np.array(
[[[1.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, -8.0]]])
np.testing.assert_allclose(out.cpu().numpy(), expected, rtol=1e-4)
if torch.cuda.is_available():
out = apply_filter(a.cuda(), torch.tensor([-1, 0, 1]).cuda())
np.testing.assert_allclose(out.cpu().numpy(), expected, rtol=1e-4)
def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
"""
Args:
img: torch tensor data to extract the contour, with shape: [channels, height, width[, depth]]
Raises:
ValueError: When ``image`` ndim is not one of [3, 4].
Returns:
A torch tensor with the same shape as img, note:
1. it's the binary classification result of whether a pixel is edge or not.
2. in order to keep the original shape of mask image, we use padding as default.
3. the edge detection is just approximate because it defects inherent to Laplace kernel,
ideally the edge should be thin enough, but now it has a thickness.
"""
img = convert_to_tensor(img, track_meta=get_track_meta())
img_: torch.Tensor = convert_to_tensor(img, track_meta=False)
spatial_dims = len(img_.shape) - 1
img_ = img_.unsqueeze(0) # adds a batch dim
if spatial_dims == 2:
kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]],
dtype=torch.float32)
elif spatial_dims == 3:
kernel = -1.0 * torch.ones(3, 3, 3, dtype=torch.float32)
kernel[1, 1, 1] = 26.0
else:
raise ValueError(
f"{self.__class__} can only handle 2D or 3D images.")
contour_img = apply_filter(img_, kernel)
contour_img.clamp_(min=0.0, max=1.0)
output, *_ = convert_to_dst_type(contour_img.squeeze(0), img)
return output
def test_2d(self):
a = torch.tensor(
[[[list(range(7)),
list(range(7, 0, -1)),
list(range(7))]]],
dtype=torch.float)
expected = np.array([
[14.0, 21.0, 21.0, 21.0, 21.0, 21.0, 14.0],
[15.0, 24.0, 27.0, 30.0, 33.0, 36.0, 25.0],
[14.0, 21.0, 21.0, 21.0, 21.0, 21.0, 14.0],
])
expected = expected[None][None]
out = apply_filter(a, torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]))
np.testing.assert_allclose(out.cpu().numpy(), expected, rtol=1e-4)
if torch.cuda.is_available():
out = apply_filter(
a.cuda(),
torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]).cuda())
np.testing.assert_allclose(out.cpu().numpy(), expected, rtol=1e-4)
def test_wrong_args(self):
with self.assertRaisesRegex(ValueError, ""):
apply_filter(torch.ones((1, 2, 3, 2)), torch.ones((2, )))
with self.assertRaisesRegex(NotImplementedError, ""):
apply_filter(torch.ones((1, 1, 1, 2, 3, 2)), torch.ones((2, )))
with self.assertRaisesRegex(TypeError, ""):
apply_filter(((1, 1, 1, 2, 3, 2)), torch.ones((2, )))