def mask_segmentation(seg_prediction, input_tensor, threshold):
    """
    Threshold and apply mask to input image
    :param seg_prediction: segmentation prediction probabilities
    :param input_tensor: input image tensor
    :param threshold: Keep input values above threshold values, otherwise 0
    :return: Tensor of segmented image
    """
    seg_channel = seg_prediction[0, 1, :, :].cpu().detach()
    mask = (seg_channel > threshold).float() * 1
    img = input_tensor[0, 1, :, :].cpu().detach()
    img[mask == 0.0] = 0.0
    return img