Перегон картинок из Pillow в NumPy/OpenCV всего за два копирования памяти

Стоп, что? В смысле «всего»? Разве преобразование из одного формата в другой нельзя сделать за одно копирование, а лучше вообще без копирования?

Да, это кажется безумием, но более привычные методы преобразования картинок работают в 1,5-2,5 раза медленнее (если нужен не read-only объект). Сегодня я покопаюсь в кишках обеих библиотек, расскажу почему так получилось и кто виноват. А также покажу финальный результат, который работает так же, только быстрее. Никаких репозиториев или пакетов не будет, только рассказ и рабочий код в конце. Но давайте обо всём по порядку.

Pillow — это библиотека для работы с изображениями на языке Python. Поддерживает разные форматы, имеет ленивую загрузку, дает доступ к метаинформации из файла. Короче делает все, что нужно для загрузки изображений.

NumPy — библиотека-комбайн для работы с многомерными массивами. Базовая библиотека для целой кучи научных библиотек, библиотек компьютерного зрения и машинного обучения.

OpenCV — самая популярная библиотека компьютерного зрения. Имеет огромное количество функций. Не имеет собственного внутреннего формата хранения для изображений, вместо этого использует массивы NumPy. Сценарий, когда нужно преобразовать изображение из Pillow в NumPy, чтобы дальше работать с ним с помощью OpenCV, чрезвычайно распространенный.

Для разнообразия сегодня я буду запускать бенчмарки на Raspberry Pi 4 1800 MHz под 64-разрядной Raspberry Pi OS. В конце концов, где ещё может понадобиться компьютерное зрение, как не на Малинке :-)

На случай, если вы не знаете как настроить окружение

Подключаетесь по SSH и ставите менеджер виртуального окружения:

$ sudo apt install python3-venv

Дальше sudo вам не понадобится. Создаете виртуальное окружение:

$ python3 -m venv pil_num_env

Активируете виртуальное окружение:

$ source ./pil_num_env/bin/activate

Обновляете pip:

$ pip install -U pip

Ставите всё, с чем мы будем сегодня работать:

$ pip install ipython pillow numpy opencv-python-headless

Всё готово, заходите в интерактивный интерпретатор:

$ ipython
Python 3.7.3 (default, Jul 25 2020, 13:03:44) 
IPython 7.21.0 -- An enhanced Interactive Python. Type '?' for help.
In [1]: _

Как работает преобразование в NumPy

Существует два общепринятых способа конвертировать изображение Pillow в NumPy, с равной вероятностью вы нагуглите один из них:

1. numpy.array(im)— делает копию из изображения в массив NumPy.

2. numpy.asarray(im)— то же самое, что numpy.array(im, copy=False), то есть якобы не делает копию, а использует память оригинального объекта. На самом деле всё несколько сложнее.

Можно было бы подумать, что во втором случае массив NumPy становится как бы вью на оригинальное изображение, и если изменять массив NumPy, то будет меняться и изображение. На деле это не так:

In [1]: from PIL import Image
In [2]: import numpy
In [3]: im = Image.open('./canyon.jpg').resize((4096, 4096))
In [4]: n = numpy.asarray(im)

In [5]: n[:, :, 0] = 255
ValueError: assignment destination is read-only

In [6]: n.flags
Out[6]: 
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : False
  WRITEABLE : False
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

Это сильно отличается от того, что будет, если использовать функцию numpy.array():

In [7]: n = numpy.array(im)

In [8]: n[:, :, 0] = 255

In [9]: n.flags
Out[9]: 
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

При этом, если провести измерение, функция asarray() действительно работает значительно быстрее:

In [10]: %timeit -n 10 n = numpy.array(im)
257 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [11]: %timeit -n 10 n = numpy.asarray(im)
179 ms ± 786 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Тут копирований явно происходит меньше, за что мы расплачиваемся невозможностью изменять массив. Но само время преобразования остается чудовищно большим по сравнению с одним копированием. Давайте разбираться, на что же оно тратится.

Интерфейс массивов NumPy

Если посмотреть на зависимости и код Pillow, там не найдется упоминаний NumPy (на самом деле найдется, но только в комментариях). То же самое верно и в обратную сторону. Как же изображения конвертируются из одного формата в другой? Оказывается, у NumPy для этого есть специальный интерфейс. Вы делаете специальное свойство у нужного объекта, в котором объясняете NumPy, как ему следует извлечь данные, а он эти данные забирает. Вот упрощенная реализация этого свойства из Pillow:

    @property
    def __array_interface__(self):
        shape, typestr = _conv_type_shape(self)
        return {
            "shape": shape,
            "typestr": typestr,
            "version": 3,
            "data": self.tobytes(),
        }

_conv_type_shape() описывает тип и размер массива, который должен получиться. А всё самое интересное происходит в методе tobytes(). Если проверить, сколько этот метод выполняется, станет понятно, что в общем-то NumPy от себя ничего не добавляет:

In [12]: %timeit -n 10 n = im.tobytes()
179 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Время точно совпадает с временем функции asarray(). Кажется виновник найден, осталось заменить вызов этой функции или ускорить её, и дело в шляпе, верно? Ну, не всё так просто.

Внутреннее устройство памяти в Pillow и NumPy

Устройство массивов в NumPy описывается чрезвычайно просто — это непрерывный кусок памяти, начинающийся с определенного указателя. Плюс есть смещения (strides), которые задаются отдельно по каждому измерению.

В Pillow всё устроено принципиально иначе. Изображение хранится чанками, в каждом чанке находится целое количество строк изображения. Каждый пиксель занимает 1 или 4 байта (не от 1 до 4, а ровно). Соответственно, для каких-то режимов изображения какие-то байты не используются. Например, для RGB не используется последний байт в каждом пикселе, а для черно-белых изображений с альфа-каналом (режим LA) не используются два средних байта для того, чтобы альфа-канал был в последнем байте пикселя.

Всё это я рассказываю, потому что не хочу, чтобы у кого-то остались иллюзии, что можно решить вопрос как-то малой кровью, не переписывая одну или другую библиотеку.

Я думаю, теперь понятно, для чего нужен метод tobytes() — он переводит внутреннее представление изображения Pillow в непрерывный поток байтов одним куском без пропусков: как раз такое, какое может использовать NumPy. NumPy уже получая на вход объект bytes, может либо сделать копию, либо использовать его в режиме read-only. Тут я не уверен, сделано ли это, чтобы нельзя было обойти неизменность объектов bytes в Python, или есть какие-то реальные ограничения на уровне C API. Но, например, если на вход вместо bytes подать bytearray, то массив не будет read-only.

Но давайте всё же посмотрим на упрощенную версию tobytes():

    def tobytes(self):
        self.load()
        # unpack data
        e = Image._getencoder(self.mode, "raw", self.mode)
        e.setimage(self.im)

        data, bufsize, s = [], 65536, 0
        while not s:
            l, s, d = e.encode(bufsize)
            data.append(d)
        if s < 0:
            raise RuntimeError(f"encoder error {s} in tobytes")
        return b"".join(data)

Тут видно, что создается "raw" энкодер и из него получаются чанки изображения не менее 65 килобайт памяти. Это и есть первое копирование: к концу функций у нас всё изображение в виде небольших чанков лежит в массиве data. Последней строкой происходит второе копирование: все чанки собираются в одну большую байтовую строку.

Кто виноват и что делать

Напомню, что библиотеки написаны так, чтобы интерфейс был, а явного использования библиотеками друг друга не было. В таких условиях, я думаю, что это почти оптимальное решение. Но что, если у нас нет такого ограничения, а скорость хочется получить максимальную?

Первое, что хочется отметить: отказываться от энкодера — не вариант. Кто знает, какие детали реализации он от нас срывает. Переносить это всё на уровень Python или переписывать часть на C — последнее дело.

Кажется, намного разумнее было бы в tobytes()заранее выделить буфер нужного размера, и уже в него записывать чанки. Но очевидно, что интерфейс энкодера так не работает: он уже возвращает чанки упакованные в объекты bytes. Тем не менее, если эти чанки не складировать, а сразу копировать в буфер, эти данные не будут вымываться из L2 кэша и быстро попадут куда надо. Что-то вроде такого:

def to_mem(im):
    im.load()
    e = Image._getencoder(im.mode, "raw", im.mode)
    e.setimage(im.im)

    mem = ... # we don't know yet

    bufsize, offset, s = 65536, 0, 0
    while not s:
        l, s, d = e.encode(bufsize)
        mem[offset:offset + len(d)] = d
        offset += len(d)
    if s < 0:
        raise RuntimeError(f"encoder error {s} in tobytes")
    return mem

Что же будет вместо mem. В идеале это должен быть массив NumPy. Создать его не представляет проблем, мы уже видели какие у него будут параметры в __array_interface__:

In [13]: shape, typestr = Image._conv_type_shape(im)
In [14]: data = numpy.empty(shape, dtype=numpy.dtype(typestr))

Но если попробовать вместо mem взять просто его плоскую версию, то ничего не выйдет:

In [15]: mem = data.reshape((data.size,))
In [16]: mem[0:4] = b'abcd'
ValueError: invalid literal for int() with base 10: b'abcd'

В данном случае кажется странным, что нельзя в массив байтов по срезу поместить байты. Но не забывайте, что, во-первых, слева могут быть не только байты, а во-вторых, библиотека называется NumPy, то есть работает с числами. К счастью, NumPy дает доступ и к непосредственной памяти массива прямо из Python. Это свойство data:

In [17]: data.data
Out[17]: <memory at 0x7f78854d68>

In [18]: data.data[0] = 255
NotImplementedError: sub-views are not implemented

In [19]: data.data.shape
Out[19]: (4096, 4096, 3)

In [20]: data.data[0, 0, 0] = 255

Там находится объект memoryview. Вот только этот memoryview какой-то странный: он тоже многомерный, как и сам массив NumPy, ещё у него такой же тип объектов, как у самого массива. К счастью, это легко исправляется методом cast:

In [21]: mem = data.data.cast('B', (data.data.nbytes,))

In [22]: mem.nbytes == mem.shape[0]
Out[22]: True

In [23]: mem[0], mem[1]
Out[23]: (255, 0)

In [24]: mem[0:4] = b'1234'

In [25]: mem[0], mem[1]
Out[25]: (49, 50)

Складываем пазл вместе:

def to_numpy(im):
    im.load()
    # unpack data
    e = Image._getencoder(im.mode, 'raw', im.mode)
    e.setimage(im.im)

    # NumPy buffer for the result
    shape, typestr = Image._conv_type_shape(im)
    data = numpy.empty(shape, dtype=numpy.dtype(typestr))
    mem = data.data.cast('B', (data.data.nbytes,))

    bufsize, s, offset = 65536, 0, 0
    while not s:
        l, s, d = e.encode(bufsize)
        mem[offset:offset + len(d)] = d
        offset += len(d)
    if s < 0:
        raise RuntimeError("encoder error %d in tobytes" % s)
    return data

Проверяем:

In [26]: n = to_numpy(im)

In [27]: numpy.all(n == numpy.array(im))
Out[27]: True

In [28]: n.flags
Out[28]: 
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

In [29]: %timeit -n 10 n = to_numpy(im)
101 ms ± 260 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Круто! Имеем ускорение в 2,5 раза с тем же функционалом и меньшее количество аллокаций.

Бенчмарки

В статье я взял достаточно большую картинку для тестов. Нет, дело не в том, что to_numpy() не дает ускорения на меньших размерах (ещё как даёт!). Дело в том, что в общем случае очень сложно добиться какого-то постоянного времени работы, когда дело касается выделения памяти. Аллокатор может затребовать новую память у системы, а может и старую сохранить. Может решить заполнить её нулями, а может и так отдать. В этом смысле работа с большими массивами хотя бы дает стабильный результат: мы всегда получаем худший случай.

Код:

In [30]: for i in range(6, 0, -1):
    ...:     i = 128 * 2 ** i
    ...:     print(f'\n\nSize: {i}x{i}   \t{i*i // 1024} KPx')
    ...:     im = Image.new('RGB', (i, i))
    ...:     print('\tnumpy.array()')
    ...:     %timeit n = numpy.array(im)
    ...:     print('\tnumpy.asarray()')
    ...:     %timeit n = numpy.asarray(im)
    ...:     print('\tto_numpy()')
    ...:     %timeit n = to_numpy(im)
    ...:     im = None
    ...: 

Результаты:

Размер	numpy.array()	numpy.asarray()	to_numpy()	Ускорение
8192x8192	995 мс	683 мс	378 мс	2,63x
4096x4096	257	179	101	2,54x
2048x2048	24,5	13,4	10,5	2,33x
1024x1024	4,84	3,45	2,74	1,77x
512x512	1,34	1,05	0,75	1,79x
256x256	0,26	0,2	0,18	1,44x

Итого, получилось избавиться от лишней аллокации памяти, ускорить работу от 1,5 до 2,5 раз, попутно немного разобраться как NumPy работает с памятью.