BodyPix - running toMask() and toColoredPartMask() in node.js throws an error: ImageData is not defined

Question

I am trying to get segmented person parts colored mask in TensorFlowJS bodypix model. Below code works fine before toColoredPartMask or toMask which throws an error "ImageData is not defined".

const tfjsnode = require('@tensorflow/tfjs-node')
const bodyPix = require('@tensorflow-models/body-pix');
const fs = require('fs');

setTimeout(async () => {
    maskImageWithBodyPix().then(response => {
        console.log(response)
    }).catch(e => {
        console.log("Error => " + e)
    })
}, 1000)

async function maskImageWithBodyPix(image = readImage("./person.jpeg")) {
    
    console.log("loadModel ...");
    if (image == null)
        return Promise.resolve("Image Not Found...")
    const resNet = {
        architecture: 'ResNet50',
        outputStride: 16,
        quantBytes: 4
    };
    let bodyModel = await bodyPix.load(resNet)

    console.log("segmentPersonParts ...");
    let segmentedPersonParts = await bodyModel.segmentPersonParts(image, {
        flipHorizontal: false,
        internalResolution: 'full',
        segmentationThreshold: 0.5,
    })
    console.log(`ImageHeight: ${segmentedPersonParts.height} | ImageWidth: ${segmentedPersonParts.width}`)
    console.log("toMaskImageData ...")
    const maskedImageData = bodyPix.toColoredPartMask(segmentedPersonParts, false);

    console.log(`maskedImageData = ${maskedImageData}`)
    
    return Promise.resolve(true)
}

const readImage = path => {
    console.log(`readImage ...`)
    if (!fs.existsSync(path))
        return null
    const imageBuffer = fs.readFileSync(path);
    const tfimage = tfjsnode.node.decodeImage(imageBuffer);
    return tfimage;
}

Answer 1

So this is a huge hack, but I got it working.

The problem lies in node_modules\@tensorflow-models\body-pix\dist\output_rendering_util.js.

If we look at this function inside:

function toColoredPartMask(partSegmentation, partColors) {
    if (partColors === void 0) { partColors = RAINBOW_PART_COLORS; }
    if (Array.isArray(partSegmentation) && partSegmentation.length === 0) {
        return null;
    }
    var multiPersonPartSegmentation;
    if (!Array.isArray(partSegmentation)) {
        multiPersonPartSegmentation = [partSegmentation];
    }
    else {
        multiPersonPartSegmentation = partSegmentation;
    }
    var _a = multiPersonPartSegmentation[0], width = _a.width, height = _a.height;
    var bytes = new Uint8ClampedArray(width * height * 4);
    for (var i = 0; i < height * width; ++i) {
        // invert mask.  Invert the segmentation mask.
        var j = i * 4;
        bytes[j + 0] = 255;
        bytes[j + 1] = 255;
        bytes[j + 2] = 255;
        bytes[j + 3] = 255;
        for (var k = 0; k < multiPersonPartSegmentation.length; k++) {
            var partId = multiPersonPartSegmentation[k].data[i];
            if (partId !== -1) {
                var color = partColors[partId];
                if (!color) {
                    throw new Error("No color could be found for part id " + partId);
                }
                bytes[j + 0] = color[0];
                bytes[j + 1] = color[1];
                bytes[j + 2] = color[2];
                bytes[j + 3] = 255;
            }
        }
    }
    return new ImageData(bytes, width, height);
}

It seemingly returns a new ImageData object. However, if we look at what ImageData actually is, we get this:

And if we dig deeper, we get to some node.js internal typescript files and see this:

declare var ImageData: {
    prototype: ImageData;
    new(sw: number, sh: number, settings?: ImageDataSettings): ImageData;
    new(data: Uint8ClampedArray, sw: number, sh?: number, settings?: ImageDataSettings): ImageData;
};

I thought that this won't work for us, so we need to overwrite that reference somehow.

What I did was install this package here and added const ImageData = require("@canvas/image-data"); line to the very top of the node_modules\@tensorflow-models\body-pix\dist\output_rendering_util.js file like this:

"use strict";
/**
 * @license
 * Copyright 2019 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
const ImageData = require("@canvas/image-data");
Object.defineProperty(exports, "__esModule", { value: true });
...

Now when we look at the ImageData again, we get this:

Just by doing that and also removing the false parameter from toColoredPartMask-call, so it's called like this:

const maskedImageData = bodyPix.toColoredPartMask(segmentedPersonParts);

We get this result:

If we were to JSON.stringify(maskedImageData), we'd get a billion lines of some numbers, so I'm not going to paste that here. Anyway, I think it works now.

I think the problem lies in the node.js implementation of ImageData. We're using some kind of weird browser implementation of it. The package I installed implements an outside of browser version of ImageData and we force bodypix to use that instead with the require.

Here is my input image:

And here is the output:

DO NOTE THAT YOU SHOULD BASICALLY NEVER EDIT ANYTHING INSIDE node_modules, BUT IF THIS IS JUST A PROJECT FOR FUN, I DON'T SEE WHY NOT PLAY AROUND A LITTLE

P.S I think the first line in toColoredPartMask that checks partColors is a bit funky, so I replaced it with if(arguments.length === 1 || !partColors) { partColors = RAINBOW_PART_COLORS; } in my code.