如何使用MediaPipe多类自拍分割模型？

Question

我正在使用 MediaPipe 库，想要从图像中检测头发和衣服。在他们的网站上，他们提供了很好的示例，您可以下载并运行代码。但是，我希望当我检测到头发或衣服时，该区域将变成白色，而所有其他区域保持不变。然而，MediaPipe 仅给出了“DeepLab-v3”模型的示例。在这里你可以看到他们的例子。

https://github.com/googlesamples/mediapipe/tree/main/examples/image_segmentation/js https://developers.google.com/mediapipe/solutions/vision/image_segmenter/index#models

这是他们的回调函数，我想检测头发或衣服并将它们设为白色。

function callback(result) {
    const cxt = canvasClick.getContext("2d");
    const { width, height } = result.categoryMask;

    let imageData = cxt.getImageData(0, 0, width, height).data;

    canvasClick.width = width;
    canvasClick.height = height;
    let category = "";

    const mask = result.categoryMask.getAsUint8Array();

    for (let i in mask) {
        if (mask[i] > 0) {
            category = labels[mask[i]];
        }
        const legendColor = legendColors[mask[i] % legendColors.length];
        imageData[i * 4] = (legendColor[0] + imageData[i * 4]) / 2;
        imageData[i * 4 + 1] = (legendColor[1] + imageData[i * 4 + 1]) / 2;
        imageData[i * 4 + 2] = (legendColor[2] + imageData[i * 4 + 2]) / 2;
        imageData[i * 4 + 3] = (legendColor[3] + imageData[i * 4 + 3]) / 2;
    }
    
    const uint8Array = new Uint8ClampedArray(imageData.buffer);
    const dataNew = new ImageData(uint8Array, width, height);
    cxt.putImageData(dataNew, 0, 0);
    const p = event.target.parentNode.getElementsByClassName("classification")[0];
    p.classList.remove("removed");
    p.innerText = "Category: " + category;
}

Answer 1

您可以选择不同的模型，如下面的示例代码所示：

const createImageSegmenter = async () => {
const vision = await FilesetResolver.forVisionTasks("https://cdn.jsdelivr.net/npm/@mediapipe/[email protected]/wasm");
imageSegmenter = await ImageSegmenter.createFromOptions(vision, {
    baseOptions: {
        modelAssetPath: "./path/to/selfie_multiclass.tflite", // Download the selfie multiclass model from their website or get the CDN URL
        delegate: "GPU"
    },
    runningMode: "IMAGE",
    outputCategoryMask: true,
    outputConfidenceMasks: true
});
labels = imageSegmenter.getLabels(); // this gives all the labels present in the model
};

然后在您的代码中进行以下更改：

function callback(result) {
const cxt = canvasClick.getContext("2d");
const { width, height } = result.categoryMask;

let imageData = cxt.getImageData(0, 0, width, height).data;

canvasClick.width = width;
canvasClick.height = height;
let category = "";

const mask = result.categoryMask.getAsUint8Array();

for (let i in mask) {
    if (mask[i] > 0) {
        category = labels[mask[i]];
    }
    if (category=="hair" || category=="clothes") {   // draw segmentation masks only for the category you want
    const legendColor = legendColors[mask[i] % legendColors.length];
    imageData[i * 4] = (legendColor[0] + imageData[i * 4]) / 2;
    imageData[i * 4 + 1] = (legendColor[1] + imageData[i * 4 + 1]) / 2;
    imageData[i * 4 + 2] = (legendColor[2] + imageData[i * 4 + 2]) / 2;
    imageData[i * 4 + 3] = (legendColor[3] + imageData[i * 4 + 3]) / 2;
   }
}

const uint8Array = new Uint8ClampedArray(imageData.buffer);
const dataNew = new ImageData(uint8Array, width, height);
cxt.putImageData(dataNew, 0, 0);
const p = event.target.parentNode.getElementsByClassName("classification")[0];
p.classList.remove("removed");
p.innerText = "Category: " + category;
}

在 legendColors 变量中，您可以给出要在选定区域上绘制的蒙版颜色的 RGBA 值。

自拍多类模型中的类/标签如下：

• 0 - 背景
• 1 - 头发
• 2 - 身体皮肤
• 3 - 面部皮肤
• 4 - 衣服
• 5 - 其他（配件）

您可以在他们的网站上找到有关该型号的更多信息： [链接]