from scipy.ndimage.measurements import label

from scipy.ndimage.filters import gaussian_filter

from scipy.ndimage.morphology import binary_fill_holes

import numpy as np

def get_main_connected_component(data):

    labeled_array, num_features = label(data)

    i = np.argmax([np.sum(labeled_array == _) for _ in range(1, num_features + 1)]) + 1

    return labeled_array == i

def postprocess_prediction(pred, gaussian_std=1, threshold=0.5, fill_holes=True, connected_component=True):

    pred = gaussian_filter(pred, gaussian_std) > threshold

    if fill_holes:

        pred = binary_fill_holes(pred)

    if connected_component:

        pred = get_main_connected_component(pred)

    return pred