tensorflow / tfjs

A WebGL accelerated JavaScript library for training and deploying ML models.
https://js.tensorflow.org
Apache License 2.0
18.34k stars 1.92k forks source link

Node.JS wrapper for EfficientDet #7867

Open lukemovement opened 1 year ago

lukemovement commented 1 year ago

I have recently written a wrapper for EfficientDet but I don't really have the time to oversee getting it added as a PR due to personal commitments. The code seems to be fully functional. The input size of the model can be changed by replacing all instances of 1536 within the sample.

Feel free to use or close the ticket :)

import * as tf from "@tensorflow/tfjs-node";
import { resolve } from "path";
import { inspect } from "util";

export class EfficientDet {
  labels = {
    1: "person",
    2: "bicycle",
    3: "car",
    4: "motorcycle",
    5: "airplane",
    6: "bus",
    7: "train",
    8: "truck",
    9: "boat",
    10: "traffic light",
    11: "fire hydrant",
    13: "stop sign",
    14: "parking meter",
    15: "bench",
    16: "bird",
    17: "cat",
    18: "dog",
    19: "horse",
    20: "sheep",
    21: "cow",
    22: "elephant",
    23: "bear",
    24: "zebra",
    25: "giraffe",
    27: "backpack",
    28: "umbrella",
    31: "handbag",
    32: "tie",
    33: "suitcase",
    34: "frisbee",
    35: "skis",
    36: "snowboard",
    37: "sports ball",
    38: "kite",
    39: "baseball bat",
    40: "baseball glove",
    41: "skateboard",
    42: "surfboard",
    43: "tennis racket",
    44: "bottle",
    46: "wine glass",
    47: "cup",
    48: "fork",
    49: "knife",
    50: "spoon",
    51: "bowl",
    52: "banana",
    53: "apple",
    54: "sandwich",
    55: "orange",
    56: "broccoli",
    57: "carrot",
    58: "hot dog",
    59: "pizza",
    60: "donut",
    61: "cake",
    62: "chair",
    63: "couch",
    64: "potted plant",
    65: "bed",
    67: "dining table",
    70: "toilet",
    72: "tv",
    73: "laptop",
    74: "mouse",
    75: "remote",
    76: "keyboard",
    77: "cell phone",
    78: "microwave",
    79: "oven",
    80: "toaster",
    81: "sink",
    82: "refrigerator",
    84: "book",
    85: "clock",
    86: "vase",
    87: "scissors",
    88: "teddy bear",
    89: "hair drier",
    90: "toothbrush",
  };

  private model?: tf.GraphModel;

  public async loadModel() {
    this.model = await tf.loadGraphModel(
      `https://tfhub.dev/tensorflow/efficientdet/d7/1`,
    );
    console.log("Model loaded successfully!");
  }

  public async predict(imageData: Uint8Array | tf.Tensor3D): Promise<any> {
    if (!this.model) {
      throw new Error("Model not loaded");
    }

    const inputTensor = this.preprocessInput(imageData);
    const predictions = (await this.model.executeAsync(inputTensor, [
      "detection_boxes",
      "detection_multiclass_scores",
    ])) as tf.Tensor[];

    const [detection_boxes, detection_multiclass_scores] = predictions;

    const detectionBoxesData = (
      (await detection_boxes.array()) as number[][][]
    )[0];
    const detectionMulticlassScoresData = (
      (await detection_multiclass_scores.array()) as number[][][]
    )[0];

    const detections: Array<{
      label: string;
      score: number;
      box: {
        originX: number;
        originY: number;
        width: number;
        height: number;
      };
    }> = [];

    for (let i = 0; i < detectionMulticlassScoresData.length; i++) {
      const scores = detectionMulticlassScoresData[i];
      const maxScore = Math.max(...scores);
      const maxScoreIndex = scores.indexOf(maxScore);

      const label =
        this.labels[(maxScoreIndex + 1) as keyof typeof this.labels];
      const score = maxScore;

      let originX = detectionBoxesData[i][1];
      let originY = detectionBoxesData[i][0];
      let width = detectionBoxesData[i][3] - originX;
      let height = detectionBoxesData[i][2] - originY;

      originX = originX * 1536;
      originY = originY * 1536;
      width = width * 1536;
      height = height * 1536;

      detections.push({
        label,
        score,
        box: {
          originX: Math.round(originX),
          originY: Math.round(originY),
          width: Math.round(width),
          height: Math.round(height),
        },
      });
    }

    return detections;
  }

  private preprocessInput(imageTensor: Uint8Array | tf.Tensor3D): tf.Tensor3D {
    if (imageTensor instanceof Uint8Array) {
      imageTensor = tf.node.decodeImage(imageTensor, 3) as tf.Tensor3D;
    }

    const resizedTensor = tf.image.resizeBilinear(imageTensor, [1536, 1536]);
    const batchedTensor = resizedTensor.expandDims(0) as tf.Tensor3D;
    return batchedTensor.toInt();
  }

  async drawBoxesOnImage(
    imageTensor: tf.Tensor3D | Uint8Array,
    detections: Array<{
      label: string;
      score: number;
      box: {
        originX: number;
        originY: number;
        width: number;
        height: number;
      };
    }>,
  ): Promise<tf.Tensor3D> {
    if (imageTensor instanceof Uint8Array) {
      imageTensor = tf.node.decodeImage(imageTensor, 3) as tf.Tensor3D;
    }

    imageTensor = tf.image.resizeBilinear(imageTensor, [1536, 1536]);

    const imageArray = await imageTensor.data();
    console.log("drawing boxes on image");
    for (const detection of detections) {
      const { originX, originY, width, height } = detection.box;

      for (let i = originX; i < originX + width; i++) {
        for (let j = originY; j < originY + height; j++) {
          // only fill the border
          if (
            i !== originX &&
            i !== originX + width - 1 &&
            j !== originY &&
            j !== originY + height - 1
          ) {
            continue;
          }

          imageArray[j * 1536 * 3 + i * 3] = 0;
          imageArray[j * 1536 * 3 + i * 3 + 1] = 0;
          imageArray[j * 1536 * 3 + i * 3 + 2] = 255;
        }
      }
    }

    const imageTensor2 = tf.tensor(imageArray, [1536, 1536, 3], "int32");

    return imageTensor2 as tf.Tensor3D;
  }
}

import { promises as fs } from "fs";

// get object detection result
const imagePath = resolve(process.cwd(), "sample_in.jpg");
const image = await fs.readFile(imagePath);

const model = new EfficientDet();
await model.loadModel();

const result = await model.predict(image);

console.log(inspect(result, false, null, true));

// draw boxes on image
const imageTensor = tf.node.decodeImage(image, 3) as tf.Tensor3D;

const highlightedImageTensor = await model.drawBoxesOnImage(
  imageTensor,
  result,
);

// write image to file
const output = await tf.node.encodeJpeg(highlightedImageTensor as tf.Tensor3D);
await fs.writeFile(resolve(process.cwd(), "sample_out.jpg"), output);
gaikwadrahul8 commented 1 year ago

Hi, @lukemovement

Thank you for writing Node.JS wrapper for EfficientDet and I was trying to run your code in my Mac system with random person color jpg image from the internet and converted index.tsfile to index.jswith the help of tsc after that I ran node index.js I'm getting below error messages, If have I missed something here please let me know ? or please guide me with the steps to run your above Node.JS wrapper for EfficientDet code. Thank you!

Here is error log for your reference :

gaikwadrahul-macbookpro:test-7867 gaikwadrahul$ node index.js
/Users/gaikwadrahul/Desktop/TFJS/test-7867/index.js:250
var image = await fs_1.promises.readFile(imagePath);
            ^^^^^

SyntaxError: await is only valid in async functions and the top level bodies of modules
    at internalCompileFunction (node:internal/vm:73:18)
    at wrapSafe (node:internal/modules/cjs/loader:1176:20)
    at Module._compile (node:internal/modules/cjs/loader:1218:27)
    at Module._extensions..js (node:internal/modules/cjs/loader:1308:10)
    at Module.load (node:internal/modules/cjs/loader:1117:32)
    at Module._load (node:internal/modules/cjs/loader:958:12)
    at Function.executeUserEntryPoint [as runMain] (node:internal/modules/run_main:81:12)
    at node:internal/main/run_main_module:23:47

Node.js v18.15.0
gaikwadrahul-macbookpro:test-7867 gaikwadrahul$ 
lukemovement commented 1 year ago

Hi, @lukemovement

Thank you for writing Node.JS wrapper for EfficientDet and I was trying to run your code in my Mac system with random person color jpg image from the internet and converted index.tsfile to index.jswith the help of tsc after that I ran node index.js I'm getting below error messages, If have I missed something here please let me know ? or please guide me with the steps to run your above Node.JS wrapper for EfficientDet code. Thank you!

Here is error log for your reference :

gaikwadrahul-macbookpro:test-7867 gaikwadrahul$ node index.js
/Users/gaikwadrahul/Desktop/TFJS/test-7867/index.js:250
var image = await fs_1.promises.readFile(imagePath);
            ^^^^^

SyntaxError: await is only valid in async functions and the top level bodies of modules
    at internalCompileFunction (node:internal/vm:73:18)
    at wrapSafe (node:internal/modules/cjs/loader:1176:20)
    at Module._compile (node:internal/modules/cjs/loader:1218:27)
    at Module._extensions..js (node:internal/modules/cjs/loader:1308:10)
    at Module.load (node:internal/modules/cjs/loader:1117:32)
    at Module._load (node:internal/modules/cjs/loader:958:12)
    at Function.executeUserEntryPoint [as runMain] (node:internal/modules/run_main:81:12)
    at node:internal/main/run_main_module:23:47

Node.js v18.15.0
gaikwadrahul-macbookpro:test-7867 gaikwadrahul$ 

It sounds like you are using an older version of NodeJS than me. Use this to test the model instead.

const main = async () => {
  // get object detection result
  const imagePath = resolve(process.cwd(), "sample_in.jpg");
  const image = await fs.readFile(imagePath);

  const model = new EfficientDet();
  await model.loadModel();

  const result = await model.predict(image);

  // draw boxes on image
  const imageTensor = tf.node.decodeImage(image, 3) as tf.Tensor3D;

  const highlightedImageTensor = await model.drawBoxesOnImage(
    imageTensor,
    result,
  );

  // write image to file
  const output = await tf.node.encodeJpeg(
    highlightedImageTensor as tf.Tensor3D,
  );
  await fs.writeFile(resolve(process.cwd(), "sample_out.jpg"), output);
};

main();
gaikwadrahul8 commented 1 year ago

Hi, @lukemovement

Thank you for sharing your code and I tried your first code snippet with latest version of Node.js v18.17.0, I'm still getting the same error message so now I'll assign this issue to our concerned team, they'll take look and appropriate action from their end. Thank you!

gaikwadrahul-macbookpro:test-7867 gaikwadrahul$ node index.js
/Users/gaikwadrahul/Desktop/TFJS/test-7867/index.js:250
var image = await fs_1.promises.readFile(imagePath);
            ^^^^^

SyntaxError: await is only valid in async functions and the top level bodies of modules
    at internalCompileFunction (node:internal/vm:73:18)
    at wrapSafe (node:internal/modules/cjs/loader:1178:20)
    at Module._compile (node:internal/modules/cjs/loader:1220:27)
    at Module._extensions..js (node:internal/modules/cjs/loader:1310:10)
    at Module.load (node:internal/modules/cjs/loader:1119:32)
    at Module._load (node:internal/modules/cjs/loader:960:12)
    at Function.executeUserEntryPoint [as runMain] (node:internal/modules/run_main:81:12)
    at node:internal/main/run_main_module:23:47

Node.js v18.17.0
gaikwadrahul-macbookpro:test-7867 gaikwadrahul$