@codait/max-image-segmenter
Version:
Identify objects in an image, additionally assigning each pixel of the image to a particular object.
466 lines (446 loc) • 9.31 kB
JavaScript
'use strict';Object.defineProperty(exports,'__esModule',{value:true});/* globals tf, Image */
const IMAGESIZE = 512;
const computeTargetSize = function (width, height) {
const resizeRatio = IMAGESIZE / Math.max(width, height);
return {
width: Math.round(resizeRatio * width),
height: Math.round(resizeRatio * height)
}
};
const getImageData = function (imageInput) {
{
return Promise.resolve(imageInput)
}
};
const imageToTensor = function (imageData) {
return tf.tidy(() => {
const imgTensor = tf.browser.fromPixels(imageData);
const targetSize = computeTargetSize(imgTensor.shape[0], imgTensor.shape[1]);
return imgTensor.resizeBilinear([targetSize.width, targetSize.height]).expandDims()
})
};
/**
* convert image to Tensor input required by the model
*
* @param {HTMLImageElement} imageInput - the image element
*/
const preprocess = function (imageInput) {
return getImageData(imageInput)
.then(imageToTensor)
.then(inputTensor => {
return Promise.resolve(inputTensor)
})
.catch(err => {
console.error(err);
return Promise.reject(err)
})
};/* globals tf */
let modelPath = null;
{
modelPath = `file://${__dirname}/../model/model.json`;
}
let model = null;
let warmed = false;
/**
* load the image segmenter model
*/
const load = function (initialize) {
if (!model) {
// console.log('loading model...')
// console.time('model load')
return tf.loadGraphModel(modelPath)
.then(m => {
// console.timeEnd('model load')
model = m;
if (istrue(initialize)) {
warmup();
}
return Promise.resolve(model)
})
.catch(err => {
// console.timeEnd('model load')
console.error(err);
return Promise.reject(err)
})
} else if (istrue(initialize) && !warmed) {
warmup();
return Promise.resolve(model)
} else {
return Promise.resolve(model)
}
};
/**
* run the model to get a prediction
*/
const run = function (imageTensor) {
if (!imageTensor) {
console.error('no image provided');
throw new Error('no image provided')
} else if (!model) {
console.error('model not available');
throw new Error('model not available')
} else {
// console.log('running model...')
return tf.tidy(() => {
// console.time('model inference')
const results = model.predict(imageTensor.toInt());
// console.timeEnd('model inference')
warmed = true;
return results
})
}
};
/**
* run inference on the TensorFlow.js model
*/
const inference = function (imageTensor) {
return load(false).then(() => {
try {
const results = run(imageTensor);
return Promise.resolve(results)
} catch (err) {
return Promise.reject(err)
}
})
};
const warmup = function () {
try {
run(tf.zeros([1, 512, 512, 3]));
} catch (err) { }
};
const istrue = function (param) {
return param === null ||
typeof param === 'undefined' ||
(typeof param === 'string' && param.toLowerCase() === 'true') ||
(typeof param === 'boolean' && param)
};const labels = [
'background',
'airplane',
'bicycle',
'bird',
'boat',
'bottle',
'bus',
'car',
'cat',
'chair',
'cow',
'dining table',
'dog',
'horse',
'motorbike',
'person',
'potted plant',
'sheep',
'sofa',
'train',
'tv'
];
const colors = [
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[],
[]
];const predictedObjs = function (segArray) {
const segLabels = {};
segArray.forEach(arr => {
arr.forEach(seg => {
if (!segLabels[labels[seg]]) {
segLabels[labels[seg]] = true;
}
});
});
return Object.keys(segLabels)
};
/**
* convert model Tensor output to image data for previewing
*
* @param {Tensor} inferenceResults - the output from running the model
*/
const postprocess = function (inferenceResults) {
return inferenceResults.unstack()[0].array()
.then(segArray => {
return Promise.resolve({
segmentationMap: segArray,
objectsDetected: predictedObjs(segArray),
imageSize: {
width: segArray[0].length,
height: segArray.length
}
})
})
};const version="0.3.0";{
global.tf = require('@tensorflow/tfjs-node');
}
const processInput = function (inputImage) {
return preprocess(inputImage)
};
const loadModel = function (init) {
return load(init)
};
const runInference = function (inputTensor) {
return inference(inputTensor)
};
const processOutput = function (inferenceResults) {
return postprocess(inferenceResults)
};
const predict = function (inputImage) {
return processInput(inputImage)
.then(runInference)
.then(processOutput)
.catch(err => {
console.error(err);
})
};exports.colorsMap=colors;exports.labelsMap=labels;exports.loadModel=loadModel;exports.predict=predict;exports.processInput=processInput;exports.processOutput=processOutput;exports.runInference=runInference;exports.version=version;