const fs = require("fs");
const tf = require("@tensorflow/tfjs-node");
const mobilenet = require("@tensorflow-models/mobilenet");
const knnClassifier = require("@tensorflow-models/knn-classifier");
const classifier = knnClassifier.create();
const dt = require("directory-tree");


class NeuralNet {
  async init() {
    let GlobalNet = await mobilenet.load();
    this.net = GlobalNet;
    this.times = 1;
  }

  async train(image,label) {
    //image is of type NeuralNet.Image
    const ac = await this.net.infer(image,true);
    classifier.addExample(ac,label);
  }

  async predict(image) {
    const ac = await this.net.infer(image,"conv_preds");
    let pred = await classifier.predictClass(ac);
    return pred
  }

  image(path) {
    const file = fs.readFileSync(path);
    return tf.node.decodeImage(file,3);
  }

  async trainDir(dir,label) {
    dt(dir).children.forEach((c)=>{
      this.train(this.image(c.path),label);
      console.log("TRAIN " + this.times);
      this.times++;
    })
  }
  
}


exports.NeuralNet = NeuralNet;