SAM от Meta: простая сегментация изображений с минимальным кодом.

Пошаговое руководство.

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

В этом уроке я покажу вам, как использовать SAM с библиотекой transformers Hugging Face. Я покажу, как создавать нулевые маски сегментации для любого изображения.

К концу этого руководства вы получите четкое представление о том, как использовать SAM с трансформерами, и сможете применить его к своей собственной сегментации изображений проекты. Давайте начнем!

Шаг 1: Импорт модели

Мы собираемся использовать пакет Transformers для локальной загрузки модели. Hugging Face предоставляет нам конвейер, которому можно назначить задачу и, при необходимости, конкретную модель. Мы тестируем модель SAM для генерации маски, поэтому нам просто нужно указать ее следующим образом:

from transformers import pipeline

generator = pipeline("mask-generation", model="facebook/sam-vit-large", device=0)

На момент написания этой статьи модель SAM еще не была доступна в пакете transforms, поэтому мне нужно было установить пакет из исходного кода, что в основном означает, что вы устанавливаете и собираете его из исходного кода на GitHub. Вы можете легко сделать это с помощью этой команды bash:

pip install -q --upgrade git+https://github.com/huggingface/transformers.git

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

Шаг 2: Подготовка изображений

Теперь мы собираемся подготовить изображения для модели. Мы будем использовать библиотеку PIL для загрузки и обработки изображений и, в конечном итоге, requests, если нам нужно получить наши изображения из Интернета.

from PIL import Image
import requests

# Local Image
local_image = Image.open("path/to/image.jpg").convert("RGB")

# Hosted image
img_url = "https://example.com/image.jpg"
web_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")

Вы можете использовать matplotlib.pyplot, чтобы убедиться, что вы загрузили правильные изображения, или, в конечном итоге, отладить любые проблемы.

import matplotlib.pyplot as plt

plt.imshow(web_image)

Мы будем использовать эту фотографию прекрасного Дахаба, расположенного в Египте, до конца урока:

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

Шаг 3: Отображение функций

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

import numpy as np
import matplotlib.pyplot as plt
import

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)
    del mask
    gc.collect()

def show_masks_on_image(raw_image, masks):
  plt.imshow(np.array(raw_image))
  ax = plt.gca()
  ax.set_autoscale_on(False)
  for mask in masks:
      show_mask(mask, ax=ax, random_color=True)
  plt.axis("off")
  plt.show()
  del mask
  gc.collect()

Теперь у нас есть способ отображать выходные данные модели на наших изображениях. Мы готовы начать сегментацию.

Шаг 4: Давайте сегментируем что угодно!

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

outputs = generator(local_image, points_per_batch=64)

masks = outputs["masks"]
show_masks_on_image(local_image, masks)

Модель вернет информацию о сегментах внутри outputs, и мы извлечем каждую маску и отобразим ее на нашем изображении. Вы должны получить что-то вроде этого:

Если вам нужно указать ограничивающие рамки или точки, вы можете найти дополнительную информацию в этих примерах блокнотов, предоставленных Meta и Hugging Face, которые содержат готовый к использованию код.

Поздравляем, теперь вы знаете, как быстро приступить к работе с SAM! Для справки в будущем вы можете использовать любую из тысяч моделей с открытым исходным кодом на HF Hub, используя конвейер, как это сделали мы. Обычно требуется очень мало строк кода, чтобы все настроить и приступить к тестированию.

Проверьте мою последнюю статью ниже, если вам понравилось!