Function _argmax_nd

tests/transforms/test_generate_heatmap.py:24–28

argmax for N-D array → returns coordinate vector (z,y,x) or (y,x).

(x)

def _argmax_nd(x) -> np.ndarray:
    """argmax for N-D array → returns coordinate vector (z,y,x) or (y,x)."""
    if isinstance(x, torch.Tensor):
        x = x.cpu().numpy()
    return np.asarray(np.unravel_index(np.argmax(x), x.shape))


# Test cases for 2D array inputs with different data types