Удаление фона с изображения с помощью Python и OpenCV
Появилась давеча задачка определить цвета на изображениях. Если точнее, то нужно было вычислить цвета товара на фотографиях, для применения в фильтрах. Товаром, в моем случае, оказались ковры.
«Ну что может быть проще?» — сказал я — «Пара дней и готово.»
Как же я был наивен тогда, две недели назад…
В качестве языка был выбран Python. Именно на нем больше всего примеров по Machine Learning (ML) и Computer Vision (CV). Так же я пользовался Python, когда делал тепловую карту стоимости коммерческой недвижимости Московской области.
В качестве CV библиотеки была выбрана OpenCV v4.1.1.
Собственно, об определении цветов, как ясно из заголовка, не в этой статье. Все потому, что нюансов там множество. Так же несколько экспериментов и визуальный осмотр некоторых фото показали, что первым делом с фотографий нужно удалить фон. Благо библиотека opencv уже имеет весь необходимый функционал, как мне казалось.
Код всех тестов собран тут.
Тут нужно добавить что опыта работы с питоном у меня крайне мало. А уж с opencv и подавно. А уж про математику за всем этим стоящую я вообще не говорю.
Поэтому большая часть тестов проводилась так:
- Найти код;
- найти доки по используемым функциям;
- поиграться с параметрами;
- проклясть все;
- снова поиграться с параметрами;
- перейти к шагу 1.
Большая часть тестов будет импортировать одни и те же пакеты, так что приведу их здесь, чтобы не тратить на них место дальше.
import numpy as np
import cv2
from matplotlib import pyplot as plt
Изображения использованные в тестах
 |
 |
 |
 |  |
 |
 |  |
 |
 |
 |
 |
Тест 1
О цвете скриншотов в этом тесте
Opencv использует цветовую модель BGR, а pyplot рисует в RGB.
Я забыл конвертировать в RGB (cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) перед выводом на plot. Именно поэтому цвета на скриншотах, как бы, с синим отливом, а маски желто-фиолетовые, а не чернобелые.
В описании дальнейших тестов я это учел.
Но маски конвертировать не буду. Желто-фиолетовые получаются более наглядными, чем черно-белые.
Используются методы cv2.threshold, cv2.morphologyEx и cv2.dilate().
Как видно на результатах приведенных далее, тест не удался. Фон (background), нормально определился только на втором изображении. А вот передний план (foreground) не определился ни где.
Здесь необходимо сделать одну оговорку: поскольку это был первый тест, «игрался» я с ним не так долго.
Пройдемся по коду.
Загружаем изображение и обесцвечиваем
img = cv2.imread('img_test/12.jpg')
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
Определение фона
# Вычисляем маску фона
_, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
# Убираем шум
kernel = np.ones((2, 2), np.uint8)
opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=50)
# создаем маску фона
sure_bg = cv2.dilate(opening, kernel, iterations=20)
На данном этапе определение фона срабатывает отлично. Проблемы возникают при удалении шума. Если поставить высокое кол-во итераций, на данном примере фон, практически исчезает, если низкое, остается много шума. Можно было бы вручную подкрутить и добиться нужного результата, но это сработало бы только для данного изображения, что неприемлемо.
На создание маски фона не очень влияет количество итераций, но, как можно увидеть на изображении 6, его увеличение дает несколько лучший результат
 |
 |
|---|---|
| Первый тест. Изображение 5 | Первый тест. Изображение 6 |
Определение переднего плана
dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 5)
_, sure_fg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0)
Здесь результат мало зависит от рассчитанной до этого маски фона opening, хотя она и задействована в вычислении. Так же изменение параметров методов практически не влияет на результат. Скорее всего причина в конкретном изображении, что играет не в пользу данного метода.
 |
|---|
| Первый тест. Изображение 7 |
Подробно останавливаться на остальном коде, мне кажется, смысла особого нет. По сути там просто «соединяются» подсчитанные маски фона и переднего плана и рисуются маркеры на изображении. Результаты работы приведены в начале.
# Вычисляем "неопределенный регион" из вычисленных фона и переднего плана
# В большинстве случаев это просто фон, т.к. передний план считается плохо.
sure_fg = np.uint8(sure_fg)
unknown = cv2.subtract(sure_bg, sure_fg)
# Размечаем маркеры. Они будут использованы для вычисления фона алгоритмом watershed
_, markers = cv2.connectedComponents(sure_fg)
# Добавляем единичку ко всем значениям, чтобы они были больше 0. Он отвечает за "неизвестную" область.
markers = markers + 1
# Добавляем "неизвестную" область на маркеры
markers[unknown == 255] = 0
# собственно считаем
markers = cv2.watershed(img, markers)
# и рисуем наши маркеры на изображении
img[markers == -1] = [255, 0, 0]
Результаты теста №1
| Первый тест. Изображение 1. Результат | Первый тест. Изображение 2. Результат |
|---|---|
 |
 |
| Первый тест. Изображение 3. Результат | Первый тест. Изображение 4. Результат |
 |
 |
Тест 2
Данный тест оказался самым удачным в моем конкретном случае, поэтому о нем напишу в конце статьи.
Тест 3
Самый быстрый в моем случае тест. Даже в коде копаться не пришлось. Написан с использованием PyTorch.
Все дело в том, что данный способ использует машинное обучение (что мне не подходило) и уже обученную модель models.segmentation.deeplabv3_resnet101 (ссылка на информацию о модели).
Находить ковры (да и не только их) эту модель, очевидно, никто не обучал.
Она отлично показала себя на определении авто и довольно плохо на мебели. На всех остальных моих тестовых изображениях, результат был плачевный, что было ожидаемо.
 |
 |
| Тест 3. Изображение 1 | Тест 3. Изображение 2 |
Остальные результаты
 |
 |
| Тест 3. Изображение 3 | Тест 3. Изображение 4 |
 |
 |
| Тест 3. Изображение 5 | Тест 3. Изображение 6 |
 |
 |
| Тест 3. Изображение 7 | Тест 3. Изображение 8 |
 |
 |
| Тест 3. Изображение 9 | Тест 3. Изображение 10 |
 |
 |
| Тест 3. Изображение 11 | Тест 3. Изображение 12 |
Тест 4. Поиск форм
В процессе экспериментов возникла теория: что если не искать фон и пытаться его удалить, а пытаться найти на изображении самый большой по площади контур и принять его за наш объект.
Для этого была написана функция process(), которая получает на входе изображение, а на выходе отдает threshold (для визуализации с чем работали) и результат в виде изображения с нарисованном на нем самым большим найденным контуром.
def process(img):
work_img = img.copy()
# считаем маску
_, threshold = cv2.threshold(work_img, 30, 200, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
# ищем на маске контуры
contours, hierarchy = cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
max_area = 0
max_contour = []
# ищем самый большой из найденных контуров
for cnt in contours:
approx = cv2.approxPolyDP(cnt, 0.1 * cv2.arcLength(cnt, True), True)
cnt_area = cv2.contourArea(cnt)
# обязательно проверяем кол-во точек.
# Иначе, часто, получается что самая большая фигура -
# это треугольник через всю диагональ изображения
if len(approx) > 3 and cnt_area > max_area:
max_area = cnt_area
max_contour = approx
# рисуем самый большой контур на изображении
cv2.drawContours(work_img, [max_contour], 0, (255, 0, 0), 2)
return work_img, threshold
Результат оказался удручающим. Он очень сильно зависит от исходного изображения, поскольку используется методы cv2.threshold() и cv2.approxPolyDP() которые нужно подкручивать вручную.
Так же удивительным оказался результат даже на очевидной, для человека, полученной маске.
Помимо прочего я, к своему стыду, так и не разобрался с методом cv2.approxPolyDP(). В этом тоже может быть проблема плохой работы скрипта. На скриншотах результатов это наглядно видно (рис. Тест 4. Изображение 1): на маске присутствует овальная область, которая, судя по всему и превращается в четырехугольник.
 |
 |
| Тест 4. Изображение 1 | Тест 4. Изображение 2 |
 |
 |
| Тест 4. Изображение 3 | Тест 4. Изображение 4 |
Тем не менее, польза от эксперимента все же есть. К примеру, если изменить коэффициент с 0.1 до 0.001 при расчете параметра точности аппроксимации (приближения, усреднения) epsilon в методе cv2.approxPolyDP(), то неплохо находится очертание слона на изображении.
approx = cv2.approxPolyDP(cnt, 0.001 * cv2.arcLength(cnt, True), True)
 |
| Тест 4. Изображение 5 |
Финал
Код теста 2 на GitHub
Код итогового cli скрипта
В качестве финального варианта был выбран Тест 2. В нем используется метод cv2.grabCut().
Именно этот вариант дал наилучшие результаты на всех тестовых изображениях, а также на 2000+ других. К сожалению, этот же метод оказался самым ресурсоемким. На обработку 2000+ изображений было потрачено порядка 6 часов.
Алгоритм работы базируется на 2 допущениях:
1. Все в рамках заданного отступа от края изображения является фоном
2. Все в рамках заданной области в центре изображения является передним планом.
Несколько примеров работы, и перейдем к коду:
 |
 |
| Финал. Изображение 1 | Финал. Изображение 2 |
 |
|
| Финал. Изображение 3 |
# инициализация и загрузка изображения
# до этого размера по минимальной стороне
# мы будем уменьшать изображение
smallestSideSize = 500
# размер прямоугольника для фона (отступ от внешней стороны изображения)
# в процентах от размера изображения
mainRectSize = .04 # 4%
# размер прямоугольника для переднего плана
# будет размещен в центре изображения
fgSize = .15 # 15%
В коде, в закомментированном виде, я оставил вызов функции quantify_colors(). Она уменьшает число цветов на изображении методом k-средних (k-means) и может пригодиться. В некоторых случаях, это улучшает результат, но мене метод не подошел. Отчасти из-за качества результата, отчасти из-за сильно возрастающей нагрузки ввиду использования метода k-средних.
def quantify_colors(img, k=32, attempts=5):
# преобразуем пиксели во float
float_pixels = np.float32(img)
# "Вытягиваем 2-мерный массив в 1-мерный"
float_pixels = float_pixels.reshape((float_pixels.shape[0] * float_pixels.shape[1], 3))
# применяем метода k-средних
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
ret, label, center = cv2.kmeans(float_pixels, k, None, criteria, attempts, cv2.KMEANS_RANDOM_CENTERS)
# Преобразуем пиксели обратно в uint
center = np.uint8(center)
# "закрашиваем" и восстанавливаем 2-мерность массива
ret = center[label.flatten()]
ret = ret.reshape(img.shape)
return ret
Далее нужно инициализировать начальную маску и исходный прямоугольник фона.
Сразу нужно отметить, что в качестве сигналов на маске grabCut() понимает следующие значения:
— 0 (cv2.GC_BGD) — это точно фон
— 1 (cv2.GC_FGD) — это точно передний план
— 2 (cv2.GC_PR_BGD) — возможно это фон
— 3 (cv2.GC_PR_FGD) — возможно это передний план
# Создаем пустую маску
mask = np.zeros(img_small.shape[:2], np.uint8)
# Создаем прямоугольник фона
bg_w = round(new_w * mainRectSize)
bg_h = round(new_h * mainRectSize)
bg_rect = (bg_w, bg_h, new_w - bg_w, new_h - bg_h)
# Создаем прямоугольник переднего плана
fg_w = round(new_w * (1 - fgSize) / 2)
fg_h = round(new_h * (1 - fgSize) / 2)
fg_rect = (fg_w, fg_h, new_w - fg_w, new_h - fg_h)
# Рисуем на маске закрашенный прямоугольник с сигналом что это передний план (объект)
cv2.rectangle(mask, fg_rect[:2], fg_rect[2:4], color=cv2.GC_FGD, thickness=-1)
Далее мы инициализируем переменные для хранения моделей фона и переднего плана. И дважды вызываем метод cv2.grabCut().
В первый раз для расчета маски в режиме cv2.GC_INIT_WITH_RECT. Я так понял (возможно ошибочно), что пропустить его нельзя, так как без него сразу вываливалось исключение. В этом режиме переданная маска используется только для вывода. На ней будет дорисован результат расчетов сделанных на основе примененной к изображению рамки фона bg_rect.
Затем на полученной маске дорисовывается рамка «возможно фон», на том же месте, где и рамка «точно фон», но чуть большей толщины.
После чего метод cv2.grabCut() вызывается второй раз, но теперь в режиме cv2.GC_INIT_WITH_MASK. В этом режиме к изображению применяется составленная нами на предыдущих шагах маска. И на нее же рисуется итоговая маска.
# инициализируем массивы для моделей
bgdModel1 = np.zeros((1, 65), np.float64)
fgdModel1 = np.zeros((1, 65), np.float64)
# Первый запуск
cv2.grabCut(img_small, mask, bg_rect, bgdModel1, fgdModel1, 3, cv2.GC_INIT_WITH_RECT)
# Рисуем рамку "возможно фон"
cv2.rectangle(mask, bg_rect[:2], bg_rect[2:4], color=cv2.GC_PR_BGD, thickness=bg_w * 3)
# Второй запуск
cv2.grabCut(img_small, mask, bg_rect, bgdModel1, fgdModel1, 10, cv2.GC_INIT_WITH_MASK)
# строим итоговую маску по которой будем вырезать фон с изображения
# значения 1 и 3 ("передний план" и "возможно передний план") будем оставлять на итоговом
# изображении. Все остальное будет удалено
mask_result = np.where((mask == 1) + (mask == 3), 255, 0).astype('uint8')
Следующие несколько строк кода не обязательны, но в моем случае оказались полезны.
Это сопоставление размера вырезаемой области с той что останется. Если будет вырезано больше, чем останется с некоторым коэффициентом (в моем случае 1.6), предполагаем, что фона на изображении просто нет и вырезать ничего не нужно. Это помогло «спасти» некоторое количество объектов занимающих все изображение целиком.
unique, counts = np.unique(mask_result, return_counts=True)
mask_dict = dict(zip(unique, counts))
if mask_dict[0] > mask_dict[255] * 1.6:
mask_result = np.where((mask == 0) + (mask != 1) + (mask != 3), 255, 0).astype('uint8')
Финальный шаг — применяем маску на изображение
# "выжигаем" маску на изображение
masked = cv2.bitwise_and(img_small, img_small, mask=mask_result)
# для себя я выбрал синий фон в качестве хрома-кея
masked[mask_result < 2] = [0, 0, 255]
Несколько рандомных результатов
 |
 |
| Финал. Изображение 4 | Финал. Изображение 5 |
 |
 |
| Финал. Изображение 6 | Финал. Изображение 7 |
 |
 |
| Финал. Изображение 8 | Финал. Изображение 9 |
 |
 |
| Финал. Изображение 10 | Финал. Изображение 11 |
 |
 |
| Финал. Изображение 12 | Финал. Изображение 13 |
 |
 |
| Финал. Изображение 14 | Финал. Изображение 15 |
На этом все. Результат, конечно, не идеальный, но меня удовлетворяет. Следующий шаг — разметить данные изображения по цветам. Что, как оказалось, не такая уж и простая задачка.
Но об этом уже в другой раз.