Как изолировать объекты сегментации с помощью Ultralytics YOLO#

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

Зачем изолировать объекты сегментации?#

Извлечение отдельных объектов из изображения открывает целый ряд последующих сценариев обработки:

• Удаление фона для предметной съемки, каталогов или творческого редактирования.
• Кадрирование отдельных объектов для создания наборов данных классификации на основе обнаруженных объектов.
• Сфокусированная обработка, чтобы последующие этапы, такие как OCR, анализ цвета или измерение, видели только объект, а не окружающую сцену.
• Экспорт в прозрачный PNG, чтобы размещать объекты на новых фонах.

Этот подход работает с любой моделью сегментации Ultralytics YOLO и состоит из четырех этапов: запустить инференс → извлечь каждый контур → изолировать объект → сохранить результат.

Запусти инференс сегментации#

Установи необходимые библиотеки, затем загрузи модель сегментации (суффикс -seg, необходимый для создания масок) и запусти предсказание для исходного изображения:

from ultralytics import YOLO

# Load a segmentation model
model = YOLO("yolo26n-seg.pt")

# Run inference on a source
results = model.predict(source="path/to/image.jpg")

Извлеки контуры объектов#

Каждый элемент в results соответствует одному изображению, а перебор результата возвращает по одному обнаружению за раз. Для каждого обнаружения скопируй исходное изображение, прочитай метку класса и нарисуй контур маски объекта на пустой бинарной маске. Белая область этой маски точно обозначает пиксели, принадлежащие объекту.

Фрагменты кода в этом и следующем разделе запускаются внутри показанного ниже цикла обнаружения; полный скрипт, готовый для копирования и вставки, находится в разделе Полный пример.

from pathlib import Path

import cv2
import numpy as np

for r in results:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem  # source image base-name

    # Iterate each detected object in the image
    for ci, c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]  # class name

        # Build a binary mask and draw the object contour onto it
        b_mask = np.zeros(img.shape[:2], np.uint8)
        contour = c.masks.xy[0].astype(np.int32).reshape(-1, 1, 2)
        cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

Изолируй объект#

Когда бинарная маска готова, объедини ее с исходным изображением. Есть два распространенных варианта, в зависимости от того, каким должен быть фон:

Выбери способ изоляции

Черный фонПрозрачный фон

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

# Isolate object with a black background
mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
isolated = cv2.bitwise_and(mask3ch, img)

Обрежь изображение по ограничивающей рамке

Чтобы оставить только область объекта, а не изображение полного размера, вырежи ее по ограничивающей рамке обнаружения:

# Bounding box coordinates
x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
# Crop the isolated image to the object region
iso_crop = isolated[y1:y2, x1:x2]

Сохрани результат (необязательно)#

Ты сам решаешь, что делать с каждым изолированным объектом. Распространенный следующий шаг — записать его на диск для дальнейшего использования:

# Save the isolated object to file
cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated)

Здесь img_name — это основа имени исходного изображения, label — имя класса, а ci — индекс обнаружения, поэтому несколько экземпляров одного класса получают уникальные имена файлов. Замени isolated на iso_crop, если выше применялась необязательная обрезка.

Полный пример#

Приведенный ниже скрипт объединяет все этапы в один исполняемый блок. По умолчанию он использует черный фон; замени единственную отмеченную строку на np.dstack([img, b_mask]), чтобы вместо этого получить прозрачный PNG:

from pathlib import Path

import cv2
import numpy as np

from ultralytics import YOLO

model = YOLO("yolo26n-seg.pt")
results = model.predict(source="path/to/image.jpg")

for r in results:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem

    for ci, c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]

        # Build a binary mask from the object contour
        b_mask = np.zeros(img.shape[:2], np.uint8)
        contour = c.masks.xy[0].astype(np.int32).reshape(-1, 1, 2)
        cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

        # Isolate the object (black background)
        mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
        isolated = cv2.bitwise_and(mask3ch, img)  # transparent PNG: isolated = np.dstack([img, b_mask])

        # Save or add your custom post-processing here
        cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated)

        # Optional: crop to the bounding box before saving
        # x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
        # cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated[y1:y2, x1:x2])

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

Заключение#

Теперь у тебя есть полный способ изолировать объекты сегментации с помощью Ultralytics YOLO: запустить инференс, создать бинарную маску для каждого контура, затем извлечь объект на черном или прозрачном фоне и при необходимости обрезать его по ограничивающей рамке. Изучи полную документацию по задаче сегментации и режиму предсказания, чтобы адаптировать рабочий процесс под свои классы.

Часто задаваемые вопросы#