Skip to content

Retraining image classification model using TensorFlow.js in Node.

Notifications You must be signed in to change notification settings

adwellj/node-tfjs-retrain

Repository files navigation

node-tfjs-retrain

Retraining image classification model using TensorFlow.js in Node.

Based on the following js.tensorflow.org tutorials:

Excellent article from James Thomas that helped (skip the stuff about downloading the model shards manually; there's a much simpler way described in the comments):
Machine Learning in Node.js with TensorFlow.js

Example Usages

  • Retrain and test model:
    node app.js --images_dir="C:/Retraining_Project/Images" --model_dir="C:/Retraining_Project/Model"
  • Skip retraining; just test model:
    node app.js --images_dir="C:/Retraining_Project/Images" --model_dir="C:/Retraining_Project/Model" --skip_training=true
  • Create sample images:
    node create_images.js C:/Retraining_Project/Images

Help

I'm quite new to Node and TF, so there's a lot of room for cleanup and code improvements. My initial goal was to get something working and worry about minor improvements later. Feel free to submit a PR :)

Initial TODOs:

  • Create a script to automatically build a training data set to enable users to immediately get started
  • Save trained model in the project directory
  • Replace ui_mock.js with a TrainingParams class

About

Retraining image classification model using TensorFlow.js in Node.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published