torch_em.metric.cldice

 1import numpy as np
 2import torch
 3from skimage.morphology import skeletonize
 4
 5from ..loss.cldice import SoftSkeletonize
 6
 7# From "clDice -- A Novel Topology-Preserving Loss Function for Tubular Structure Segmentation":
 8# https://arxiv.org/abs/2003.07311
 9
10
11def cl_score(img, skel):
12    """Compute the skeleton volume intersection.
13
14    Args:
15        img: image
16        skel: skeleton
17
18    Returns:
19        Skeleton volume intersection.
20    """
21    return np.sum(img * skel) / np.sum(skel)
22
23
24def clDice(input_, target, skeletonize_method="skimage", num_iter=5):
25    """Compute the clDice score between binary input and target.
26
27    Args:
28        input_: The binary input.
29        target: The binary target.
30        skeletonize_method: The skeletonziation method. Either `skimage` for
31            `skimage.morphology.skeletonize` or `soft` for `torch_em.loss.SoftSkeletonize`.
32        num_iter: Number of iterations for soft skeletonization.
33            Only used if skeletonize_method is `soft`
34
35    Returns:
36        The clDice score.
37    """
38    if input_.shape != target.shape:
39        raise ValueError(f"Expect input and target of same shape, got: {input_.shape}, {target.shape}.")
40
41    if skeletonize_method == "skimage":
42        skel_input = skeletonize(input_)
43        skel_target = skeletonize(target)
44
45    elif skeletonize_method == "soft":
46        soft_skeletonize = SoftSkeletonize(num_iter=num_iter)
47
48        # add batch and channel dims for `SoftSkeletonize`
49        input_tensor = torch.from_numpy(input_).float().unsqueeze(0).unsqueeze(0)
50        target_tensor = torch.from_numpy(target).float().unsqueeze(0).unsqueeze(0)
51
52        # convert skeletons back to numpy
53        skel_input = soft_skeletonize(input_tensor).squeeze().numpy()
54        skel_target = soft_skeletonize(target_tensor).squeeze().numpy()
55    else:
56        raise ValueError("Unknown option for `skeletonize_method`. Valid options are `skimage` and `soft`.")
57
58    # Tprec = |S_P ∩ V_L| / |S_P|
59    # Tsens = |S_L ∩ V_P| / |S_L|
60    t_prec = cl_score(target, skel_input)
61    t_sens = cl_score(input_, skel_target)
62
63    return 2.*(t_prec*t_sens) / max(t_prec+t_sens, 1e-7)