Bria AI выпустили RMBG 2.0 — нейросеть для удаления фона на фотографиях

Разработчики организации Bria AI выпустили RMBG 2.0 — обновление нейросети для удаления фона на фотографиях. Модель машинного обучения теперь лучше распознаёт объекты на переднем плане и поддерживает больше сценариев работы.

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

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

С мелкими объектами нейросеть справляется плохо

Нейросеть RMBG опубликовали на площадке Hugging Face. Есть версии для PyTorch и Safetensors. В тестовом пространстве разработчики развернули бесплатное демо. Для некоммерческого использования RMBG доступна по лицензии Creative Commons. 

Пример кода для запуска:

from PIL import Image import matplotlib.pyplot as plt import torch from torchvision import transforms from transformers import AutoModelForImageSegmentation model = AutoModelForImageSegmentation.from_pretrained(‘briaai/RMBG-2.0’, trust_remote_code=True) torch.set_float32_matmul_precision([‘high’, ‘highest’][0]) model.to(‘cuda’) model.eval() # Data settings image_size = (1024, 1024) transform_image = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) image = Image.open(input_image_path) input_images = transform_image(image).unsqueeze(0).to(‘cuda’) # Prediction with torch.no_grad(): preds = model(input_images)[-1].sigmoid().cpu() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize(image.size) image.putalpha(mask) image.save(«no_bg_image.png»)

Источник: habr.com

Похожие новости:

Вам врут, когда говорят, что китайские смартфоны во всём превзошли iPhone. Вот в чём они конкретно лучше Компактный топ-смартфон Xiaomi с быстрейшим процессором и даже хорошей автономностью. Но… Наш вердикт по новым смартфонам среднего класса Infinix Note 40 Лучший смартфон среднего класса в 2024 году — Realme 12 Pro+