Удаление фона с фотографии (AI-технологии)

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

Честно говоря, про AI я немного слукавил и для этих операций используется обученная модель которая в большинстве случаев уберет фон на вашем фото и вы получите PNG файл с альфа-каналом.

Есть несколько сайтов для удаления фона с фотографии например https://www.photoroom.com, но они платные и, честно говоря, они используют тот же подход который мы сейчас и рассмотрим.

Использовать мы будем проект Rembg и его репозиторий на GitHub находится по адресу https://github.com/danielgatis/rembg.

Первым делом мы скачиваем репозиторий проекта и собираем Docker образ содержащий утилиту rembr и пред обученную модель для удаления фона с фотографии.

$ git clone https://github.com/danielgatis/rembg.git
$ cd ./rembg
$ docker build -t rembg:1 ./

Создаем два каталога input и output, соответственно в input скидываем исходные фотографии, а в output будем помещать обработанные фотографии в которых будет удален фон.

Конвертация выполняется командой.

$ docker run -v ./input:/rembg -v ./output:/output rembg:1 i 00124.jpg /output/00124.png

Таким образом мы запускаем утилиту в докере с параметром на конвертацию и двумя подключенными томами. Давайте посмотрим, что у нас получилось.

За основу взял довольно сложное изображение, вот его исходник.

Платная утилита от photoroom справилась вот так (она платная только в том случае если вам FullHD надо).

Утилита с моделью по умолчанию справилась так.

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

#!/bin/sh

#docker run -v ./input:/rembg -v ./output:/output rembg:1 i 00124.jpg /output/00124.png

find ./input/*.jpg | while read file_2_convert
    do
    source_file=`basename "$file_2_convert"`
    file_name=`basename "$file_2_convert" | cut -d "." -f 1`
    echo "$source_file" - "$file_name".png
    docker run -v ./input:/rembg -v ./output:/output rembg:1 i "$source_file" /output/"$file_name".png
    done

exit 0

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

Результат скажем так не сильно обрадовал и попробуем другую обученную модель.

На выбор у нас несколько моделей.

• u2net (download, source): A pre-trained model for general use cases.
• u2netp (download, source): A lightweight version of u2net model.
• u2net_human_seg (download, source): A pre-trained model for human segmentation.
• u2net_cloth_seg (download, source): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
• silueta (download, source): Same as u2net but the size is reduced to 43Mb.
• isnet-general-use (download, source): A new pre-trained model for general use cases.
• isnet-anime (download, source): A high-accuracy segmentation for anime character.
• sam (download encoder, download decoder, source): A pre-trained model for any use cases.
• birefnet-general (download, source): A pre-trained model for general use cases.
• birefnet-general-lite (download, source): A light pre-trained model for general use cases.
• birefnet-portrait (download, source): A pre-trained model for human portraits.
• birefnet-dis (download, source): A pre-trained model for dichotomous image segmentation (DIS).
• birefnet-hrsod (download, source): A pre-trained model for high-resolution salient object detection (HRSOD).
• birefnet-cod (download, source): A pre-trained model for concealed object detection (COD).
• birefnet-massive (download, source): A pre-trained model with massive dataset.

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

$ docker run -v ./input:/rembg -v ./output:/output rembg:1 i -m birefnet-massive  00124.jpg /output/00124.png
Downloading data from 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx' to file '/root/.u2net/birefnet-massive.onnx'.

Все модели хранятся внутри контейнера в каталоге /root/.u2net и логично его вытащить в еще один внешний том, чтобы не загружать их на каждый чих и для этого добавим еще один том опцией -v ./models:/root/.u2net.

$ mkdir ./models
$ docker run -v ./input:/rembg -v ./output:/output -v ./models:/root/.u2net rembg:1 i -m silueta  00124.jpg /output/00124.png

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

Давайте сейчас попробуем все модели и посмотрим какая справится лучше с фотографией с зонтиками с корпоратива 4Финанс (стандартный u2net пробовали в начале статьи).

u2netp

Моделька в 4МБ, в принципе результат предсказуем.

u2net_human_seg

Модель 176МБ и прям только люди и остались, зонтики исчезли.

Попробуем еще одно фото, где я на относительно ровном фоне.

Фото само по себе неудачное, а моделька справилась просто замечательно.

u2net_cloth_seg

Аналогично моделька на 176МБ. Но получилось нечто странное.

Судя по результату по фото в костюме эта модель пытается отдельно одежду выбрать.

silueta

У меня зонтик забрали. Но все равно ореол остается. Моделька так-же 170 с копейками МБ.

isnet-general-use

Модель 179МБ. По описанию новая модель для общих целей.

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

Пока из всех испробованных моделей самый лучший результат.

isnet-anime

Модель 176МБ для обработки аниме-картинок.

sam

А вот тут уже солидные 360МБ. Странная моделька, остались одни брюки.

birefnet-general

Серьезный датасет в 973МБ.