Classificador de imagem de aprendizado por transferência do TensorFlow.js

Neste codelab, você aprenderá como criar uma Máquina que Aprende simples, um classificador de imagens personalizadas que será treinado rapidamente no navegador usando o TensorFlow.js, uma biblioteca de machine learning avançada e flexível para JavaScript. Primeiro, você carregará e executará um modelo pré-treinado conhecido, chamado MobileNet, para classificação de imagens no navegador. Em seguida, você usará uma técnica chamada "aprendizado por transferência", que inicializa nosso treinamento com o modelo MobileNet pré-treinado e o personaliza para treinar para seu aplicativo.

Este codelab não abordará a teoria por trás do aplicativo de Máquina que Aprende. Caso esteja curioso sobre isso, confira este tutorial.

Você aprenderá a fazer o seguinte

  • Como carregar um modelo MobileNet pré-treinado e fazer uma previsão sobre novos dados
  • Como fazer previsões pela webcam
  • Como usar ativações intermediárias do MobileNet para fazer aprendizado por transferência em um novo conjunto de classes que você define dinamicamente com a webcam

Então, vamos começar!

Para concluir este codelab, você precisará dos seguintes requisitos:

  1. Uma versão recente do Chrome ou de outro navegador moderno.
  2. Um editor de texto executado localmente na sua máquina ou na Web por meio de algo como Codepen ou Glitch.
  3. Conhecimento sobre HTML, CSS, JavaScript e Chrome DevTools (ou as DevTools do seu navegador preferido).
  4. Conhecimento conceitual de alto nível sobre redes neurais. Se você precisar de uma introdução ou revisão, assista a este vídeo da 3blue1brown ou este vídeo sobre aprendizado profundo em JavaScript por Ashi Krishnan (links em inglês).

Abra index.html em um editor e adicione este conteúdo:

<html>
  <head>
    <!-- Load the latest version of TensorFlow.js -->
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet"></script>
  </head>
  <body>
    <div id="console"></div>
    <!-- Add an image that we will use to test -->
    <img id="img" crossorigin src="https://i.imgur.com/JlUvsxa.jpg" width="227" height="227"/>
    <!-- Load index.js after the content of the page -->
    <script src="index.js"></script>
  </body>
</html>

Em seguida, abra/crie o arquivo index.js em um editor de código e inclua o seguinte código:

let net;

async function app() {
  console.log('Loading mobilenet..');

  // Load the model.
  net = await mobilenet.load();
  console.log('Successfully loaded model');

  // Make a prediction through the model on our image.
  const imgEl = document.getElementById('img');
  const result = await net.classify(imgEl);
  console.log(result);
}

app();

Para executar a página da Web, basta abrir index.html em um navegador da Web. Se você estiver usando o Console do Cloud, basta atualizar a página de visualização.

Você verá uma imagem de um cachorro e, no Console JavaScript nas Ferramentas para desenvolvedores, as principais previsões do MobileNet. O download do modelo pode demorar um pouco.

A imagem foi classificada corretamente?

Vale lembrar que isso também funcionará em um celular.

Agora, vamos tornar esse processo mais interativo e em tempo real. Vamos configurar a webcam para fazer previsões sobre imagens que chegam por ela.

Primeiro configure o elemento de vídeo da webcam. Abra o arquivo index.html, adicione a seguinte linha na seção <body> e exclua a tag <img> que usamos para carregar a imagem do cachorro:

<video autoplay playsinline muted id="webcam" width="224" height="224"></video>

Abra o arquivo index.js e adicione o webcamElement ao topo do arquivo

const webcamElement = document.getElementById('webcam');

Agora, na função app() que você adicionou antes, é possível remover a previsão por meio da imagem e, em vez disso, criar um loop infinito que faz previsões pelo elemento da webcam.

async function app() {
  console.log('Loading mobilenet..');

  // Load the model.
  net = await mobilenet.load();
  console.log('Successfully loaded model');

  // Create an object from Tensorflow.js data API which could capture image
  // from the web camera as Tensor.
  const webcam = await tf.data.webcam(webcamElement);
  while (true) {
    const img = await webcam.capture();
    const result = await net.classify(img);

    document.getElementById('console').innerText = `
      prediction: ${result[0].className}\n
      probability: ${result[0].probability}
    `;
    // Dispose the tensor to release the memory.
    img.dispose();

    // Give some breathing room by waiting for the next animation frame to
    // fire.
    await tf.nextFrame();
  }
}

Se você abrir o console na página da Web, verá uma previsão do MobileNet com probabilidade de cada frame coletado na webcam.

Isso pode não ser muito intuitivo porque o conjunto de dados do ImageNet não parece muito com imagens que normalmente apareceriam em uma webcam. Uma maneira de testar isso é segurando uma foto do seu cachorro na frente da câmera do laptop.

Agora, vamos tornar esse modelo mais útil. Criaremos um classificador de objeto personalizado de três classes usando a webcam. Faremos uma classificação por meio do MobileNet, mas desta vez usaremos uma representação interna (ativação) do modelo para uma determinada imagem de webcam e a usaremos para classificação.

Usaremos um módulo chamado "K-Nearest Neighbors Classifier", que permite colocar imagens da webcam (na verdade, as ativações do MobileNet) em diferentes categorias (ou "classes") e, quando o usuário solicitar uma previsão, basta escolher a classe que tem a ativação mais semelhante à da previsão que está sendo realizada.

Adicione uma importação do KNN Classifier ao final das importações na tag <head> de index.html. Você ainda precisará do MobileNet, então não remova essa importação:

...
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/knn-classifier"></script>
...

Adicione três botões para cada um dos botões em index.html abaixo do elemento de vídeo. Esses botões serão usados para adicionar imagens de treinamento ao modelo.

...
<button id="class-a">Add A</button>
<button id="class-b">Add B</button>
<button id="class-c">Add C</button>
...

Na parte superior de index.js, crie o classificador:

const classifier = knnClassifier.create();

Atualize a função do app:

async function app() {
  console.log('Loading mobilenet..');

  // Load the model.
  net = await mobilenet.load();
  console.log('Successfully loaded model');

  // Create an object from Tensorflow.js data API which could capture image
  // from the web camera as Tensor.
  const webcam = await tf.data.webcam(webcamElement);

  // Reads an image from the webcam and associates it with a specific class
  // index.
  const addExample = async classId => {
    // Capture an image from the web camera.
    const img = await webcam.capture();

    // Get the intermediate activation of MobileNet 'conv_preds' and pass that
    // to the KNN classifier.
    const activation = net.infer(img, true);

    // Pass the intermediate activation to the classifier.
    classifier.addExample(activation, classId);

    // Dispose the tensor to release the memory.
    img.dispose();
  };

  // When clicking a button, add an example for that class.
  document.getElementById('class-a').addEventListener('click', () => addExample(0));
  document.getElementById('class-b').addEventListener('click', () => addExample(1));
  document.getElementById('class-c').addEventListener('click', () => addExample(2));

  while (true) {
    if (classifier.getNumClasses() > 0) {
      const img = await webcam.capture();

      // Get the activation from mobilenet from the webcam.
      const activation = net.infer(img, 'conv_preds');
      // Get the most likely class and confidence from the classifier module.
      const result = await classifier.predictClass(activation);

      const classes = ['A', 'B', 'C'];
      document.getElementById('console').innerText = `
        prediction: ${classes[result.label]}\n
        probability: ${result.confidences[result.label]}
      `;

      // Dispose the tensor to release the memory.
      img.dispose();
    }

    await tf.nextFrame();
  }
}

Agora, ao carregar a página index.html, você pode usar objetos comuns ou gestos faciais/corporais para capturar imagens para cada uma das três classes. Cada vez que você clicar em um dos botões "Adicionar", uma imagem é adicionada à classe como exemplo de treinamento. Enquanto você faz isso, o modelo continua fazendo previsões sobre imagens de webcam recebidas e mostrando os resultados em tempo real.

Tente adicionar outra classe que não represente nenhuma ação.

Neste codelab, você implementou um app da Web de machine learning simples usando o TensorFlow.js. Você carregou e usou um modelo MobileNet pré-treinado para classificar imagens da webcam. Depois, você personalizou o modelo para classificar as imagens em três categorias personalizadas.

Acesse js.tensorflow.org para ver mais exemplos e demonstrações com código para saber como você pode usar o TensorFlow.js nos seus aplicativos.