import matplotlib.pyplot as plt
from astropy.convolution import convolve
from astropy.visualization import simple_norm
from photutils.background import Background2D, MedianBackground
from photutils.datasets import make_100gaussians_image
from photutils.segmentation import (SourceFinder,
                                    make_2dgaussian_kernel)

data = make_100gaussians_image()

bkg_estimator = MedianBackground()
bkg = Background2D(data, (50, 50), filter_size=(3, 3),
                   bkg_estimator=bkg_estimator)
data -= bkg.background

threshold = 1.5 * bkg.background_rms

kernel = make_2dgaussian_kernel(3.0, size=5)
convolved_data = convolve(data, kernel)

finder = SourceFinder(n_pixels=10, progress_bar=False)
segment_map = finder(convolved_data, threshold)

fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(10, 12.5))
segment_map.imshow(ax=ax1)
ax1.set_title('Segmentation Image')
segment_map.plot_patches(ax=ax1, edgecolor='white', lw=1.5)
norm = simple_norm(data, 'sqrt', percent=99.5)
ax2.imshow(data, norm=norm, origin='lower')
ax2.set_title('Background-subtracted Data')
segment_map.plot_patches(ax=ax2, edgecolor='white', lw=1.5)
fig.tight_layout()