mediapipe 分割segmentation
基于图像
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2 as cv
import numpy as np
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,
output_confidence_masks=True
)
segmenter = vision.ImageSegmenter.create_from_options(options)
image = cv.imread("image/s1.jpg")
rgb_image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_image)
segmentation_result = segmenter.segment(mp_image)
confidence_masks = segmentation_result.confidence_masks
mask_np = confidence_masks[0].numpy_view()
mask_np = np.where(mask_np > 0.5, 255, 0).astype(np.uint8)
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
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2 as cv
import numpy as np
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
)
segmentor = vision.ImageSegmenter.create_from_options(options)
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)
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()