const ort = require("onnxruntime-node");
const sharp = require("sharp");

async function loadImageAndResize(imagePath) {
  const image = await sharp(imagePath).resize(299, 299).raw().toBuffer();

  const normalizedImage = new Float32Array(1 * 299 * 299 * 3);
  for (let i = 0; i < image.length; i++) {
    normalizedImage[i] = image[i] / 255.0;
  }

  return normalizedImage;
}

async function runModel(imagePath, modelPath) {
  const startTime = new Date().getTime();
  const normalizedImage = await loadImageAndResize(imagePath);
  const loadImgTime = new Date().getTime();
  console.log(`Load image time: ${loadImgTime - startTime} ms`);
  const session = await ort.InferenceSession.create(modelPath);
  const modelLoadTime = new Date().getTime();
  console.log(`Load model time: ${modelLoadTime - loadImgTime} ms`);

  const inputName = session.inputNames[0];

  const options = {
    [inputName]: new ort.Tensor("float32", normalizedImage, [1, 299, 299, 3]),
  };
  const makeTensorTime = new Date().getTime();
  console.log(`Make tensor time: ${makeTensorTime - modelLoadTime} ms`);
  const feeds = {};
  feeds[inputName] = options[inputName];

  const results = await session.run(feeds);
  const runModelTime = new Date().getTime();
  console.log(`Run model time: ${runModelTime - makeTensorTime} ms`);

  const outputName = session.outputNames[0];
  const output = results[outputName];
  const categories = ["drawings", "hentai", "neutral", "porn", "sexy"];
  const sortedIndices = output.data
    .map((value, index) => index)
    .sort((a, b) => output.data[b] - output.data[a]);
  const imagePreds = {};

  imagePreds[imagePath] = {};
  for (const index of sortedIndices) {
    imagePreds[imagePath][categories[index]] = output.data[index].toString();
  }

  console.log(JSON.stringify(imagePreds, null, 2));
}

const imagePath = "./images/mnzl.jpg";
const modelPath = "./model.onnx";

runModel(imagePath, modelPath);