103 lines
3.0 KiB
JavaScript
103 lines
3.0 KiB
JavaScript
const { PNG } = require('pngjs');
|
|
const fs = require('fs');
|
|
|
|
async function segmentationMask(input1, input2, output, threshold = 0.02) {
|
|
let img1 = await loadPngFile(input1);
|
|
let img2 = await loadPngFile(input2);
|
|
|
|
if (img1.width !== img2.width || img1.height !== img2.height) {
|
|
throw new Error("Cannot compute mask on images with different sizes");
|
|
}
|
|
|
|
for (let i = 0; i < img1.data.length; i += 4) {
|
|
let r1 = img1.data[i + 0] / 255;
|
|
let g1 = img1.data[i + 1] / 255;
|
|
let b1 = img1.data[i + 2] / 255;
|
|
|
|
let r2 = img2.data[i + 0] / 255;
|
|
let g2 = img2.data[i + 1] / 255;
|
|
let b2 = img2.data[i + 2] / 255;
|
|
|
|
// Test difference
|
|
let difference = Math.sqrt((r1 - r2) * (r1 - r2) + (g1 - g2) * (g1 - g2) + (b1 - b2) * (b1 - b2));
|
|
let pixelsAreDifferent = difference > threshold;
|
|
|
|
if (pixelsAreDifferent) {
|
|
img1.data[i + 0] = 255;
|
|
img1.data[i + 1] = 255;
|
|
img1.data[i + 2] = 255;
|
|
} else {
|
|
img1.data[i + 0] = 0;
|
|
img1.data[i + 1] = 0;
|
|
img1.data[i + 2] = 0;
|
|
}
|
|
}
|
|
|
|
let outputStream = fs.createWriteStream(output);
|
|
await img1.pack().pipe(outputStream);
|
|
}
|
|
|
|
// The following is fixed from https://www.npmjs.com/package/png-quality
|
|
|
|
async function loadPngFile(pathOrBuffer) {
|
|
// Load buffer of path
|
|
if (!(pathOrBuffer instanceof Buffer)) {
|
|
pathOrBuffer = await new Promise((resolve, reject) => {
|
|
fs.readFile(pathOrBuffer, (err, data) => {
|
|
if (err) return reject(err)
|
|
resolve(data)
|
|
});
|
|
})
|
|
}
|
|
|
|
// Load PNG from buffer
|
|
return await new Promise((resolve, reject) => {
|
|
const png = new PNG()
|
|
png.parse(pathOrBuffer, err => {
|
|
if (err) return reject(err)
|
|
resolve(png)
|
|
})
|
|
})
|
|
}
|
|
|
|
async function mse(png1, png2) {
|
|
const pngs = [png1, png2];
|
|
for (let i in pngs) {
|
|
if (!(pngs[i] instanceof PNG)) pngs[i] = await loadPngFile(pngs[i])
|
|
}
|
|
|
|
if (pngs[0].width !== pngs[1].width || pngs[0].height !== pngs[1].height) {
|
|
throw new Error('Width or height does not equal')
|
|
}
|
|
|
|
if (pngs[0].data.length !== pngs[1].data.length) {
|
|
throw new Error('Data buffer length does not equal')
|
|
}
|
|
|
|
const square = (a) => a * a,
|
|
channelIndex = [0, 1, 2],
|
|
channelMax = 255 * 255,
|
|
area = pngs[0].width * pngs[1].height
|
|
|
|
let mse = 0
|
|
for (let i = 0; i < pngs[0].data.length; i += 4) {
|
|
const rgbas = pngs.map(png => png.data.slice(i, i + 4))
|
|
const rgbs = rgbas.map(rgba => channelIndex.map(i => rgba[i] * rgba[3]))
|
|
channelIndex.forEach(i => mse += square(rgbs[0][i] - rgbs[1][i]))
|
|
}
|
|
|
|
return mse / 3.0 / (channelMax * channelMax) / area
|
|
}
|
|
|
|
async function psnr(png1, png2) {
|
|
const m = await mse(png1, png2)
|
|
return 10 * Math.log10(1 / m)
|
|
}
|
|
|
|
module.exports = {
|
|
loadPngFile,
|
|
mse,
|
|
psnr,
|
|
segmentationMask,
|
|
};
|