def __init__(
self,
keys: KeysCollection,
spatial_dims: int = 2,
sigma: Union[Sequence[float], float, Sequence[torch.Tensor], torch.Tensor] = 0.0,
prob_threshold: float = 0.5,
box_size: Union[int, Sequence[int]] = 48,
allow_missing_keys: bool = False,
) -> None:
super().__init__(keys, allow_missing_keys)
self.prob_nms = ProbNMS(
spatial_dims=spatial_dims, sigma=sigma, prob_threshold=prob_threshold, box_size=box_size
)
def __call__(self,
probs_map: Union[np.ndarray, torch.Tensor],
resolution_level: int = 0):
"""
probs_map: the input probabilities map, it must have shape (H[, W, ...]).
resolution_level: the level at which the probabilities map is made.
"""
resolution = pow(2, resolution_level)
org_outputs = ProbNMS.__call__(self, probs_map)
outputs = []
for org_output in org_outputs:
prob = org_output[0]
coord = np.asarray(org_output[1:])
coord_wsi = ((coord + 0.5) * resolution).astype(int)
outputs.append([prob] + list(coord_wsi))
return outputs
def test_output(self, class_args, probs_map, expected):
nms = ProbNMS(**class_args)
output = nms(probs_map)
assert_allclose(output, expected)