Изолирование объектов сегментации

После выполнения задачи Segment Task иногда хочется извлечь изолированные объекты из результатов умозаключений. В этом руководстве приводится общий рецепт, как это сделать с помощьюрежима Ultralytics PredictMode.

Прогулка по рецептам

1. Начни с необходимого импорта

   from pathlib import Path
   
   import cv2
   import numpy as np
   from ultralytics import YOLO
   

   Ultralytics Install

   Ознакомься с разделом Ultralytics Quickstart Installation, где описана краткая инструкция по установке необходимых библиотек.

   ---

2. Загрузи модель и запусти predict() метод на источнике.

   from ultralytics import YOLO
   
   # Load a model
   model = YOLO('yolov8n-seg.pt')
   
   # Run inference
   results = model.predict()
   

   Нет аргументов в пользу предсказаний?

   Без указания источника будут использоваться примеры изображений из библиотеки:

   'ultralytics/assets/bus.jpg'
   'ultralytics/assets/zidane.jpg'
   

   Это полезно для быстрого тестирования с помощью predict() метод.

   Чтобы узнать больше о моделях сегментации, посетите Сегментная задача Страница. Чтобы узнать больше о predict() метод, см. Режим предсказания в разделе "Документация".

   ---

3. Теперь пройдись итерацией по результатам и контурам. Для рабочих процессов, которые хотят сохранить изображение в файл, исходное изображение base-name и обнаружение class-label извлекаются для последующего использования (необязательно).

   # (2) Iterate detection results (helpful for multiple images)
   for r in res:
       img = np.copy(r.orig_img)
       img_name = Path(r.path).stem # source image base-name
   
       # Iterate each object contour (multiple detections)
       for ci,c in enumerate(r):
           # (1) Get detection class name
           label = c.names[c.boxes.cls.tolist().pop()]
   

   1. Чтобы узнать больше о работе с результатами обнаружения, смотри раздел "Вставки" для режима Predict.
   2. Чтобы узнать больше о predict() Смотри результаты Работа с результатами для режима Predict
   For-Loop

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

   ---

4. Начни с создания бинарной маски из исходного изображения, а затем нарисуй на ней заполненный контур. Это позволит изолировать объект от других частей изображения. Пример из bus.jpg для одного из обнаруженных person Объекты класса показаны справа.

   

   # Create binary mask
   b_mask = np.zeros(img.shape[:2], np.uint8)
   
   # (1) Extract contour result
   contour = c.masks.xy.pop()
   # (2) Changing the type
   contour = contour.astype(np.int32)
   # (3) Reshaping
   contour = contour.reshape(-1, 1, 2)
   
   
   # Draw contour onto mask
   _ = cv2.drawContours(b_mask,
                       [contour],
                       -1,
                       (255, 255, 255),
                       cv2.FILLED)
   

   1. Для получения дополнительной информации о c.masks.xy смотри Раздел "Маски" из режима предсказания.

   2. Здесь значения приводятся к np.int32 для совместимости с drawContours() функция из OpenCV.

   3. OpenCV drawContours() Функция ожидает, что контуры будут иметь форму [N, 1, 2] расширьте раздел ниже, чтобы узнать больше подробностей.

   Расшифруй, чтобы понять, что происходит при определении contour переменная.

   • c.masks.xy :: Предоставляет координаты точек контура маски в формате (x, y). Для получения более подробной информации смотри Раздел "Маски" из режима предсказания.

   • .pop() :: Как masks.xy является списком, содержащим один элемент, этот элемент извлекается с помощью pop() метод.

   • .astype(np.int32) :: Использование masks.xy вернется с типом данных float32но это не будет совместимо с OpenCV. drawContours() Функция изменит тип данных на int32 для совместимости.

   • .reshape(-1, 1, 2) :: Переформатируй данные в нужную форму [N, 1, 2] где N это количество точек контура, каждая из которых представлена одной записью 1И запись состоит из 2 Значения. На сайте -1 обозначает, что количество значений по этому измерению гибкое.

   Разверни, чтобы получить объяснение drawContours() конфигурация.

   • Инкапсуляция contour переменная в квадратных скобках, [contour]во время тестирования было обнаружено, что он эффективно генерирует желаемую контурную маску.

   • Значение -1 указанный для drawContours() Параметр указывает функции нарисовать все контуры, присутствующие на изображении.

   • The tuple (255, 255, 255) Представляет собой белый цвет, который является желаемым цветом для рисования контура в этой бинарной маске.

   • Добавление cv2.FILLED окрасит все пиксели, заключенные в границу контура, в один и тот же цвет, в данном случае все заключенные пиксели будут белыми.

   • Смотри Документация по OpenCV на drawContours() за дополнительной информацией.

   ---

5. Далее есть 2 варианта того, как двигаться дальше с изображением с этой точки, и последующий вариант для каждого.

   Параметры изоляции объектов

   Пиксели черного фонаПиксели прозрачного фона

   # Create 3-channel mask
   mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
   
   # Isolate object with binary mask
   isolated = cv2.bitwise_and(mask3ch, img)
   

   Как это работает?

   • Во-первых, бинарная маска сначала преобразуется из одноканального изображения в трехканальное. Это преобразование необходимо для последующего шага, на котором маска и оригинальное изображение объединяются. Оба изображения должны иметь одинаковое количество каналов, чтобы быть совместимыми с операцией смешивания.

   • Оригинальное изображение и трехканальная бинарная маска объединяются с помощью функции OpenCV bitwise_and(). Эта операция сохраняет только Значения пикселей, которые больше нуля (> 0) из обоих изображений. Поскольку пиксели маски больше нуля (> 0) только В пределах контурной области от исходного изображения остаются пиксели, которые перекрывают контур.

   Изолировать с помощью черных пикселей: Подварианты

   Полноразмерное изображение

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

   

   Обрезанное изображение объекта

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

   

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

   1. Подробнее о результатах работы с ограничительными рамками читай в разделе Боксы из режима Predict.

   Что делает этот код?

   • The c.boxes.xyxy.cpu().numpy() Вызов извлекает ограничительные рамки в виде массива NumPy в xyxy формат, в котором xmin, ymin, xmax, и ymax представляют собой координаты прямоугольника ограничительной рамки. Смотри Секция коробок из режима предсказания для более подробной информации.

   • The squeeze() Операция удаляет лишние размеры из массива NumPy, обеспечивая ему ожидаемую форму.

   • Преобразование значений координат с помощью .astype(np.int32) Измени тип данных координат коробки с float32 на int32что делает их совместимыми для обрезки изображений с помощью индексных срезов.

   • Наконец, область ограничительной рамки вырезается из изображения с помощью индексной нарезки. Границы определяются с помощью [ymin:ymax, xmin:xmax] координаты ограничительной рамки обнаружения.

   # Isolate object with transparent background (when saved as PNG)
   isolated = np.dstack([img, b_mask])
   

   Как это работает?

   • Использование NumPy dstack() Функция (укладка массива по оси глубины) в сочетании со сгенерированной бинарной маской создаст изображение с четырьмя каналами. Это позволит сделать все пиксели за пределами контура объекта прозрачными при сохранении в формате PNG файл.

   Изолировать с помощью прозрачных пикселей: Подварианты

   Полноразмерное изображение

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

   

   Обрезанное изображение объекта

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

   

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

   1. Подробнее о результатах работы с ограничительными рамками читай в разделе Боксы из режима Predict.

   Что делает этот код?

   • При использовании c.boxes.xyxy.cpu().numpy()Граничные поля возвращаются в виде массива NumPy, используя xyxy формат координат бокса, которые соответствуют точкам xmin, ymin, xmax, ymax Ограничительную рамку (прямоугольник) см. Секция коробок из режима предсказания за дополнительной информацией.

   • Добавление squeeze() гарантирует, что все лишние размеры будут удалены из массива NumPy.

   • Преобразование значений координат с помощью .astype(np.int32) Измени тип данных координат коробки с float32 на int32 что будет совместимо при обрезке изображения с помощью индексных срезов.

   • Наконец, область изображения для ограничительного поля обрезается с помощью индексной нарезки, где границы задаются с помощью [ymin:ymax, xmin:xmax] координаты ограничительной рамки обнаружения.

   А если я хочу получить обрезанный объект, включая фон?

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

   ---

6. Что делать дальше, полностью зависит от тебя как разработчика. Здесь показан базовый пример одного из возможных следующих шагов (сохранение изображения в файл для дальнейшего использования).

   • ПРИМЕЧАНИЕ: этот шаг необязателен и может быть пропущен, если не требуется для твоего конкретного случая использования.
   Пример последнего шага

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

   • В этом примере img_name базовое имя исходного файла изображения, label это имя обнаруженного класса, а ci это индекс обнаружения объекта (в случае нескольких экземпляров с одним и тем же именем класса).

Полный код примера

Здесь все шаги из предыдущего раздела объединены в один блок кода. Для повторного использования оптимальным будет определить функцию, которая будет выполнять некоторые или все команды, содержащиеся в for-петли, но это упражнение остается на усмотрение читателя.

from pathlib import Path

import cv2
import numpy as np
from ultralytics import YOLO

m = YOLO('yolov8n-seg.pt')#(4)!
res = m.predict()#(3)!

# iterate detection results (5)
for r in res:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem

    # iterate each object contour (6)
    for ci,c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]

        b_mask = np.zeros(img.shape[:2], np.uint8)

        # Create contour mask (1)
        contour = c.masks.xy.pop().astype(np.int32).reshape(-1, 1, 2)
        _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

        # Choose one:

        # OPTION-1: Isolate object with black background
        mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
        isolated = cv2.bitwise_and(mask3ch, img)

        # OPTION-2: Isolate object with transparent background (when saved as PNG)
        isolated = np.dstack([img, b_mask])

        # OPTIONAL: detection crop (from either OPT1 or OPT2)
        x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
        iso_crop = isolated[y1:y2, x1:x2]

        # TODO your actions go here (2)

1. Строка, заполняющая contour Здесь она объединена в одну строку, там, где выше она была разделена на несколько.
2. Что здесь будет, зависит только от тебя!
3. Дополнительную информацию смотри в разделе Режим предсказания.
4. Подробнее об этом читай в разделе " Сегментная задача ".
5. Узнай больше о работе с результатами
6. Узнай больше о результатах применения маски сегментации

Комментарии