使用BodyPix和TensorFlow.js实现Color Pop效果

站长

2024年04月13日 18:55 · 阅读数 51

我最喜欢的谷歌照片应用程序的一个小功能是它的彩色弹出效果。色彩弹出(又名色彩飞溅)效果使主体(通常是一个人)从图像的其余部分中脱颖而出。主题仍然是彩色的，但背景是灰度的。在大多数情况下，这给人一种愉快的感觉。

使用BodyPix和TensorFlow.js实现Color Pop效果

虽然这个功能的效果非常好，但Google Photos只对一些它认为容易检测到人类的照片应用这个效果。这限制了它的潜力，并且不允许用户手动选择图像来应用此效果。这让我思考，有没有什么方法可以达到类似的效果，是我选择指定的图像?

大多数应用程序不提供自动化解决方案。它需要用户手动在图像上绘制这种效果，这既耗时又容易出错。我们能做得更好吗?像谷歌照片这样聪明的东西?是的!😉

如何手动实现此效果，我发现以下两个主要步骤:

在图像中的人周围创建一个遮罩(又名分割)。
使用蒙版来保留人物的颜色，同时使背景灰度化。

从图像中分割人物

这是这个过程中最重要的一步。一个好的结果在很大程度上取决于分割掩码的创建有多好。这一步需要一些机器学习，因为它已经被证明在这样情况下工作得很好。

从头开始构建和训练机器学习模型会花费太多时间，😛快速搜索后，我找到了BodyPix，这是一个用于人物分割和姿势检测的Tensorflow.js模型。

Tensorflow.js的BodyPix模型：

tfjs-models/body-pix at master · tensorflow/tfjs-models (github.com)

使用BodyPix和TensorFlow.js实现Color Pop效果

正如您所看到的，它可以很好地检测图像中的一个人(包括多个人)，并且在浏览器上运行相对较快。彩虹色的区域是我们需要的分割图。🌈

让我们用Tensorflow.js和BodyPix CDN脚本设置一个基本的HTML文件。

<html>  
<head>  
<title>Color Pop using Tensorflow.js and BodyPix</title>  
</head>  
  
<body>  
<!-- Canvas for input and output -->  
<canvas></canvas>  
  
    <!-- Load TensorFlow.js -->  
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.2"></script>  
    <!-- Load BodyPix -->  
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/body-pix@2.0"></script>  

    <!-- Color Pop code-->  
    <script src="colorpop.js"></script>  
  
</body>  
</html>

在画布中加载图像

在分割之前，了解如何在JavaScript中操作图像的像素数据是很重要的。一个简单的方法是使用HTML Canvas。Canvas使它易于读取和操作图像的像素数据，一旦加载。同时，它也兼容BodyPix，双赢!

function loadImage(src) {  
    const img = new Image();  
    const canvas = document.querySelector('canvas');  
    const ctx = canvas.getContext('2d');  

    // Load the image on canvas  
    img.addEventListener('load', () => {  
    // Set canvas width, height same as image  
    canvas.width = img.width;  
    canvas.height = img.height;  
    ctx.drawImage(img, 0, 0);  

    // TODO: Implement pop()  
    pop();  
    });  

    img.src = src;  
}

加载BodyPix模型

BodyPix的README很好地解释了如何使用模型。加载模型的一个重要部分是您使用的 architecture 。ResNet更准确但速度更慢，而MobileNet准确性较低但速度更快。在构建和测试这种效果时，我将使用MobileNet。稍后我将切换到ResNet并比较结果。

async function pop() {  
    // Loading the model  
    const net = await bodyPix.load({  
    architecture: 'MobileNetV1',  
    outputStride: 16,  
    multiplier: 0.75,  
    quantBytes: 2  
    });  
}

执行分割👥

BodyPix有多种功能来分割图像。有些适合身体部位分割，有些适合单人/多人分割。所有这些都在它们的README中有详细的解释。 segmentPerson() ，它在一个单独的地图中为图像中的每个人创建一个分割地图。而且，它比其他方法相对更快。

segmentPerson() 接受一个Canvas元素作为输入图像，以及一些配置设置。 internalResolution setting指定分割前调整输入图像大小的因子。我将使用 full 这个设置，因为我想要清晰的分割地图。

async function pop() {  
    // Loading the model  
    const net = await bodyPix.load({  
    architecture: 'MobileNetV1',  
    outputStride: 16,  
    multiplier: 0.75,  
    quantBytes: 2  
    });  
  
    // Segmentation  
    const canvas = document.querySelector('canvas');  
    const { data:map } = await net.segmentPerson(canvas, {  
    internalResolution: 'full',  
    });  
}

分割后的结果是一个对象(如下所示)。结果对象的主要部分是 data ，它是一个 Uint8Array ，将分割映射表示为一个数字数组

{  
    width: 640,  
    height: 480,  
    data: Uint8Array(307200) [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, …],  
    allPoses: [{"score": 0.4, "keypoints": […]}, …]  
}

制作背景灰度

准备好分割数据后，下一节使用分割实现颜色弹出，以使背景灰度化并保留图像中人物的颜色。为此，需要对图像进行像素级操作，而这正是canvas元素发挥作用的地方。 getImageData() 函数返回 ImageData ，其中包含RGBA格式的每个像素的颜色。

async function pop() {  
// ... previous code  
  
// Extracting image data  
const ctx = canvas.getContext('2d');  
const { data:imgData } = ctx.getImageData(0, 0, canvas.width, canvas.height);  
}

应用效果

所有的食材都准备好了。让我们从创建新的图像数据开始。 createImageData() 创建一个新的ImageData。

接下来，我们迭代映射中的像素，其中每个元素为1或0。

1表示在该像素处检测到一个人。
0表示在该像素处没有检测到人。

为了使代码更具可读性，我使用解构将颜色数据提取到 r g b a 变量中。

最后，基于分割图值（0或1），可以将灰度或实际RGBA颜色分配给新的图像数据。

每个像素处理后，使用 putImageData() 函数将新的图像数据绘制回画布上。

async function pop() {  
// ... previous code  
  
// Creating new image data  
const newImg = ctx.createImageData(canvas.width, canvas.height);  
const newImgData = newImg.data;  
  
// Apply the effect  
for(let i=0; i<map.length; i++) {  
// Extract data into r, g, b, a from imgData  
const [r, g, b, a] = [  
    imgData[i*4],  
    imgData[i*4+1],  
    imgData[i*4+2],  
    imgData[i*4+3]  
];  
  
// Calculate the gray color  
const gray = ((0.3 * r) + (0.59 * g) + (0.11 * b));  
  
// Set new RGB color to gray if map value is not 1  
// for the current pixel in iteration  
   [  
        newImgData[i*4],  
        newImgData[i*4+1],  
        newImgData[i*4+2],  
        newImgData[i*4+3]  
   ] = !map[i] ? [gray, gray, gray, 255] : [r, g, b, a];  
}  
  
// Draw the new image back to canvas  
ctx.putImageData(newImg, 0, 0);  
}

可以看到具有颜色弹出效果的最终图像应用于原始图像。好耶！🎉

探索其他架构和设置

我对ResNet和MobileNet架构进行了一些测试。在所有示例图像中，图像的一个尺寸(宽度或高度)的大小为1080px。注意，分割的内部分辨率设置为 full 。

在我的测试中，我在加载BodyPix模型时使用了以下设置。

// MobileNet architecture  
const net = await bodyPix.load({  
    architecture: 'MobileNetV1',  
    outputStride: 16,  
    quantBytes: 4,  
});  
  
  
// ResNet architecture  
const net = await bodyPix.load({  
    architecture: 'ResNet50',  
    outputStride: 16,  
    quantBytes: 2,  
});