Перегон картинок из 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. В конце концов, где ещё может понадобиться компьютерное зрение, как не на Малинке :-)
На случай, если вы не знаете как настроить окружение
Как работает преобразование в NumPy
Существует два общепринятых способа конвертировать изображение Pillow в NumPy, с равной вероятностью вы нагуглите один из них:
numpy.array(im)
— делает копию из изображения в массив NumPy.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 работает с памятью.