Программа BiRefNet для удаления фона на фотографиях

Задача убрать фон была всегда, и тем или иным способом с ней справлялись. Вспоминаю случай начала 80х. Надо было срочно сдавать фотографии сотрудников лаборатории на паспорт моряка. Это для тропиков голова на колу галстука в овале. Оказалось, что часть имеющихся негативов сняты не на те документы и либо фон не белый, либо пиджак присутствует. Нет проблем. Есть темная комната с увеличителем Беларусь-2, и я приступаю к работе. Печатаю, проявляю, сушу, вырезаю, размещаю чуть выше экспонируемой фотобумаги, чтобы размыть края, и печатаю. Задача решена: документы прошли. С появлением компьютера маску стало делать проще, но все равно нужна внимательность и время. Если надо убрать фон с двух фотографий для статьи, то нет проблем, а если для сотен фотографий образцов в каталоге товаров, то становится грустно. Поэтому, когда компьютер стал сам учится быть внимательным к определенным деталям, результат превзошел самые смелые ожидания.

Полное название проекта : Bilateral Reference for High-Resolution Dichotomous Image Segmentation. Насколько он получился удачным можно убедиться, запустив программу в сети на платформе Gradio. Наибольшее количество примеров и возможность использовать разные варианты тренировки машинного обучения можно посмотреть здесь. Можно загрузить и свою картинку и получить результат без ограничения размеров. Правда, фон будет черным. Если нужен белый фон, то запускаем другой вариант, без расширенных настроек. Для запуска на локальной машине идем на сайт и скачиваем репозиторий. Удовлетворяем зависимости, перечисленные в файле requirements.txt. Это может вызвать некоторые трудности, поскольку для моей ОС manjaro некоторые можно найти только в AUR и там без правки настроек установить не удается. Но вот мы все установили и хотим заняться обработкой своих изображений. Не тут-то было. На сайте подробно описано как тренировать программу, но нет указания на программу, которую можно запустить и воспользоваться результатами чужих тренировок. Впрочем, все не так плохо, нас отсылают к возможности это сделать в сети на платформе colab. Google Colab — предоставляет возможность работать с кодом на языке Python в браузере. Запускаем и убеждаемся, что все работает. Можно и свое изображение подгрузить, однако велика вероятность, что вам скажет, что нет свободных бесплатных ресурсов GPU. Программа может работать и без графического процессора, но это требует изменения в коде. Поэтому просто сохраняем программу на локальном компьютере и немного видоизменяем. Например так, как это сделал я.

Программа

# -*- coding: utf-8 -*-


import argparse
 
parser = argparse.ArgumentParser(description="Приветствие пользователя")
parser.add_argument("name", help="cpu/cuda")
parser.add_argument("name2", help="Имя файла") 
parser.add_argument("value", type=int, help="N(56)x16")
args = parser.parse_args()
scl=args.value*16
import matplotlib.pyplot as plt  # для вывода картинки
from PIL import Image
import torch
from torchvision import transforms
# from IPython.display import display

from models.birefnet import BiRefNet


# # Option 1: loading BiRefNet with weights:
# from transformers import AutoModelForImageSegmentation
# birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)

# # Option-2: loading weights with BiReNet codes:
# birefnet = BiRefNet.from_pretrained('zhengpeng7/BiRefNet')

# Option-3: Loading model and weights from local disk:
from utils import check_state_dict

birefnet = BiRefNet(bb_pretrained=False)
state_dict = torch.load('BiRefNet-DIS-epoch_590.pth', map_location='cpu')
state_dict = check_state_dict(state_dict)
birefnet.load_state_dict(state_dict)


# Load Model
device = args.name
torch.set_float32_matmul_precision(['high', 'highest'][0])

birefnet.to(device)
birefnet.eval()
print('BiRefNet is ready to use.')

# Input Data
transform_image = transforms.Compose([
    transforms.Resize((scl, scl)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

imagepath = args.name2 +'.jpg'
image = Image.open(imagepath)
input_images = transform_image(image).unsqueeze(0).to(args.name)

# Prediction
with torch.no_grad():
    preds = birefnet(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()

# Show Results
pred_pil = transforms.ToPILImage()(pred)
# Scale proportionally with max length to 1024 for faster showing
scale_ratio = scl / max(image.size)
scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))
image_masked = image.resize((scl, scl))
image_masked.putalpha(pred_pil)
im=image_masked.resize(scaled_size)
im.save(args.name2 +".png")
imp=pred_pil.resize(scaled_size)
imp.save(args.name2 +"m.png")


plt.imshow(image_masked.resize(scaled_size))
plt.show()
# plt.imshow(image.resize(scaled_size))
# plt.show()
plt.imshow(pred_pil.resize(scaled_size))
plt.show()

Вариант 1 - это BiRefNet-portrait-epoch_150.pth - волосы обрезаны великолепно, но коробку на заднем плане от платья не отделили и часть циркуля удалили. Вариант 2 - это BiRefNet-DIS-epoch_590.pth - циркуль и платье идеально, у волос черная кайма.