Создание мозаичной картинки / Хабр

Наверняка вы неоднократно видели в интернете такие картинки:

Я решил написать универсальный скрипт для создания подобных изображений.

Теоретическая часть

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

Здесь встаёт вопрос о том, как понять, на какое изображение из датасета нам следует заменить некоторую область. Разумеется, идеальным замощением некоторой области будет эта же область. Каждую область размером $m\times n$ можно задать $3 \times n \times m$ числами (здесь каждому пикселю соответствуют три числа — его R, G и B компоненты). Иными словами, каждая область задается трехмерным тензором. Теперь становится понятно, что нам для определения качества замощения области картинкой, при условии совпадения их размеров, нужно посчитать некоторую функцию потерь. В данной задаче можно считать MSE двух тензоров:

$MSE(x, y) = \frac{\sum\limits_{i=1}^N (x_{i} - y_{i}) ^ 2}{N}$

Здесь $inline$ — количество признаков, в нашем случае $3 \times n \times m$ .

Однако эта формула малоприменима к реальным случаям. Дело в том, что когда датасет довольно большой, а области, на которые поделено исходное изображение, довольно маленькие, придется делать непозволительно много действий, а именно каждое изображение из датасета сжимать до размеров области и считать MSE по $3 \times n \times m$ характеристикам. Точнее говоря, в данной формуле плохо то, что мы вынуждены сжимать абсолютно каждое изображение для сравнения, причем не один раз, а число, равное количеству областей, на которые поделена исходная картинка.

Я предлагаю следующее решение проблемы: мы немного пожертвуем качеством и теперь каждую картинку из датасета будем характеризовать только 3 числами: средними RGB по изображению. Конечно, отсюда вытекает несколько проблем: во-первых, теперь идеальным замощением области является не только лишь она сама, а, например, она же, но перевернутая (очевидно, что это замощение хуже первого), во-вторых, после подсчета среднего цвета мы можем получить такие R, G и B, что на изображении даже не будет пикселя с такими компонентами (проще говоря, сложно сказать, что наш глаз воспринимает изображение как смесь всех его цветов). Однако лучшего способа я не придумал.

Получается, что теперь нам остается лишь один раз провести расчет средних RGB для картинок из датасета, а потом пользоваться полученной информацией.

Резюмируя вышенаписанное, получаем, что нам теперь необходимо некоторой области подобрать наиболее близкий ей RGB пиксель из набора, а затем замостить область тем изображением из датасета, которому принадлежат такие найденные средние RGB. Для сравнения области и пикселя поступим аналогичным образом: область преобразуем в три числа и найдем наиболее близкие средние RGB. Получается, что нам остается лишь по известным $inline$ найти в наборе такие $R_{i}, G_{i}, B_{i}$ , что Евклидово расстояние между этими двумя точками в трехмерном пространстве будет минимально:

$\sqrt{(R - R_{i}) ^ 2 + (G - G_{i}) ^ 2 + (B - B_{i}) ^ 2} = min$

Предобработка датасета

Вы можете собрать свой собственный датасет картинок. Я использовал слияние датасетов с изображениями котиков и собачек.

Как я писал выше, мы можем единожды посчитать средние RGB показатели для изображений из датасета и просто их сохранить. Что мы и делаем:

import os
import cv2
import numpy as np
import pickle

items = {}

# cv2 по умолчанию открывает картинки в BGR, а не RGB, поэтому будем их переводить

for path in os.listdir('dogs_images_dataset'):  # у датасета с собаками есть подпапки, поэтому два цикла
    for file in os.listdir(os.path.join('dogs_images_dataset', path)):
        file1 = os.path.join('dogs_images_dataset', path + '/' + file)
        img = np.array(cv2.cvtColor(cv2.imread(file1), cv2.COLOR_BGR2RGB))
        r = round(img[:, :, 0].mean())
        g = round(img[:, :, 1].mean())
        b = round(img[:, :, 2].mean())
        items[file1] = (r, g, b,)

for file in os.listdir('cats_images_dataset'):  # у датасета с кошками подпапок немного, поэтому все изображения можно вручную сохранить в одну директорию
    file1 = os.path.join('cats_images_dataset', file)
    img = np.array(cv2.cvtColor(cv2.imread(file1), cv2.COLOR_BGR2RGB))
    r = round(img[:, :, 0].mean())
    g = round(img[:, :, 1].mean())
    b = round(img[:, :, 2].mean())
    items[file1] = (r, g, b,)

with open('data.pickle', 'wb') as f:
    pickle.dump(items, f)

Этот скрипт будет выполнятся сравнительно долго, после чего необходимая нам информация будет сохранена в файл data.pickle.

Создание мозаики

Наконец, перейдем к созданию мозаики. Вначале напишем необходимые import-ы, а также объявим несколько констант:

import os
import cv2
import pickle
import numpy as np
from math import sqrt

PATH_TO_PICTURE = ''  # путь к директории с исходным изображением
PICTURE = 'picture.png'  # имя файла исходного изображения

VERTICAL_SECTION_SIZE = 7  # размер области по горизонтали в пикселях
HORIZONTAL_SECTION_SIZE = 7  # размер области по вертикали в пикселях

Достаём из файла сохраненные данные:

with open('data.pickle', 'rb') as f:
    items = pickle.load(f)

Описываем функцию потерь:

def lost_function(r_segm, g_segm, b_segm, arg):
    r, g, b = arg[1]
    return sqrt((r - r_segm) ** 2 + (g - g_segm) ** 2 + (b - b_segm) ** 2)

Открываем исходное изображение:

file = os.path.join(PATH_TO_PICTURE, PICTURE)
img = np.array(cv2.cvtColor(cv2.imread(file), cv2.COLOR_BGR2RGB))
size = img.shape
x, y = size[0], size[1]

Теперь обратим внимание, что замощение возможно тогда и только тогда, когда $(x\_orig \space \vdots \space x) \space \wedge \space (y\_orig \space \vdots \space y)$ , где $x\_orig, y\_orig$ — размеры исходного изображения, а $inline$ — размеры области замощения. Разумеется, приведенное условие выполняется не всегда. Поэтому мы обрежем исходное изображение до подходящих размеров, вычтя из размеров изображения их остатки от деления на размеры области:

img = cv2.resize(img, (y - (y % VERTICAL_SECTION_SIZE), x - (x % HORIZONTAL_SECTION_SIZE)))
size = img.shape
x, y = size[0], size[1]

Теперь перейдем непосредственно к замощению:

for i in range(x // HORIZONT
AL_SECTION_SIZE):
    for j in range(y // VERTICAL_SECTION_SIZE):
        sect = img[i * HORIZONTAL_SECTION_SIZE:(i + 1) * HORIZONTAL_SECTION_SIZE,
                j * VERTICAL_SECTION_SIZE:(j + 1) * VERTICAL_SECTION_SIZE]
        r_mean, g_mean, b_mean = sect[:, :, 0].mean(), sect[:, :, 1].mean(), sect[:, :, 2].mean()

Здесь в предпоследней строчке выбирается нужная область картинки, а в последней строчке считаются её средние RGB составляющие.

Сейчас рассмотрим одну из важнейших строк:

current = sorted(items.items(), key=lambda argument: lost_function(r_mean, g_mean, b_mean, argument))[0]

Данная строка сортирует все изображения датасета по возрастанию по значению функции потерь для них и достает argmin.

Теперь нам остается только обрезать изображение и заменить область на него:

resized = cv2.resize(cv2.cvtColor(cv2.imread(current[0]), cv2.COLOR_BGR2RGB),
                             (VERTICAL_SECTION_SIZE, HORIZONTAL_SECTION_SIZE,))
        img[i * HORIZONTAL_SECTION_SIZE:(i + 1) * HORIZONTAL_SECTION_SIZE,
        j * VERTICAL_SECTION_SIZE:(j + 1) * VERTICAL_SECTION_SIZE] = resized

Ну и наконец выведем получившуюся картинку на экран:

img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imshow('ImageWindow', img)
cv2.waitKey(0)

Еще немного про функцию потерь

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

$|\Delta R| + |\Delta G| + |\Delta B| \\ \sqrt{\Delta R ^ 2 + \Delta G ^ 2 + \Delta B ^ 2} \\ \sqrt{0.2126 \Delta R ^ 2 + 0.7152 \Delta G ^ 2 + 0.0722 \Delta B ^ 2} \\ \sqrt{0.2126^2 \Delta R ^ 2 + 0.7152^2 \Delta G ^ 2 + 0.0722^2 \Delta B ^ 2}$