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

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

Да, это кажется безумием, но более привычные методы преобразования картинок работают в 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, с равной вероятностью вы нагуглите один из них:

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

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

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

In [1]: from PIL import ImageIn [2]: import numpyIn [3]: im = Image.open('./canyon.jpg').resize((4096, 4096))In [4]: n = numpy.asarray(im)In [5]: n[:, :, 0] = 255ValueError: assignment destination is read-onlyIn [6]: n.flagsOut[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] = 255In [9]: n.flagsOut[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.dataOut[17]: <memory at 0x7f78854d68>In [18]: data.data[0] = 255NotImplementedError: sub-views are not implementedIn [19]: data.data.shapeOut[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]: TrueIn [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]: TrueIn [28]: n.flagsOut[28]:   C_CONTIGUOUS : True  F_CONTIGUOUS : False  OWNDATA : True  WRITEABLE : True  ALIGNED : True  WRITEBACKIFCOPY : False  UPDATEIFCOPY : FalseIn [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 работает с памятью.