seanchas116 / node-tflite

TensorFlow Lite bindings for Node.js
https://www.npmjs.com/package/node-tflite
MIT License
14 stars 6 forks source link

Usage with tflite file from Lobe.ai #6

Closed jlarmstrongiv closed 3 years ago

jlarmstrongiv commented 3 years ago

こんばんは 👋 hi

I’m trying to use node-tflite to run a tflite model generated from Lobe.ai. I’m just not sure how to get it running.

I have looked at the documentation in the readme and the testing script for inspiration.

I had problems with size mismatch, but solved that by changing the UintArray size. But, I think haphazardly doing what I did breaks the final classification. I’m just unsure how to work with float32.

Here is a Google Drive link to a sample tflite export (the file is too big to be uploaded to GitHub). The sample is just a demo classification that classifies whether an image is a rubber duck or a bunch of rocks.

Anyway, really love this package! Can’t wait to get it to work :)

P.S. I also considered using the tfjs converter, but it throws an error too.

seanchas116 commented 3 years ago

Thank you for trying out node-tflite!

I tried your model and looks like this works:

import fs from "fs";
import path from "path";
import { createCanvas, loadImage } from "canvas";
import { Interpreter } from "node-tflite";

const modelPath = path.resolve(__dirname, "saved_model.tflite");
//const imagePath = path.resolve(__dirname, "rock/optimized-3.4s-1-1.2RiverRock-1.jpg");
const imagePath = path.resolve(__dirname, "rubber-duck/optimized-01-rubberduck-hongkong.jpg");

function createInterpreter() {
  const modelData = fs.readFileSync(modelPath);
  return new Interpreter(modelData);
}

async function getImageInput(size: number) {
  const canvas = createCanvas(size, size);
  const context = canvas.getContext("2d");
  const image = await loadImage(imagePath);
  context.drawImage(image, 0, 0, size, size);
  const data = context.getImageData(0, 0, size, size);

  const inputData = new Float32Array(size * size * 3); // Use Float32Array instead of Uint8Array

  for (let i = 0; i < size * size; ++i) {
    inputData[i * 3] = data.data[i * 4] / 255; // Convert 0...255 int to 0...1 float
    inputData[i * 3 + 1] = data.data[i * 4 + 1] / 255;
    inputData[i * 3 + 2] = data.data[i * 4 + 2] / 255;
  }

  return inputData;
}

async function invoke() {
    const interpreter = createInterpreter();
    interpreter.allocateTensors();

    const inputData = await getImageInput(224);
    interpreter.inputs[0].copyFrom(inputData);

    interpreter.invoke();

    // get string output
    const outputData = Buffer.alloc(interpreter.outputs[0].byteSize);
    interpreter.outputs[0].copyTo(outputData);
    console.log(outputData.toString())
}

invoke()

The points are:

I hope this helps!

jlarmstrongiv commented 3 years ago

Wow!!!! Thank you so much 😄 helps a ton