使用tensorflow和nodejs训练cocossd

时间:2019-06-29 20:40:17

标签: node.js tensorflow.js

我有这个有效的代码,它实际上有效,基本上它的作用是 借助于node.js将模型cocossd加载到tensorflow中,并对其进行预测。

但是我想做的是使用我自己的类和标签重新训练该模型。

所有这些都使用node.js。我怎样才能做到这一点?我已经有2周的时间,但没有在搜索中找到结果。请帮忙。

const tf = require('@tensorflow/tfjs')
const mn = require('@tensorflow-models/mobilenet');
const ccsd = require('@tensorflow-models/coco-ssd');
              require('@tensorflow/tfjs-node')

const fs = require('fs');
const jpeg = require('jpeg-js');

const COCOSSD_PATH = 'model/mobilenet_v2/model.json';
const IMAGE_PATH = 'panda.jpg';
const NUMBER_OF_CHANNELS = 3;

//LOAD DE MODEL
const loadModel = async path => {
  const model = new ccsd.ObjectDetection();
        model.modelPath = `file://${path}`;
        await model.load();
  return model;
}

//TRANSFORMA LA IMAGEN EN rgba VALUES
const readImage = path => {
  const buf = fs.readFileSync(path)
  const pixels = jpeg.decode(buf, true)
  return pixels
}

//CONVERT THE IMAGE FROM rgba VALUES TO rgb VALUES
const imageByteArray = (image, numChannels) => {
  const pixels = image.data
  const numPixels = image.width * image.height;
  const values = new Int32Array(numPixels * numChannels);

  for (let i = 0; i < numPixels; i++) {
    for (let channel = 0; channel < numChannels; ++channel) {
      values[i * numChannels + channel] = pixels[i * 4 + channel];
    }
  }
  return values
}

//CONVERT THE IMAGE TO A TENSOR
const imageToInput = (image, numChannels) => {
  const values = imageByteArray(image, numChannels)
  const outShape = [image.height, image.width, numChannels];
  const input = tf.tensor3d(values, outShape, 'int32');

  return input
}

//EVERY THING START HERE AND ENDS HERE
const inic = async (cocossdPath, imagePath) => {
  const image = readImage(imagePath)
  const input = imageToInput(image, NUMBER_OF_CHANNELS);

  const cocossd = await loadModel(cocossdPath);
  const detection = await cocossd.detect(input);

  console.log('classification results: ', detection);

}

inic(COCOSSD_PATH,IMAGE_PATH)

0 个答案:

没有答案