Mediapipe学习 分割Segmentation

mediapipe 分割segmentation

基于图像

# step 1
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2 as cv
import numpy as np

# BG_COLOR = (192, 192, 192) # gray
# MASK_COLOR = (255, 255, 255) # white

# step 2 create the segmentor
base_options = python.BaseOptions(model_asset_path='models/selfie_segmenter_landscape.tflite')
options = vision.ImageSegmenterOptions(base_options=base_options,
                                       running_mode=vision.RunningMode.IMAGE,
                                       output_category_mask=False,  # category_mask 判断前景的所属范畴
                                       output_confidence_masks=True # confidence_masks 确定度 0-1之间
                                       )
segmenter = vision.ImageSegmenter.create_from_options(options)

# step 3 read the image
image = cv.imread("image/s1.jpg")
rgb_image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
# convert image to mp_image
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_image)
# step 4 segment the image
segmentation_result = segmenter.segment(mp_image)
# step 5 create a mask and show it
confidence_masks = segmentation_result.confidence_masks # list[MPIMAGE]
mask_np = confidence_masks[0].numpy_view()
# mask_up 做为image的mask
mask_np = np.where(mask_np > 0.5, 255, 0).astype(np.uint8)
# opencv mask
# mask_np = cv.bitwise_not(mask_np)
# mask_np 用于 image
image = cv.bitwise_and(image, image, mask=mask_np)
cv.namedWindow("mask", cv.WINDOW_KEEPRATIO)
cv.imshow("mask", mask_np)
cv.namedWindow("image", cv.WINDOW_KEEPRATIO)
cv.imshow("image", image)

cv.waitKey(0)
cv.destroyAllWindows()

基于视频Video

# step 1 import library
import mediapipe as mp
from  mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2 as cv
import numpy as np

# step 2 create segmentor
base_options = python.BaseOptions(
    model_asset_path='models/selfie_segmenter_landscape.tflite'  # 从官网下载的模型文件
)
options = vision.ImageSegmenterOptions(
    base_options=base_options,
    running_mode=vision.RunningMode.VIDEO,
    output_category_mask=True,                       # 两种模式二选一，这个返回一个掩码
    output_confidence_masks=False                    # 如果选择，返回list
)
segmentor = vision.ImageSegmenter.create_from_options(options)

# step 3 get cap
# cv 读取视频
cap = cv.VideoCapture('videos/yoga.mp4')  # 视频文件
cv.namedWindow('frame', cv.WINDOW_KEEPRATIO)
cv.namedWindow('mask', cv.WINDOW_KEEPRATIO)
fps = cap.get(cv.CAP_PROP_FPS)
frame_index = 0
while True:
    ret, frame = cap.read()
    if not ret:
        break
    frame_index += 1
    # 生成时间戳
    frame_timestamp_ms = int(1000 / fps * frame_index)
# get segmentation result
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
    # 获取分成结果
    segmentation_result = segmentor.segment_for_video(mp_image, frame_timestamp_ms)
    # 计算掩码
    category_mask = segmentation_result.category_mask.numpy_view()
    category_mask = np.array(category_mask, dtype=np.float32)
    category_mask = np.where(category_mask > 0.5, 0, 255).astype(np.uint8)
    frame = cv.bitwise_and(frame, frame, mask=category_mask)
    # 显示
    cv.imshow('frame', frame)
    cv.imshow('mask', category_mask)
    if cv.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv.destroyAllWindows()