Meta AI представляет революционную модель сегментации изображений, обученную на 1 миллиарде масок - UPROGER

Вступление

После революционного шага, сделанного ChatGPT от OpenAI в области NLP, развитие искусственного интеллекта продолжается, и Meta AI вносит поразительный прогресс в компьютерное зрение. Исследовательская группа Meta AI представила модель под названием Segment Anything Model (SAM) и набор данных из 1 миллиарда масок на 11 миллионах изображений. Сегментация изображения – это разбиение изображения на множество покрывающих его областей.

Meta AI представляет революционную модель сегментации изображений, обученную на 1 миллиарде масок

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

1. Сегментируйте любую задачу

Основной целью команды Meta AI было создать модель сегментации изображений, которая работала бы с подсказками пользовательского ввода так же, как она работает с ChatGPT. Поэтому они придумали решение для интеграции пользовательского ввода с изображением для создания масок сегментации. Подсказкой для сегментации может быть любая информация, указывающая, что нужно сегментировать в изображении. Например, набор точек переднего плана или фона, прямоугольник, текст произвольной формы и т.д. Таким образом, выходные данные модели представляют собой допустимую маску сегментации для любого пользовательского запроса.

2. Сегментируйте любую модель

Segment Anything Model (SAM) состоит из трёх компонентов, показанных на рисунке ниже:

Архитектура модели высокого уровня состоит из кодировщика изображений, кодировщика подсказок и декодировщика масок. Для кодировщика изображений они использовали предварительно обученную модель MAE [1], имеющую архитектуру Vision Transformer(ViT) [2]. Модели ViT – это самые современные модели для решения задач классификации и сегментации изображений. Что касается подсказок, они разделили их на два типа — один тип подсказок разреженный, такой как точки, прямоугольники и текст, а другой тип плотный, такой как маски. На шаге кодирования подсказок создаются вложения для каждого типа подсказок. Что касается декодировщика масок, он просто сопоставляет вложения изображений, встраивания подсказок и выходные токены с маской.

3. Сегментируйте любые данные

3.1 Сегментируйте любой механизм обработки данных

Принцип GIGO применим и к области искусственного интеллекта. Если входные данные низкого качества, результат, сгенерированный моделью, также не будет хорошим. Именно поэтому команда Meta постаралась отобрать высококачественные изображения для обучения своей модели. Команда создала механизм обработки данных для фильтрации набора необработанных изображений. Создание механизма обработки данных делится на три этапа:

Ручной этап: Для нанесения масок на изображение вручную были привлечены профессиональные аннотаторы.
Полуавтоматический этап: Они обучили модель на аннотированных изображениях и сделали вывод на основе остальных изображений. Затем аннотаторов попросили пометить дополнительные немаркированные объекты, которые не были обнаружены моделью, или исправить сегменты с низкими показателями достоверности.
Полностью автоматический этап: Этот этап включает в себя автоматическую генерацию масок и автоматическую фильтрацию, которая пытается убрать неоднозначные маски и сохранить маски, основанные на достоверности, стабильности и размере.

3.2 Сегментируйте любой набор данных

Механизм обработки данных Segment Anything создал набор данных в 1 миллиард масок (SA-1B) на 11 миллионах разнообразных изображений с высоким разрешением (в среднем 3300×4900 пикселей). Стоит отметить, что 99,1% масок были сгенерированы автоматически, однако качество настолько высокое, потому что они тщательно отобраны.

Вывод — Почему модель является эволюционной

Команда Meta AI, вместе с другими командами крупной компании, добивается больших успехов в разработке искусственного интеллекта. Segment Anything Model (SAM) обладает возможностями для работы приложений во многих областях, требующих поиска и сегментации любого объекта на любом изображении. Например:

SAM мог бы быть компонентом большой мультимодальной модели, которая интегрировала бы изображения, текст, аудио и т.д.
SAM мог бы позволить выбирать объект в области AR / VR на основе взгляда пользователя.
SAM может улучшить творческие приложения, такие как извлечение областей изображения для редактирования видео.
и многое другое.

Демонстрация сегментации изображений

В этой части я попытаюсь использовать официальный код GitHub, чтобы поиграть с алгоритмом с использованием Google Colab и выполнить два типа сегментации изображения. Сначала я выполню сегментацию с помощью пользовательского запроса, а затем я выполню полностью автоматическую сегментацию.

Часть 1: Сегментация изображения с использованием пользовательского запроса

1. Настройка (импорт библиотек):

from IPython.display import display, HTML
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2


display(HTML(
"""
<a target="_blank" href="https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb">
  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
</a>
"""
))

using_colab = True

if using_colab:
    import torch
    import torchvision
    print("PyTorch version:", torch.__version__)
    print("Torchvision version:", torchvision.__version__)
    print("CUDA is available:", torch.cuda.is_available())
    import sys
    !{sys.executable} -m pip install opencv-python matplotlib
    !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'
    
    !mkdir images
    !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg
    !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/groceries.jpg
        
    !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

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

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))

3. Входное изображение (исходное изображение для сегментации). Давайте попробуем выбрать маску из первого пакета с продуктами:

image = cv2.imread('/content/images/groceries.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 
plt.figure(figsize=(5,5))
plt.imshow(image)
plt.axis('on')
plt.show()

4. Загрузите предварительно обученную модель с именем sam_vit_h_4b8939.pth, которая является моделью по умолчанию. Существуют и другие облегчённые версии моделей, такие как sam_vit_l_0b3195.pth и sam_vit_b_01ec64.pth.

sam_checkpoint = "/content/sam_vit_h_4b8939.pth"
device = "cuda"
model_type = "default"

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

predictor = SamPredictor(sam)

predictor.set_image(image)

5. Визуализируйте точку на изображении (подсказка пользователя), которая поможет идентифицировать наш целевой объект — первую сумку с продуктами.

input_point = np.array([[465, 300]])
input_label = np.array([1])
plt.figure(figsize=(10,10))
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
plt.axis('on')
plt.show()

6. Сделайте прогноз, чтобы сгенерировать маску объекта.

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)
print(masks.shape)  # (number_of_masks) x H x W

7. Покажите 3 лучших сгенерированных маски. Когда multimask_output=True, алгоритм возвращает три маски. Позже мы сможем выбрать тот, который наберёт наибольшее количество баллов.

for i, (mask, score) in enumerate(zip(masks, scores)):
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_mask(mask, plt.gca())
    show_points(input_point, input_label, plt.gca())
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
    plt.axis('off')
    plt.show()

Выделенные объекты – это маски, предсказанные моделью. Как показывает результат, модель сгенерировала три выходные маски со следующими показателями прогнозирования: mask1 — 0,990, Mask2 — 0,875 и Mask3 — 0,827. Мы выбираем mask1, которая набрала наибольшее количество баллов. Вуаля!!!! Маска прогнозирования модели – это целевой объект, который мы изначально хотели сегментировать. Результат потрясающий, модель работает довольно хорошо!

Часть 2: Полностью автоматическая сегментация изображений — Продолжение.

1. Функция построения графиков сегментов:

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

2. Автоматически создавайте маски:

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)

masks = mask_generator.generate(image)
print(len(masks))

3. Вывод результата:

plt.figure(figsize=(5,5))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

Алгоритм идентифицировал 137 различных объектов (масок), используя параметры по умолчанию. Каждая маска содержит информацию о площади сегмента, координатах ограничивающей рамки, показателях прогнозирования и стабильности, которые могут быть использованы для фильтрации необычных сегментов.

Я надеюсь, что вам понравился данный контент!. Теперь вы сможете начать создавать красивые приложения самостоятельно!

Источники

[1] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollar, and Ross Girshick. Masked autoencoders are scalable vision learners. CVPR, 2022.

[2] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16×16 words: Transformers for image recognition at scale. ICLR, 2021.

[3] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollar, Ross Girshick. Segment Anything, 2023