BodyPix:Real-time 人物分割

BodyPix: Real-time Person Segmentation

BodyPix 是一个 open-source 机器学习模型,它允许在浏览器中使用 TensorFlow.js 进行人物和 body-part 分割。 我想将模型转换为 .pb 冻结图以便在 Python.

上使用它

我该怎么做?

我尝试在不同的 places 上找到解决方案,但没有用。

  • 下载 model.json 文件

例如:https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float/model-stride16.json

  • 下载对应权重

https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float/group1-shard1of23.bin

...

https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float/group1-shard23of23.bin

  • 安装tfjs_graph_converter

来自 https://github.com/patlevin/tfjs-to-tf

  • 将模型转换为 .pb 文件

tfjs_graph_converter path/to/js/model path/to/frozen/model.pb

在 Python 中使用 bodypix 进行分割。但是只有当人在墙前而不是其他物体时才能得到更好的结果。

from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
import cv2
bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))

cap = cv2.VideoCapture(0) 
while cap.isOpened(): 
    ret, frame = cap.read()
    # BodyPix Segmentation
    result = bodypix_model.predict_single(frame)
    mask = result.get_mask(threshold=0.5).numpy().astype(np.uint8)
    seg = result.get_colored_part_mask(mask)

    tf.keras.preprocessing.image.save_img(
    pwd+"\output-colored-mask.jpg",
    seg
)