tensorflowjs-binarization/src/threshold.js

Realization of binarization algorithms with Tensorflowjs

//let tf = require('@tensorflow/tfjs-node-gpu');

let threshold = function (tf) {
	
	return {
		
		binary: async function binary(imageTensor, coeff, ...inv){
			try{
				let threshold = coeff*255;
				let img_shape = imageTensor.shape;
				let rank = parseInt(imageTensor.rankType, 10);
				if(imageTensor.shape[3]>1) throw new ChannelsException();
				if(rank!=2 && rank!=3) throw new RankException(rank);
		
		//if(imageTensor.dType != 'float32') imageTensor = imageTensor.div(tf.scalar(255));
	
				let arrayed_image = Array.from(imageTensor.dataSync());
	
				arrayed_image.forEach((item, i) =>{ 
					if (!inv[0]){
					arrayed_image[i] = (item < threshold) ? 0 : 255;
					}else arrayed_image[i] = (item > threshold) ? 0 : 255;
				});
				console.log(imageTensor.dtype)
				return tf.tensor3d(arrayed_image, img_shape);
			}
			catch(error){
				console.log(error)
			}
},
		
		otsu: async function otsu(imageTensor){
	try{
		let img_shape = imageTensor.shape;
		let rank = parseInt(imageTensor.rankType, 10);
		if(imageTensor.shape[3]>1) throw new ChannelsException();
		if(rank!=2 && rank!=3) throw new RankException(rank);
		
		//if(imageTensor.dType != 'float32') imageTensor = imageTensor.div(tf.scalar(255));
		
		let histogram = Array(256).fill(0);
		let arrayed_image = Array.from(imageTensor.dataSync());
		
		for (i = 0; i < arrayed_image.length; i++) {

			let gray = arrayed_image[i]
			histogram[Math.round(gray)] += 1;
  
		}
		
		let threshold = otsu_alg(histogram, arrayed_image.length);
	console.log(`threshold = ${threshold}`);
		
		arrayed_image.forEach(function(item, i) { 
			arrayed_image[i] = (item < threshold) ? 0 : 255;
		});


		return tf.tensor(arrayed_image, img_shape);

	}
	catch(error){
		console.log(error)
	}
}
		
	}
    this.otsu = otsu;
    this.binary = binary;
}

module.exports = threshold;



function otsu_alg(histData /* Array of 256 greyscale values */, total /* Total number of pixels */) {
	
	let sum = 0;
    for (let t=0 ; t<256 ; t++) sum += t * histData[t];

    let sumB = 0;
    let wB = 0;
    let wF = 0;

    let varMax = 0;
    let threshold = 0;

    for (let t=0 ; t<256 ; t++) {
        wB += histData[t];               // Weight Background
        if (wB == 0) continue;
		
        wF = total - wB;                 // Weight Foreground
        if (wF == 0) break;

        sumB += t * histData[t];

        let mB = sumB / wB;            // Mean Background
        let mF = (sum - sumB) / wF;    // Mean Foreground

        // Calculate Between Class Variance
        let varBetween = wB * wF * (mB - mF) * (mB - mF);

        // Check if new maximum found
        if (varBetween > varMax) {
            varMax = varBetween;
            threshold = t;
        }
    }

    return threshold;
}












function RankException(rank_number) {
   let message = `Number of dimensions of image should be equal to 2 or 3, got ${rank_number}`;
   const error = new Error(message);
   return error;
}

function ChannelsException() {
   let message = `Expected grayscale image}`;
   const error = new Error(message);
   return error;
}