I am following a Edx course named as Google AI for JavaScript developers with TensorFlow.js . Following the course there is a project where classification is taught using Mnist Dataset and the lecturer leaves it upto us if we want to extend the project where we could manually draw numbers on a canvas and see if the model can predict what number we wrote. I gave this project a try and I have implemented everything my model has a 0.97 accuracy and I am able to draw on canvas using JS and have converted the image to grayscale and then normalized and then resized it from 168 x 168 to 28 by 28. I have even included a debugging canvas to see if the converted image is correct and it displays the image what I drew but somehow I am unable to predict them correctly and I cant seem to figure out what is wrong with my code. here is the relevent function that is doing everything related to making predictions. Any help would be much appreciated, I have been stuck on this for 2 days now

async function makePrediction() {
let ctx = canvasRef.current.getContext("2d");
let imageData = ctx.getImageData(
  0,
  0,
  canvasRef.current.width,
  canvasRef.current.height
);
console.log("imageData is : ", imageData);

//converting to a tensor
let imageTensor = tf.browser.fromPixels(imageData);
console.log("imageTensor.shape :", imageTensor.shape);

//converting to grayscale
let greyscaleImageTensor = tf.image.rgbToGrayscale(imageTensor);
// let greyscaleImageTensor = tf.image.rgbToGrayscale(normalizedInput);
console.log("greyscaleImageTensor.shape : ", greyscaleImageTensor.shape);

//normalizing
let normalizedInput = greyscaleImageTensor.div(tf.scalar(255));

let resizedImageTensor = tf.image.resizeBilinear(
  normalizedInput,
  [28, 28],
  true
);
console.log("resizedImageTensor.shape : ", resizedImageTensor.shape);

// let newInput = tf.reshape(resizedImageTensor, [784]);
let newInput = resizedImageTensor.flatten();
console.log("newInput.shape: ", newInput.shape);
console.log("newInput : ", newInput);

// Debugging
let CTX = canvasDebug.current.getContext("2d");
// let testTensor = tf.reshape(imageTensor, [112896]);

let imageData2 = CTX.getImageData(0, 0, 28, 28);
console.log("imageData2.data.length : ", imageData2.data.length);
let data = await newInput.data();

console.log("data.length : ", data.length);

for (let i = 0; i < data.length; i++) {
  // console.log("data[i] : ", data[i]);
  imageData2.data[i * 4] = data[i] * 255;
  imageData2.data[i * 4 + 1] = data[i] * 255;
  imageData2.data[i * 4 + 2] = data[i] * 255;
  imageData2.data[i * 4 + 3] = 255;
}

CTX.putImageData(imageData2, 0, 0);
//END Debugging

let opt = model.predict(newInput.expandDims());
opt.print();

let answer = opt.squeeze().argMax();

let index = await answer.array();
setIsPrediction(true);
setPrediction(index);
if (index === 0) {
  setPrediction("0");
}

}
I am getting the training data from https://storage.googleapis.com/jmstore/TensorFlowJS/EdX/TrainingData/mnist.js

And since I followed the tutorial from the course so I know I trained the model correctly and it was making predictions correctly for data within the training data. So I know I am doing something wrong here but cant seem to figure out what cuz I could get the converted image display on debug canvas and the image appears to be correct