relabeller.py 文件源码

def relabeller(in_data, list_old_labels, list_new_labels, verbose=True):
    """
    :param in_data: array corresponding to an image segmentation.
    :param list_old_labels: list or tuple of labels
    :param list_new_labels: list or tuple of labels of the same len as list_old_labels
    :param verbose:
    :return: array where all the labels in list_new_labels are substituted with list_new_label in the same order.
    """

    if isinstance(list_new_labels, int):
        list_new_labels = [list_new_labels, ]
    if isinstance(list_old_labels, int):
        list_old_labels = [list_old_labels, ]

    # sanity check: old and new have the same number of elements
    if not len(list_old_labels) == len(list_new_labels):
        raise IOError('Labels list does not have the same length.')

    new_data = copy.deepcopy(in_data)

    for k in range(len(list_new_labels)):
        places = in_data == list_old_labels[k]
        if np.any(places):
            np.place(new_data, places, list_new_labels[k])
            if verbose:
                print('Label {0} substituted with label {1}'.format(list_old_labels[k], list_new_labels[k]))
        else:
            print('Label {0} not present in the array'.format(list_old_labels[k]))

    return new_data