Retraining image classification model using TensorFlow.js in Node.
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
.gitignore
README.md
app.js
augment_images.js
create_images.js
data.js
model.js
package-lock.json
package.json
ui_mock.js

README.md

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