{"id":14291,"url":"\/distributions\/14291\/click?bit=1&hash=257d5375fbb462be671b713a7a4184bd5d4f9c6ce46e0d204104db0e88eadadd","hash":"257d5375fbb462be671b713a7a4184bd5d4f9c6ce46e0d204104db0e88eadadd","title":"\u0420\u0435\u043a\u043b\u0430\u043c\u0430 \u043d\u0430 Ozon \u0434\u043b\u044f \u0442\u0435\u0445, \u043a\u0442\u043e \u043d\u0438\u0447\u0435\u0433\u043e \u0442\u0430\u043c \u043d\u0435 \u043f\u0440\u043e\u0434\u0430\u0451\u0442","buttonText":"","imageUuid":""}
Vitaliy Hramchenko

Как удалить фон за человеком на изображении с помощью проектов с открытым исходным кодом

Перевод моей статьи "How to cut a person out of a picture using python" на Medium.

В этом примере мы рассмотрим как сделать конвейер для высококачественного автоматического удаления фона вокруг человека с помощью ИИ.

Для проекта потребуется python с установленными пакетами pytorch 1.6, torchvision и opencv, git.

Видеоинструкцию по установке этих пакетов можно посмотреть здесь.

Конвейер включает следующие этапы:

  1. Выделение фоновых областей.
  2. Фильтрация.
  3. Финальная очистка.

Выделение фоновых областей

Для выделения фоновых областей мы будем использовать нейронную сеть DeeplabV3(импорт библиотек и код загрузки модели приведен ниже) из библиотеки PyTorch, которая присваивает каждому пикселю изображения определенный класс.

import numpy as np import cv2 import matplotlib.pyplot as plt import torch from torchvision.models.segmentation import deeplabv3_resnet101 from torchvision import transforms def make_deeplab(device): deeplab = deeplabv3_resnet101(pretrained=True).to(device) deeplab.eval() return deeplab device = torch.device("cpu") deeplab = make_deeplab(device)

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

k = min(1.0, 1024/max(img_orig.shape[0], img_orig.shape[1])) img = cv2.resize(img_orig, None, fx=k, fy=k, interpolation=cv2.INTER_LANCZOS4)

Далее приведена процедура выделения фона на изображении. Deeplab работает с нормализованными данными, поэтому перед использованием модели к изображению применяется преобразование deeplab_preprocess. Передний план определяется как набор областей с классом №15 (изображение человека).

deeplab_preprocess = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) def apply_deeplab(deeplab, img, device): input_tensor = deeplab_preprocess(img) input_batch = input_tensor.unsqueeze(0) with torch.no_grad(): output = deeplab(input_batch.to(device))['out'][0] output_predictions = output.argmax(0).cpu().numpy() return (output_predictions == 15) mask = apply_deeplab(deeplab, img, device)

В итоге получается маска как на рисунке.

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

Мягкое выделение переднего плана или матирование изображения

Image Matting — это процесс точного выделения переднего плана и фона на изображениях и видео. Если при сегментации для каждого пикселя изображения определяется единственная метка — класс объекта, которому он принадлежит (в нашем случае, передний план/фон), то при матировании изображения некоторые пиксели могут иметь сразу несколько меток с разными долями, так называемые частичные или смешанные пиксели. Чтобы полностью отделить передний план от фона на изображении, производится точная оценка Альфа-значений для смешанных пикселей.

В нашем проекте мы будем использовать код F, B, Alpha Matting из репозитория Matting Marco Forte, François Pitié.

Скачаем в папку проекта код из репозитория:

git clone https://github.com/MarcoForte/FBA_Matting.git

и файл с весами модели из официального Google Drive:

FBA.pth
FBA.pth
drive.google.com

Импортируем исходные коды FBA_Matting в скрипт обработки изображения

import sys sys.path.append("./FBA_Matting") from demo import np_to_torch, pred, scale_input from dataloader import read_image, read_trimap from networks.models import build_model

Сформируем конфигурацию и загрузим модель FBA

class Args: encoder = 'resnet50_GN_WS' decoder = 'fba_decoder' weights = 'FBA.pth' args=Args() model = build_model(args)

Фильтрация результатов сегментации и генерация trimap(троичных карт)

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

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

trimap = np.zeros((mask.shape[0], mask.shape[1], 2)) trimap[:, :, 1] = mask > 0 trimap[:, :, 0] = mask == 0 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(25,25)) trimap[:, :, 0] = cv2.erode(trimap[:, :, 0], kernel) trimap[:, :, 1] = cv2.erode(trimap[:, :, 1], kernel)

Применяем модель F, B, Alpha Matting.

fg, bg, alpha = pred((img/255.0)[:, :, ::-1], trimap, model)

Возвращаемся к исходному размеру и сохраняем изображение

img_ = img_orig.astype(np.float32)/255 alpha_ = cv2.resize(alpha, (img_.shape[1], img_.shape[0]), cv2.INTER_LANCZOS4) fg_alpha = np.concatenate([img_, alpha_[:, :, np.newaxis]], axis=2) cv2.imwrite(f_name + ".png", (fg_alpha*255).astype(np.uint8))

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

Все исходные тексты можно скачать из репозитория на GitHub.

Список источников

  1. Vitaliy Hramchenko How to cut a person out of a picture using python.
  2. https://pytorch.org/.
  3. Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam Rethinking Atrous Convolution for Semantic Image Segmentation.
  4. DeepLabV3 model with a ResNet-101 backbone Pytorch Team https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/.
  5. Marco Forte, François Pitié F, B, Alpha Matting.
  6. https://github.com/MarcoForte/FBA_Matting.
3 показа
904 открытия
0
1 комментарий
Dmitrii Nionov

че за модуль демо.. выложит его тоже на гитхаб..

Ответить
Развернуть ветку
Написать комментарий...
-2 комментариев
Раскрывать всегда