How to cluster a score / probability map and get the modes (with variable numbers) in pytorch or numpy

Question

I have a 2d probability map (please correct me if I use any term wrong). Something like this:

Here yellow is a high value and violet is zero. Please ignore the re