node fromPNG toPNG and generic fromPixels toPixels

I won’t use fromPixels and toPixels in node since It uses Cairo and I’m focused on Windows and Mac. Cairo on Windows apart of having to maintain additional binaries it’s not as easy to install as other OS.

If you are like me that want to keep the binaries to the bare minimum here is a very fast fromPNG and toPNG for node purely using js.

npm install pngjs

const fs      = require('fs'),
const PNG  = require('pngjs').PNG; 

async function fromPNG(path, numChannels)
{
    numChannels = numChannels || 3;
    
    return new Promise((resolve, reject) => 
    {
        fs.createReadStream(path).pipe(
        new PNG({colorType: 6} ))
        .on('parsed', function() 
        {
            resolve( fromPixels(this.data, this.width, this.height, numChannels) );
        })
        .on('error', reject);
    });
}

function toPNG(tensor, path)
{
    let dst     = new PNG({width: tensor.shape[0], height: tensor.shape[1]});
    dst.data    = toPixels(tensor);
    dst.pack().pipe( fs.createWriteStream(path) );
}

function fromPixels(data, width, height, numChannels)
{
    const values    = [];
    const numPix    = width * height;
    
    for (let i = 0; i < numPix; i++) 
        for (let channel = 0; channel < numChannels; ++channel)
            values[i * numChannels + channel] = data[i * 4 + channel];

    let tensor = tf.tensor3d(values, [width, height, numChannels], 'int32');
    
    return tensor;
}

function toPixels(tensor) 
{
    const pixels    = tensor.dataSync();
    const data      = [];
    
    if (tensor.shape.length == 2 || tensor.shape[2] == 1) 
    {
        // Grayscale
        for (let i = 0; i < pixels.length; i++) {
          data[i*4+0] = pixels[i];
          data[i*4+1] = pixels[i];
          data[i*4+2] = pixels[i];
          data[i*4+3] = 255;
        }
        
    } else if (tensor.shape[2] == 3) 
    {
        // RGB
        for (let i = 0; i < pixels.length / 3; i++) 
        {
          data[i*4+0] = pixels[i*3+0];
          data[i*4+1] = pixels[i*3+1];
          data[i*4+2] = pixels[i*3+2];
          data[i*4+3] = 255;
        }
        
    }else if (tensor.shape[2] == 4) 
    {
        // RGBA
        data = pixels;
    }
    
    return data;
}

JPG could be the same using node-jpg. Of course like fromPixels or toPixels assumes that your shapes are going to be [W, H, Channels] maybe you have to swap them if your shapes are [Channels, W, H]. Also you have data on the range [0, 255] so don’t forget to preprocess and postprocess to fit your activations.

[davidsoergel]

Closing based on @nkreeger’s comment. Thanks for the helpful feedback & code, though!

[nkreeger]

@davidsoergel @DavidGOrtega cool nice work - we’re going to be shipping the generated tf.fromPixels() API soon. That’s the official one we want to support since it uses the TF C-library to convert JPEG/PNG/etc byte streams to Tensors in C-land.