lgorSL 25 ноя 2021 в 02:29

Как можно взять tensorflow и смешать две картинки в одну

6 мин

5.4K

Туториал

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

Не знаю, как делались оригинальные, но я попробовал сделать нечто похожее с помощью tensorflow.

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

Преобразование Фурье - прошлый век, новые проблемы требуют инновационных решений.

Код и исходные картинки лежат на гитхабе.

Для запуска понадобится tensorflow, numpy и Pillow. Для последней версии tensorflow понадобится python 3.7 или новее.

Чтобы результат получался за пару минут при вычислении на процессоре, я уменьшил картинки до размера 512х512 пикселей. Взял их из википедии: шахматы и девушка Лена.

Шахматы

Девушка Лена

Загрузим их и переведём во float в интервале от 0 до 1:

lena = Image.open("imgs/Lenna512.png")
chess = Image.open("imgs/ChessSet512.png")
imgs = [np.array(i) / 255.0 for i in [lena, chess]]

Гамма-коррекция

Воспринимаемая глазом яркость света может отличаться на много-много порядков. В изображениях с одним байтом яркости на канал есть 256 значений яркости и в физическое количество фотонов они отображаются нелинейно. Например, от двух пикселей с яркостью 127 света будет меньше, чем от одного с яркостью 254.

Для перевода в физическое значение яркости существует гамма-коррекция: возведение яркости в степень 2.2. Получится величина, пропорциональная количеству фотонов с экрана. Обратное преобразование тоже делается просто - возведением в степень 1/2.2.

Зачем всё это нужно? Если я хочу узнать видимую разницу цветов, можно просто взять разницу значений в обычном RGB пространстве.

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

Размытие по Гауссу

При размытии картинки каждый пиксель превращается в пятнышко. На математическом языке пятнышко называется ядром, а сам процесс - свёрткой. Ядро выглядит как-то так:

$Ce^{- \frac{dx^2+dy^2}{2 \sigma^2}} = C e^{-\frac{dx^2}{2 \sigma^2}} e^{-\frac{dy^2}{2 \sigma^2}}$

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

Особенностью свёртки с этим ядром является то, что его её можно сделать сначала по одной оси, а затем по другой, так что достаточно посчитать ядро для одномерного случая:

def make_gauss_blur_kernel(size: int, sigma: float) -> np.ndarray:  
    result = np.zeros(shape=[size], dtype=float)  
    center = (size - 1) // 2  
    div = 2 * (sigma ** 2)  
    for i in range(size):  
        x2 = (center - i)**2 
        result[i] = math.exp( -x2 / div)  
    return result / np.sum(result)  
  
make_gauss_blur_kernel(size=11, sigma=2)

Математически ядро бесконечное, но мы ограничим его размеры. Например, для sigma =2 и размера ядра в 11 получится так:

array([0.00881223, 0.02714358, 0.06511406, 0.12164907, 0.17699836, 0.20056541, 0.17699836, 0.12164907, 0.06511406, 0.02714358, 0.00881223])

Tensorflow 2.0

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

Магия выглядит вот так:

with tf.GradientTape() as tape:
    # computations
    loss = ...
    
gradient = tape.gradient(loss, trainable_variables)

Это как раз то, что нам нужно.

Создаём модель

class MyModel:
    def __init__(self, img_h: int, img_w: int, gauss_kernel_size: int, gauss_sigma: float, image_source: Optional[np.ndarray] = None):
        if image_source is None:
            image_source = np.zeros(shape=(1, img_h, img_w, 3), dtype=float)
        self.trainable_image = tf.Variable(initial_value=image_source, trainable=True)
        gauss_blur_kernel = make_gauss_blur_kernel(gauss_kernel_size, gauss_sigma)
        self.gauss_kernel_x = tf.constant(gauss_blur_kernel[np.newaxis, :, np.newaxis, np.newaxis] * np.ones(shape=(1, 1, 3, 1)))
        self.gauss_kernel_y = tf.constant(gauss_blur_kernel[:, np.newaxis, np.newaxis, np.newaxis] * np.ones(shape=(1, 1, 3, 1)))
        
    ...

trainable_image - переменные. Из них состоит наша "обучаемая" картинка. Их мы будем подгонять, чтобы было похоже на нужные две картинки. Так же сделаем константы для ядер свёрток по осям х и у. Их учить не будем, они и так получены умным путём.

Для свёрток используется аж четырёхмерное ядро:

ось У изображения
ось Х изображения
количество входных каналов (3)
количество выходных каналов на каждый входной канал (1, каждый цвет переходит сам в себя)

Кроме того, и в numpy и в tensorflow есть идея броадкастинга. Например, массив с размерностями (1, 256, 512, 1) можно интерпретировать как (N, 256, 512, C). Как будто по первой и последней оси размер произвольный, а числа одни одни и те же и не зависят от координаты вдоль той оси. Броадкастинг в этих библиотеках иногда работает по-разному, функция свёртки хочет увидеть ядро размерностью именно (size_x, size_y, 3, 1) и почему-то недовольно массивом (size_x, size_y, 1, 1), так что пришлось в numpy умножить на массив единичек размерностью (1, 1, 3, 1). Если бы мы сделали разные ядра свёрток для разных цветов, то нам бы пригодилась эта размерность, но у нас всё одинаково.

Я использую поканальные свёртки (чтобы канал с цветом при размытии влиял только на себя и превращался в новый канал с тем же цветом). Такие свёртки считаются быстрее обычных. А ещё внутри нет никакого преобразования Фурье, и сложность вычисления свёртки зависит линейно от размера ядра. По этой же причине свёртка с квадратным ядром размера 15x15 будет считаться в несколько раз дольше, чем две свёртки с ядрами 1х15 и 15х1.

class MyModel:
    ...

    def run(self, img_precise: np.ndarray, img_blurred: np.ndarray, m_precise = 1, m_blurred = 1) -> Report:
        with tf.GradientTape() as tape:
            # next code will be here

Для шага обучения возьмём две картинки и коэффициенты для важности ошибок у каждой.

Обучаемые переменные могут быть любыми, хоть -1, хоть 9000. Но в качестве цветов картинки хочется получить значения в интервале от 0 до 1. Для этого применим сигмоиду: около нуля она более-менее линейно растёт, но на больших входных значениях рост замедляется и результат никогда не превысит единицу.

trainable_image01 = tf.math.sigmoid(self.trainable_image)

Для первой картинки (которая должна быть резкой) разница - просто сумма квадратов разностей яркости для каждого канала каждого пикселя. По-идее, каждый пиксель учится независимо, и шаг обучения в 0.1 будет вполне нормальным.

Вместо суммы можно было бы использовать reduce_mean(), но тогда градиенты были бы меньше на площадь картинки (512x512) и пришлось бы градиенты умножать на что-то типа 10^5.

loss_precise = tf.reduce_sum(tf.square(trainable_image01 - img_precise[np.newaxis, :, :, :]))

Для второй картинки сделаем размытие (свёртку), а потом точно так же попиксельно сравним. И не забудем про гамма-коррекцию до свёртки и обратную после:

blurred = self.gauss_blur(trainable_image01 ** 2.2) ** (1.0 / 2.2)
blurred_label = self.gauss_blur(img_blurred[np.newaxis, :, :, :] ** 2.2) ** (1.0 / 2.2)
loss_gauss = tf.reduce_sum(tf.square(blurred - blurred_label))

Домножим ошибки на коэффициенты и сложим. В зависимости от соотношения коэффициентов результат будет больше стремиться к одной картинке или другой.

loss = loss_precise * m_precise + loss_gauss * m_blurred

Вжух и получим градиенты:

gradient = tape.gradient(loss, self.trainable_image)

Для обучения есть разные оптимизаторы, но я прямо руками сделал простой градиентный спуск:

def apply_gradient(self, grad: np.ndarray, lr: float):
    self.trainable_image.assign(self.trainable_image.numpy() - lr * grad)

Вызовем функцию много-много раз. Код в статье для иллюстрации, запускабельный вариант можно найти на гитхабе.

class MyModel:
    ...
    def train(self, steps_count: int, print_loss_steps: int, lr: float, **run_kwargs) -> Report:  
        for i in range(steps_count):  
            r = self.run(**run_kwargs)  
            model.apply_gradient(r.gradient, lr)  
            if i % print_loss_steps == print_loss_steps - 1:  
                print(f"{i}: loss = {r.loss}, precise = {r.loss_precise}, gauss = {r.loss_gauss}")  
        return r

Теперь всё готово для обучения модели:

model = MyModel(512, 512, gauss_kernel_size=15, gauss_sigma=3)
r = model.train(steps_count=200, print_loss_steps=50, lr=0.3, img_precise=imgs[0], img_blurred=imgs[1], m_precise=0.1, m_blurred=1.0)  
Image.fromarray(np.uint8(r.image * 255.0))

Я попробовал разные коэффициенты (0.1, 0.3, 1.0).

Картинки тут

А зачем всё это?

Просто потому что могу. Мне нравится возможность задать произвольную функцию ошибки и обучить модель, не задумываясь над тем, как получить результат аналитически.

Если Вам кажется, что я забиваю гвозди микроскопом и tensorflow совсем не для этого, то это не так. Библиотека даёт возможность легко считать градиенты, я ей пользуюсь как хочу. Машинное обучение не обязано происходить на кластерах с топовыми GPU на датасетах терабайтных размеров.

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

Вариант с обучением на порядки медленнее Фурье и на моём ноутбуке занимает несколько минут. На мой взгляд это не страшно, так как намного больше времени я потратил на написание кода. У меня нет задачи генерировать тысячи картинок, одной-двух вполне достаточно.

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

Теги:

Хабы:

Как можно взять tensorflow и смешать две картинки в одну

Гамма-коррекция

Размытие по Гауссу

Tensorflow 2.0

Создаём модель

А зачем всё это?

Публикации

Истории

Работа

Ближайшие события