Pull to refresh

История оптимизации alpha_composite в Pillow 2.0

Reading time 7 min
Views 5.9K
Недавно вышла вторая версия питоновской библиотеки для работы с изображениями Pillow. Как многие знают, это форк хорошо известной библиотеки PIL, которая, несмотря на свой солидный возраст, до недавнего времени оставалась самым вменяемым способом работы с изображениями в Питоне. Авторы Pillow наконец-то решили не только поддерживать старый код, но и добавлять новые возможности. И одной из этих возможностей стала функция alpha_composite().

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

С тех пор я уже смог написать для своих нужд более оптимальную реализацию, чем приведена в конце той статьи. И оказалось, что эта реализация быстрее alpha_composite() из новой версии Pillow, написанной на Си. Конечно, мне это польстило, но я все-таки решил попытаться улучшить реализацию из Pillow.

Для начала нужен тестовый стенд.

bench.py
from PIL import Image
from timeit import repeat
from image import paste_composite

im1 = Image.open('in1.png')
im1.load()

im2 = Image.open('in2.png')
im2.load()

print repeat(lambda: paste_composite(im1.copy(), im2), number=100)
print repeat(lambda: Image.alpha_composite(im1, im2), number=100)

out1 = im1.copy()
paste_composite(out1, im2)
out1.save('out1.png')

out2 = Image.alpha_composite(im1, im2)
out2.save('out2.png')

Скрипт берет лежащие в текущей папке файлы in1.png и in2.png и смешивает их сто раз сначала функцией paste_composite(), потом alpha_composite(). Я взял файлы из предыдущей статьи (раз, два). Время работы paste_composite() у меня составило в среднем 235мс. Время работы оригинальной alpha_composite() составило 400мс.

Почему alpha_composite() такая медленная, долго гадать не пришлось. Глядя на реализацию, видно, что функция производит все вычисления на числах с плавающей точкой. При работе с графикой числа с плавающей точкой позволяют достаточно просто описать алгоритм, но производительность у них всегда меньше, чем у целых чисел. Первое, что я сделал, перевел все на работу с целыми числами:

typedef struct {
    UINT8 r;
    UINT8 g;
    UINT8 b;
    UINT8 a;
} rgba8;

Imaging
ImagingAlphaComposite(Imaging imDst, Imaging imSrc)
{
    Imaging imOut = ImagingNew(imDst->mode, imDst->xsize, imDst->ysize);

    for (int y = 0; y < imDst->ysize; y++) {
        rgba8* dst = (rgba8*) imDst->image[y];
        rgba8* src = (rgba8*) imSrc->image[y];
        rgba8* out = (rgba8*) imOut->image[y];

        for (int x = 0; x < imDst->xsize; x ++) {
            if (src->a == 0) {
                *out = *dst;
            } else {
                UINT16 blend = dst->a * (255 - src->a);
                UINT8 outa = src->a + blend / 255;
                UINT8 coef1 = src->a * 255 / outa;
                UINT8 coef2 = blend / outa;

                out->r = (src->r * coef1 + dst->r * coef2) / 255;
                out->g = (src->g * coef1 + dst->g * coef2) / 255;
                out->b = (src->b * coef1 + dst->b * coef2) / 255;
                out->a = outa;
            }
            dst++; src++; out++;
        }
    }
    return imOut;
}

И сразу же получил прирост в 5,5 раз. Тест отрабатывал за 75 мс.

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

diff.py
#!/usr/bin/env python
import sys
from PIL import Image, ImageMath

def chanel_diff(c1, c2):
    c = ImageMath.eval('127 + c1 - c2', c1=c1, c2=c2).convert('L')
    return c.point(lambda c: c if 126 <= c <= 128 else 127 + (c - 127) * 10)

im1 = Image.open(sys.argv[1])
im2 = Image.open(sys.argv[2])

diff = map(chanel_diff, im1.split(), im2.split())

Image.merge('RGB', diff[:-1]).save('diff.png', optimie=True)
diff[-1].convert('RGB').save('diff.alpha.png', optimie=True)

На вход скрипт принимает 2 имени файла с картинками. Между каждой парой каналов этих картинок находятся отличия: сначала просто ищется разница и складывается с серым фоном (значение 127), чтобы можно было увидеть отклонения в обе стороны. Затем все отклонения, сильнее чем в одно значение, усиливаются в 10 раз, чтобы они были заметны визуально. Результат сравнения записывается в 2 файла: rgb каналы в diff.png, альфа-канал в виде оттенков серого в diff.alpha.png.

Оказалось, что разница между результатом alpha_composite() и фотошопом достаточно существенная:


Еще раз внимательно посмотрев на код, я понял, что забыл про округление: целые числа делятся с отбрасыванием остатка. Если мы хотим получить округленный результат, нужно к делимому прибавить половину делителя:

UINT16 blend = dst->a * (255 - src->a);
UINT8 outa = src->a + (blend + 127) / 255;
UINT8 coef1 = src->a * 255 / outa;
UINT8 coef2 = blend / outa;

out->r = (src->r * coef1 + dst->r * coef2 + 127) / 255;
out->g = (src->g * coef1 + dst->g * coef2 + 127) / 255;
out->b = (src->b * coef1 + dst->b * coef2 + 127) / 255;
out->a = outa;

Время работы увеличилось до 94 мс. Но и отличий от референса стало намного меньше. Впрочем, они не пропали совсем.

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

Очевидно, что полученные неточности — результат недостаточной точности коэффициентов при умножении. Можно хранить дополнительные значащие биты для каждой переменной, что поднимает точность при делении. Переменная blend вычисляется без деления, с ней ничего делать не нужно. Переменная outa получается делением blend на 255. При расчете сдвигаем значения blend и src->a на 4 бита вверх. В результате в outa будет 12 значащих бит. В обоих коэффициентах происходит деление на outa, которая теперь больше на 4 бита. Плюс сами коэффициенты хотелось бы получить на 4 бита больше. В результате делимое сдвигаем уже на 8 бит.

UINT16 blend = dst->a * (255 - src->a);  // 16 bit max
UINT16 outa = (src->a << 4) + ((blend << 4) + 127) / 255;  // 12
UINT16 coef1 = ((src->a * 255) << 8) / outa;  // 12
UINT16 coef2 = (blend << 8) / outa;  // 12

out->r = ((src->r * coef1 + dst->r * coef2 + 0x7ff) / 255) >> 4;
out->g = ((src->g * coef1 + dst->g * coef2 + 0x7ff) / 255) >> 4;
out->b = ((src->b * coef1 + dst->b * coef2 + 0x7ff) / 255) >> 4;
out->a = (outa + 0x7) >> 4;

Само собой, на эти же 4 бита нужно сдвинуть результат всех вычислений, прежде чем помещать его в пиксели изображения. Такой вариант работает еще чуть медленнее: 108 мс, но зато уже практически не имеет пикселей, отличающихся более чем на одно значение от эталона. Результат по качеству достигнут, можно опять подумать о производительности.

Очевидно, что такой код не самый оптимальный с точки зрения процессора. Деление — одна из самых сложных задач, а здесь на каждый пиксель выполняется целых 6 делений. Оказалось, что избавиться от большинства из них достаточно просто. В той же PIL в файле Paste.c описан способ быстрого деления с округлением:
a + 127 / 255 ≈ ((a + 128) + ((a + 128) >> 8)) >> 8

В результате одно деление заменяется на 2 сдвига и одно сложение, что быстрее на большинстве процессоров. Сгруппировав один из сдвигов с уже имеющимся сдвигом на 4, я получил такой код:

UINT16 blend = dst->a * (255 - src->a);
UINT16 outa = (src->a << 4) + (((blend << 4) + (blend >> 4) + 0x80) >> 8);
UINT16 coef1 = (((src->a << 8) - src->a) << 8) / outa;  // 12
UINT16 coef2 = (blend << 8) / outa;  // 12

UINT32 tmpr = src->r * coef1 + dst->r * coef2 + 0x800;
out->r = ((tmpr >> 8) + tmpr) >> 12;
UINT32 tmpg = src->g * coef1 + dst->g * coef2 + 0x800;
out->g = ((tmpg >> 8) + tmpg) >> 12;
UINT32 tmpb = src->b * coef1 + dst->b * coef2 + 0x800;
out->b = ((tmpb >> 8) + tmpb) >> 12;
out->a = (outa + 0x7) >> 4;

Он выполняется уже за 65 мс и не вносит изменений в результат по сравнению с предыдущим вариантом.

Итого скорость работы была улучшена в 6 с небольшим раз, в репозиторий был отправлен пул-реквест. Можно было бы попытаться добиться большего сходства с эталоном. Например, мне удалось идеально точно повторить альфа-канал, генерируемый фотошопом, но при этом в пиксели вносилось больше искажений.

Большой апдейт

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

diff.py
#!/usr/bin/env python
import sys
from PIL import Image, ImageMath

def chanel_diff(c1, c2):
    return ImageMath.eval('127 + c1 - c2', c1=c1, c2=c2).convert('L')

def highlight(c):
    return c.point(lambda c: c if 126 <= c <= 128 else 127 + (c - 127) * 10)

im1 = Image.open(sys.argv[1])
im2 = Image.open(sys.argv[2])

diff = map(chanel_diff, im1.split(), im2.split())

if len(sys.argv) >= 4:
    highlight(Image.merge('RGB', diff[:-1])).save('%s.png' % sys.argv[3])
    highlight(diff[-1]).convert('RGB').save('%s.alpha.png' % sys.argv[3])

def stats(ch):
    return sorted((c, n) for n, c in ch.getcolors())

for ch, stat in zip(['red ', 'grn ', 'blu ', 'alp '], map(stats, diff)):
    print ch, '  '.join('{}: {:>5}'.format(c, n) for c, n in stat)

Так же я подготовил очень серьезный тесткейс из двух картинок на 16 мегапикселей, где каждое сочетание пикселей было уникальным.

make_testcase.py
def prepare_test_images(dim):
    """Plese, be careful with dim > 32. Result image is have dim ** 4 pixels
    (i.e. 1Mpx for 32 dim or 4Gpx for 256 dim).
    """
    i1 = bytearray(dim ** 4 * 2)
    i2 = bytearray(dim ** 4 * 2)
    res = 255.0 / (dim - 1)
    rangedim = range(dim)
    pos = 0
    for l1 in rangedim:
        for l2 in rangedim:
            for a1 in rangedim:
                for a2 in rangedim:
                    i1[pos] = int(res * l1)
                    i1[pos + 1] = int(res * a1)
                    i2[pos] = int(res * l2)
                    i2[pos + 1] = int(res * a2)
                    pos += 2
        print '%s of %s' % (l1, dim)

    i1 = Image.frombytes('LA', (dim ** 2, dim ** 2), bytes(i1))
    i2 = Image.frombytes('LA', (dim ** 2, dim ** 2), bytes(i2))
    return i1.convert('RGBA'), i2.convert('RGBA')

im1, im2 = prepare_test_images(63)
im1.save('im1.png')
im2.save('im2.png')

После этого я решил начать с другого конца: не пытаться увеличить точность у целочисленного алгоритма, а пытаться привести алгоритм с плавающей точкой к работе с целыми. Вот с чего я начал:

double dsta = dst->a / 255.0;
double srca = src->a / 255.0;

double blend = dsta * (1.0 - srca);
double outa = srca + blend;
double coef1 = srca / outa;
double coef2 = 1 - coef1;

double tmpr = src->r * coef1 + dst->r * coef2;
out->r = (UINT8) (tmpr + 0.5);
double tmpg = src->g * coef1 + dst->g * coef2;
out->g = (UINT8) (tmpg + 0.5);
double tmpb = src->b * coef1 + dst->b * coef2;
out->b = (UINT8) (tmpb + 0.5);
out->a = (UINT8) (outa * 255.0 + 0.5);

И вот к чему я пришел:

UINT16 blend = dst->a * (255 - src->a);
UINT16 outa255 = src->a * 255 + blend;
// There we use 7 bits for precision.
// We could use more, but we go beyond 32 bits.
UINT16 coef1 = src->a * 255 * 255 * 128 / outa255;
UINT16 coef2 = 255 * 128 - coef1;

#define SHIFTFORDIV255(a)\
    ((a >> 8) + a >> 8)

UINT32 tmpr = src->r * coef1 + dst->r * coef2 + (0x80 << 7);
out->r = SHIFTFORDIV255(tmpr) >> 7;
UINT32 tmpg = src->g * coef1 + dst->g * coef2 + (0x80 << 7);
out->g = SHIFTFORDIV255(tmpg) >> 7;
UINT32 tmpb = src->b * coef1 + dst->b * coef2 + (0x80 << 7);
out->b = SHIFTFORDIV255(tmpb) >> 7;
out->a = SHIFTFORDIV255(outa255 + 0x80);
Tags:
Hubs:
+21
Comments 11
Comments Comments 11

Articles