Building AI-Powered Features with TensorFlow.js

Learn how to integrate machine learning capabilities into your web applications using TensorFlow.js.

Setting Up TensorFlow.js

First, install TensorFlow.js:

npm install @tensorflow/tfjs

Image Classification Example

import * as tf from '@tensorflow/tfjs';
import { loadLayersModel } from '@tensorflow/tfjs';

async function loadModel() {
  const model = await loadLayersModel('https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/2');
  return model;
}

async function classifyImage(imageElement: HTMLImageElement) {
  const model = await loadModel();
  
  // Preprocess the image
  const tensor = tf.browser.fromPixels(imageElement)
    .resizeNearestNeighbor([224, 224])
    .toFloat()
    .expandDims();
  
  // Make prediction
  const predictions = await model.predict(tensor).data();
  
  // Get top 5 predictions
  const top5 = Array.from(predictions)
    .map((p, i) => ({ probability: p, className: IMAGENET_CLASSES[i] }))
    .sort((a, b) => b.probability - a.probability)
    .slice(0, 5);
    
  return top5;
}

Text Classification Example

import * as use from '@tensorflow-models/universal-sentence-encoder';

async function loadEncoder() {
  const model = await use.load();
  return model;
}

async function getSentenceEmbeddings(sentences: string[]) {
  const model = await loadEncoder();
  const embeddings = await model.embed(sentences);
  return embeddings;
}

async function findSimilarSentences(query: string, corpus: string[]) {
  const model = await loadEncoder();
  
  // Get embeddings for query and corpus
  const queryEmbedding = await model.embed([query]);
  const corpusEmbeddings = await model.embed(corpus);
  
  // Calculate cosine similarity
  const similarities = tf.matMul(
    queryEmbedding, corpusEmbeddings.transpose()
  ).dataSync();
  
  // Return sorted results
  return Array.from(similarities)
    .map((score, idx) => ({ score, text: corpus[idx] }))
    .sort((a, b) => b.score - a.score);
}