Inconsistent prediction results when adapting Model to ONNX runtime Node

[mau-io]

Issue Description

I am currently trying to adapt the JoyTag model from this repository (https://github.com/fpgaminer/joytag) for use with onnxruntime-node. However, I'm encountering an issue where the predicted tags from the model do not match the expected results for given images.

Expected Tags

For a football match image, the expected tags based on the model should closely align with terms like short_hair, brown_hair, black_hair, photoshop_(medium), standing, male_focus, multiple_boys, shorts, dark_skin, blurry, tattoo, facial_hair, grass, dark-skinned_male, beard, ball, 6+boys, third-party_edit, 4boys, sportswear, emblem, mustache, motion_blur, photo_(medium), logo, 5boys, bald, 3d, real_life, soccer_uniform, crowd, soccer_ball, etc., encompassing various attributes and actions observable in the image.

Actual Tags

The tags received are unrelated and nonsensical, not matching the expected context of the image.

Code Example

Below is the main code where I handle image preprocessing and make predictions using the ONNX model:

import ort from 'onnxruntime-node';
import fs from 'fs';
import Jimp from 'Jimp';

const MODEL_PATH = '../joytag/model.onnx';
const IMAGE_PATH = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/football-match.jpg';
const THRESHOLD = 0.4;

const topTags = fs.readFileSync(`../joytag/top_tags.txt`, 'utf8')
  .split('\n')
  .filter(line => line.trim())
  .map(line => line.trim());

  async function prepareImage(imagePath, targetSize) {

    const image = await Jimp.read(imagePath);

    const { width, height } = image.bitmap;
    const maxDim = Math.max(width, height);
    const padLeft = Math.floor((maxDim - width) / 2);
    const padTop = Math.floor((maxDim - height) / 2);

    let paddedImage = new Jimp(maxDim, maxDim, 0xFFFFFFFF);
    paddedImage.composite(image, padLeft, padTop);

    if (maxDim !== targetSize) {
      paddedImage = await paddedImage.resize(targetSize, targetSize, Jimp.RESIZE_BICUBIC);
    }

    const imageTensor = new Float32Array(3 * targetSize * targetSize);
    await paddedImage.scan(0, 0, targetSize, targetSize, function (x, y, idx) {
        const pos = (y * targetSize + x) * 3;
        imageTensor[pos] = this.bitmap.data[idx + 0] / 255.0;       // R
        imageTensor[pos + 1] = this.bitmap.data[idx + 1] / 255.0; // G
        imageTensor[pos + 2] = this.bitmap.data[idx + 2] / 255.0; // B
    });

    const mean = [0.485, 0.456, 0.406];
    const std = [0.229, 0.224, 0.225];
    for (let i = 0; i < imageTensor.length; i += 3) {
        imageTensor[i] = (imageTensor[i] - mean[0]) / std[0]; // R
        imageTensor[i + 1] = (imageTensor[i + 1] - mean[1]) / std[1]; // G
        imageTensor[i + 2] = (imageTensor[i + 2] - mean[2]) / std[2]; // B
    }

    const tensor = new ort.Tensor('float32', imageTensor, [1, 3, targetSize, targetSize]);
    return tensor;
}

async function predict(imagePath) {
  const imageTensor = await prepareImage(imagePath, 448);
  console.log(imageTensor)
  const session = await ort.InferenceSession.create(MODEL_PATH);
  const feeds = { input: imageTensor };
  const results = await session.run(feeds);
  console.log(results)
  const output = results.output.data;

  // Helper function to apply the sigmoid function
  function sigmoid(x) {
    return 1 / (1 + Math.exp(-x));
  }
  // Apply sigmoid to each output score
  const probabilities = output.map(sigmoid);
  const scores = {};
  probabilities.forEach((prob, i) => {
    scores[topTags[i]] = prob;
  });
  const predictedTags = Object.keys(scores).filter(tag => scores[tag] > THRESHOLD);
  const tagString = predictedTags.join(', ');
  return { tagString, scores, predictedTags };
}

async function main() {
  const { tagString, scores, predictedTags } = await predict(IMAGE_PATH);
  console.log({predictedTags})
}

main();

Request for Assistance:

Any guidance or suggestions on why the tag predictions might be diverging so significantly would be greatly appreciated.

Environment

• ONNX Runtime Version: 18.12.1
• Node.js Version: 1.17.3