Update README.md

README.md

@@ -74,6 +74,56 @@ print('\033[1;32m' + prediction + '\033[0m' if prediction == 'safe' else '\033[1
 Output Example:
 ![prediction](./output_example.png)
 
+# Tensorflow.js
+```bash
+npm i @tensorflow/tfjs-node
+```
+```javascript
+const tf = require('@tensorflow/tfjs-node');
+const fs = require('fs');
+const { pipeline } = require('stream');
+const { promisify } = require('util');
+
+
+const download = async (url, path) => {
+    // Taken from https://levelup.gitconnected.com/how-to-download-a-file-with-node-js-e2b88fe55409
+
+    const streamPipeline = promisify(pipeline);
+    const response = await fetch(url);
+
+    if (!response.ok) {
+        throw new Error(`unexpected response ${response.statusText}`);
+    }
+
+    await streamPipeline(response.body, fs.createWriteStream(path));
+};
+
+
+async function run() {
+    // download saved model and variables from https://huggingface.co/FredZhang7/google-safesearch-mini/tree/main/tensorflow
+    if (!fs.existsSync('tensorflow')) {
+        fs.mkdirSync('tensorflow');
+        await download('https://huggingface.co/FredZhang7/google-safesearch-mini/resolve/main/tensorflow/saved_model.pb', 'tensorflow/saved_model.pb');
+        fs.mkdirSync('tensorflow/variables');
+        await download('https://huggingface.co/FredZhang7/google-safesearch-mini/resolve/main/tensorflow/variables/variables.data-00000-of-00001', 'tensorflow/variables/variables.data-00000-of-00001');
+        await download('https://huggingface.co/FredZhang7/google-safesearch-mini/resolve/main/tensorflow/variables/variables.index', 'tensorflow/variables/variables.index');
+    }
+
+    // load model and image
+    const model = await tf.node.loadSavedModel('./tensorflow/');
+    const image = tf.node.decodeImage(fs.readFileSync('cat.jpg'), 3);
+
+    // predict
+    const input = tf.expandDims(image, 0);
+    const tensor = model.predict(input);
+    const max = tensor.argMax(1);
+    const classes = ['nsfw_gore', 'nsfw_suggestive', 'safe'];
+    console.log('\x1b[32m%s\x1b[0m', classes[max.dataSync()[0]], '\n');
+}
+
+run();
+```
+
 <br>
 
 # Bias and Limitations