Перегон картинок из 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     ...: 

Результаты:

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