import matplotlib.pyplot as plt
from astropy.convolution import convolve
from photutils.background import Background2D, MedianBackground
from photutils.datasets import make_100gaussians_image
from photutils.segmentation import (deblend_sources, detect_sources,
                                    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  # subtract the background

threshold = 1.5 * bkg.background_rms

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

n_pixels = 10
segment_map = detect_sources(convolved_data, threshold, n_pixels=n_pixels)
deblended_segment_map = deblend_sources(convolved_data, segment_map,
                                        n_pixels=n_pixels,
                                        progress_bar=False)

fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, figsize=(10, 4))
slc = (slice(273, 297), slice(425, 444))
ax1.imshow(data[slc], origin='lower')
ax1.set_title('Background-subtracted Data')

segm_cutout = segment_map[slc]
segm_cutout.imshow(ax=ax2, cmap=segment_map.cmap)
ax2.set_title('Original Segment')

deblended_segm_cutout = deblended_segment_map[slc]
deblended_segm_cutout.imshow(ax=ax3, cmap=deblended_segment_map.cmap)
ax3.set_title('Deblended Segments')
fig.tight_layout()