Перегон картинок из 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, с равной вероятностью вы нагуглите один из них:

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 работает с памятью.