Numpy, векторизованная функция на масках с несколькими метками

Question

У меня есть следующий код для замены каждого значения изображения значением медианы (или любой другой функции) на основе масок из изображения меток, созданного на этапе сегментации. Такое ощущение, что цикл for можно векторизовать. Какой лучший подход сделать это?

Я смотрел на создание отдельного индексного массива для каждой метки, но в конечном итоге не видел, как это поможет.

import numpy as np
from skimage.segmentation import slic
from skimage import data, io

def create_segment_image(original_image, labels_image):
    segment_image = np.zeros(original_image.shape, original_image.dtype)

    for label in np.unique(labels_image):
        segment_image[labels_image==label] = np.median(original_image[labels_image==label], axis=0)

    return segment_image

if __name__ == '__main__':

    original_image = data.astronaut()
    labels_image = slic(original_image, n_segments=1000, max_iter=10, compactness=7, sigma=1, convert2lab=True, enforce_connectivity=True, min_size_factor=0.1, max_size_factor=3, slic_zero=False)
    segment_image = create_segment_image(original_image, labels_image)
    # io.imsave('images/segment_image.png', segment_image)

Answer 1

Я не знаю о векторизации самого внутреннего цикла. Вызов median вычисляет разное количество элементов каждый раз, что затруднит размещение всех вызовов в одном массиве.

С другой стороны, есть довольно низко висящие фрукты с точки зрения того, как вы выбираете элементы по меткам. Вы находите индексы для каждой метки дважды в своей исходной функции, вычисляя массив индексов только один раз, когда сбивает примерно 25% времени выполнения

def create_segment_image_2(original_image, labels_image):
    segment_image = np.zeros(original_image.shape, original_image.dtype)
    for label in np.unique(labels_image):
        inds = np.where(labels_image == label)
        segment_image[inds] = np.median(original_image[inds], axis=0)    
    return segment_image

Вы можете получить еще большие улучшения, отсортировав индексы массива по меткам, а затем воспользовавшись этой сортировкой, выбрав элементы изображения для медианы. Замена множества поисков одной сортировкой дает ускорение в 20 раз.

def create_segment_image_3(original_image, labels_image):
    segment_image = np.zeros(original_image.shape, original_image.dtype)
    # sort the indices by their labels
    labelinds = np.argsort(labels_image, None)
    labels = np.unique(labels_image)
    # use the searchsorted to find the indices for each label
    rights = np.searchsorted(labels_image.flatten(), labels, side='right', sorter=labelinds)
    left = 0
    for right in rights:
        # choose our block of the image array
        inds = labelinds[left:right]
        # convert back to a two dimensional index array
        inds = [inds // original_image.shape[1], inds % original_image.shape[1]]
        segment_image[inds] = np.median(original_image[inds], axis=0)
        # update our boundaries
        left = right
    return segment_image

Бенчмаркинг в ipython

In [54]: %timeit create_segment_image(original_image, labels_image)
2.15 s ± 29.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [55]: %timeit create_segment_image_2(original_image, labels_image)
1.48 s ± 4.68 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [56]: %timeit create_segment_image_3(original_image, labels_image)
121 ms ± 561 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Подтверждение того, что наши новые решения дают тот же результат, что и старое.

In [57]: np.all(create_segment_image_2(original_image, labels_image) == create_segment_image(original_image, labels_image))
Out[57]: True

In [58]: np.all(create_segment_image_3(original_image, labels_image) == create_segment_image(original_image, labels_image))
Out[58]: True