[Isolate a galaxy in an image] #astro #visualisation #python

isolate_galaxy_recipe.py

import numpy as np
from scipy.ndimage import binary_dilation


def isolate_galaxy(image, seg_map, galaxy_value,
                   n_dilation = 5, 
                   use_background_pixels = False, 
                   noise_factor = 0.0):
    """
    Isolate a galaxy in an image by replacing its neighbors with background noise

    Basic recipe to replace detected sources around a selected galaxy
    with either randomly selected pixels from the background, or a random
    realisation of the background noise.

    A segmentation map of the image containing all the detected sources is required
    for the task. On this segmentation map, it is assumed that the background has a
    value of 0 and that each detected source has an specific integer value.

    Parameters
    ----------
    image : ndarray
        postage stamp with several galaxies
    seg_map : ndarray
        segmentation map of the image
    galaxy_value : int
        value of the segmentation map corresponding to the selected galaxy
    n_dilation : int (default 5)
        number of dilation iterations for the neighbor masks
    use_background_pixels : bool (default False)
        select the strategy for filling in neighbor pixels:
        - if `True`, the neighbor pixels are replaced with random pixels 
          from the background
        - if `False`, the neighbor pixels are replaced with samples from 
          a Gaussian distribution following the background pixel properties
    noise_factor : float (default 0.0)
        this parameter allows to add a Gaussian noise image to the final image
        up to a given amplitude, in order to smooth out the edges of the neighbors

    Returns
    -------
    masked_image : ndarray

    """
    masked_image = image.copy()

    # Create binary mask of the selected source and of all sources
    sources = binary_dilation(seg_map, iterations=n_dilation)
    selected_source = binary_dilation(np.where(seg_map == galaxy_value, 1, 0),
                                      iterations=n_dilation)

    # Compute the binary mask of all sources BUT the central galaxy
    sources_except_central = np.logical_xor(sources, selected_source)

    # Create binary mask for background and compute its properties
    background_mask = np.logical_not(sources)
    background_std = np.std(image * background_mask)

    if use_background_pixels:
        # Select random pixels from the background and fill in the voids
        n_pixels_to_fill_in = sources_except_central.sum()
        random_background_pixels = np.random.choice(
            image[background_mask],
            size=n_pixels_to_fill_in
        )
        masked_image[sources_except_central] = random_background_pixels
    else:
        # Create a realisation of the background from its properties
        random_background = np.random.normal(scale=background_std, size=image.shape)
        masked_image[sources_except_central] = random_background[sources_except_central]
    
    # Optionally add a layer of noise to smooth the image around the neighbor edges
    masked_image += noise_factor * np.random.normal(scale=background_std, size=image.shape)

    return masked_image.astype(image.dtype)