Сегментация селфи в Python с использованием OpenCV и Mediapipe

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

ИМПОРТ ЗАВИСИМОСТЕЙ

import mediapipe as mp
import cv2
import numpy as np
import matplotlib.pyplot as plt

СОЗДАНИЕ МОДЕЛИ

mp_selfie = mp.solutions.selfie_segmentation

mp.solutions позволяет нам получить доступ к различным предварительно обученным моделям в Mediapipe. В нашем случае мы реализовали сегментацию селфи.

СОЗДАНИЕ Прямых трансляций

cap = cv2.VideoCapture(2)
# Create with statement for model 
with mp_selfie.SelfieSegmentation(model_selection=0) as model: 
 while cap.isOpened():
 ret, fram = cap.read()
 frame = cv2.resize(fram, (1280,650))
 
 # Apply segmentation
 frame.flags.writeable = False
 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 results = model.process(frame)
 frame.flags.writeable = True
cv2.imshow(‘Webcam feed’, frame)
if cv2.waitKey(1) & 0xFF == ord(‘q’):
 break
cap.release()
cv2.destroyAllWindows()

Здесь мы сначала получили доступ к нашей веб-камере, используя VideoCapture() method OpenCV, затем мы определили модель, используя with statement, и назвали ее model. Теперь, пока наша веб-камера открыта, мы берем изображение, изменяем его размер до желаемого соотношения сторон, делаем изображение доступным для записи для лучшей обработки, преобразуем цвет изображения из RGB в BGR, поскольку OpenCV работает с изображениями BGR. Затем мы отправили это изображение в нашу модель для обработки и присвоили его results, который содержит массив вероятностей того, где мы находимся в кадре в определенный момент времени. Затем мы устанавливаем изображение как недоступное для записи (false). Затем мы показали это через веб-камеру. Оператор if используется для выхода и закрытия веб-камеры, если мы нажмем клавишу «q».

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

background = np.zeros(frame.shape, dtype=np.uint8)
mask = np.stack((results.segmentation_mask,)*3, axis=-1) > 0.5
segmented_image = np.where(mask, frame, background)
plt.imshow(segmented_image);

Итак, здесь мы сначала создали background массив нулей той же формы, что и наш кадр, указав тип данных как uint8. Это означает, что фон представляет собой просто черное изображение. Затем мы создали mask, а затем использовали метод np.where, в котором мы передали нашу маску, рамку и изображение. Теперь, если мы построим это, оно покажет изображение, на котором все будет черным, кроме нашего селфи.

Теперь мы можем легко добавить все это в живую рабочую модель.

СОЕДИНЯЕМ ВСЕ ВМЕСТЕ!

cap = cv2.VideoCapture(0)
# Create with statement for model 
with mp_selfie.SelfieSegmentation(model_selection=0) as model: 
 while cap.isOpened():
 ret, frame = cap.read()
 frame = cv2.resize(frame, (1280,650))
 
 frame.flags.writeable = False
 
 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 results = model.process(frame)
 
 mask = np.stack((results.segmentation_mask,)*3, axis=-1) > 0.5 
 
 
 background = np.zeros(frame.shape, dtype=np.uint8)
 imgs = np.where(mask, frame, background)
 frame.flags.writeable = True
 final = cv2.cvtColor(imgs, cv2.COLOR_RGB2BGR)
 
 cv2.imshow(‘Webcam feed’, final)
if cv2.waitKey(1) & 0xFF == ord(‘q’):
 break
cap.release()
cv2.destroyAllWindows()

Снова тот же метод, мы открыли нашу камеру, взяли наш кадр, отправили его в нашу модель и присвоили его переменной результатов. Затем мы создали маску и фон, отправили рамку, маску и фон методу where Numpy и присвоили его img.. Затем мы преобразовали цвет imgback в RGB и показали его в нашей прямой трансляции в режиме реального времени.

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

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

Вы можете увидеть рабочую модель в моем репозитории GitHub здесь.