OCR en el edge | Will Hackett

La tecnología avanzada se vuelve cada vez más corriente, y la visión por ordenador se ha convertido en una herramienta casi imprescindible, con muchas aplicaciones.

Sirve para identificar objetos, patrones, personas, plantas y mascotas. Da información útil a los usuarios, ayuda a tomar decisiones y facilita rastrear enormes conjuntos de datos.

Súmale la demanda creciente de análisis en tiempo real y de respuestas con baja latencia, y el edge compute se ha vuelto el mejor sitio para desplegar tu modelo entrenado.

Qué es el OCR

El reconocimiento óptico de caracteres (OCR) convierte imágenes de texto en texto codificado para máquinas. Se usa sobre todo para pasar imágenes de documentos a texto en el que se pueda buscar y que se pueda indexar.

Con él puedes escanear documentos, tickets, tarjetas de visita y más. Se entiende bien por qué sería útil en Blinq.

Qué es el edge compute

El edge compute consiste en ejecutar tu código lo más cerca posible del usuario. Normalmente se hace sobre una CDN, como Cloudflare Workers o AWS Lambda@Edge.

Es distinto de la nube tradicional, donde tu código corre en un centro de datos que suele estar lejos del usuario.

Por qué el edge

Se puede ejecutar el modelo en el dispositivo, pero dejarlo en el edge tiene varias ventajas:

Baja latencia: el modelo está más cerca del usuario, así que la respuesta llega antes. Parece que corre en el propio dispositivo.
Bajo coste: no pagas por los recursos de cómputo que ejecutan el modelo, solo por las peticiones.
Fácil de actualizar: puedes actualizar el modelo sin actualizar la app. Viene muy bien si tienes muchos usuarios con versiones antiguas.

En Blinq mantenemos una larga cola de versiones de la app móvil, así que poder actualizar el modelo sin tocar la app nos facilita mucho la vida.

Cómo ejecutar tu modelo en el edge

Los runtimes del edge, como Cloudflare Workers y Lambda@Edge, se parecen poco a los tradicionales, como Node.js y Golang. Están pensados para ser rápidos y ligeros, así que no tienen acceso a las mismas librerías ni a las mismas APIs.

Se parecen mucho a ejecutar tu código en un navegador, así que no basta con copiar y pegar y esperar que funcione.

El modelo

Para este ejemplo uso un modelo de OCR sencillo, Keras OCR. Es un modelo de TensorFlow que recibe una imagen y devuelve el texto que encuentra en ella.

También uso Constellation, de Cloudflare (en beta), que permite ejecutar modelos ONNX en el edge. Es una opción nueva muy buena, porque admite modelos que TensorFlow.js no soporta y no te limita la cuota de 10MB de un Worker de Cloudflare.

Lambda@Edge tiene límites mucho más altos, pero conviene igualmente mantener el modelo lo más pequeño posible para reducir el arranque en frío.

En este ejemplo he usado el modelo crnn_kurapan, que puedes descargar aquí.

Convertir Keras a ONNX

El primer paso es convertir el modelo de Keras a ONNX. Es un proceso bastante sencillo, pero tiene varios pasos. Doy por hecho que tienes Conda instalado y usas Python 3.10.

Si estás en Apple Silicon, necesitarás también:

tensorflow-macos

Primero, crea un entorno de conda.

$ conda create -n edge-onnx python=3.10 pip
$ conda activate edge-onnx

Después, instala las dependencias.

$ python -m tensorflow tf2onnx keras-ocr

Si estás en Apple Silicon, tendrás que instalar tensorflow-macos en lugar de tensorflow.

$ python -m pip install tensorflow-macos

Por último, convierte el modelo con este fragmento.

import tf2onnx
import onnx
import keras_ocr

recognition = keras_ocr.recognition.Recognizer(
    alphabet='abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', weights='kurapan')

output_path = 'crnn_kurapan.onnx'
model_proto, _ = tf2onnx.convert.from_keras(
    recognition.model, opset=None, output_path=output_path)

onnx.save(model_proto, output_path)

Esto crea un modelo ONNX en el directorio actual, en crnn_kurapan.onnx.

Ejecutar el modelo en el edge

Para ejecutar el modelo ONNX en el edge he usado Cloudflare Constellation, porque el modelo es demasiado grande para Cloudflare Workers.

Puedes hacerlo desde el panel de Cloudflare o desde la CLI. Yo uso la CLI porque es más fácil de automatizar.

Primero hay que crear un proyecto.

$ npx wrangler constellation project create "ocr" ONNX
$ npx wrangler constellation project list

Apunta el ID del proyecto.

A continuación, sube el modelo:

$ wrangler constellation model upload ocr crnn_kurapan crnn_kurapan.onnx
$ npx wrangler constellation model list ocr

Apunta el ID del modelo.

Luego, prepara el proyecto e inicializa wrangler.

$ npm create cloudflare@2

Yo llamé al proyecto ocr y usé la plantilla ‘Hello World’.

En el archivo wrangler.toml, añade el binding de constellation y activa node_compat.

name = "ocr"
main = "src/index.js"
node_compat = true
workers_dev = true
compatibility_date = "2023-07-01"

constellation = [
    {binding = 'OCR', project_id = '{{ project ID from earlier }}'},
]

Por último, instala el cliente de constellation y las dependencias para pasar las imágenes al modelo en streaming.

npm install --save @cloudflare/constellation string-to-stream pngjs

Ya podemos escribir el código que ejecuta el modelo.

import str from 'string-to-stream';
import { PNG } from 'pngjs/browser';

import { Tensor, run } from '@cloudflare/constellation';

const MODEL_ID = '{{ model ID from earlier }}';

function normalizeImage(data) {
  return new Promise(async (resolve, reject) => {
    const stream = str(data);

    const png = new PNG({ filterType: 4 });

    stream
      .pipe(png)
      .on('parsed', function () {
        const [r, g, b] = new Array(3).fill([]);

        for (let i = 0; i < this.data.length; i += 4) {
          r.push(this.data[i] / 255.0);
          g.push(this.data[i + 1] / 255.0);
          b.push(this.data[i + 2] / 255.0);
        }

        resolve({
          input: [...r, ...g, ...b],
          shape: [3, this.height, this.width, 3],
        });
      })
      .on('error', function (error) {
        reject({ err: error.toString() });
      });
  });
}

export default {
  async fetch(request, env) {
    if (request.method !== 'POST') {
      return new Response('Method not allowed', { status: 405 });
    }

    const formData = await request.formData();
    const image = formData.get('image');

    if (!image) {
      return new Response('No image found', { status: 400 });
    }

    const buffer = await image.arrayBuffer();

    try {
      const normalized = await normalizeImage(buffer);

      if (!normalized) {
        return new Response('Unable to normalize image', { status: 500 });
      }

      const input = new Tensor(
        'float32',
        normalized.shape,
        normalized.input,
        'input_1'
      );

      const output = await run(env.OCR, MODEL_ID, input);

      return new Response(JSON.stringify(output), { status: 200 });
    } catch (err) {
      return new Response(err, { status: 500 });
    }
  },
};

Mientras siga en beta, puede que tengas que desactivar los avisos de constellation con export NO_CONSTELLATION_WARNING=true.

Cuando estés listo para desplegar, ejecuta wrangler deploy.

Probar el modelo

Ahora puedes hacer una petición HTTP al Worker para probar el modelo.

$ curl -X POST -F "
image=@/path/to/image.png" https://ocr.example.workers.dev

Devuelve una respuesta JSON con la salida del modelo.

{
  "output_1": [...]
}

La salida del modelo es un array de probabilidades por cada carácter. Para convertirlo en texto puedes usar algo parecido a este fragmento.

const charset =
  'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789';

const text = output
  .map((probabilities) => {
    const max = Math.max(...probabilities);
    const index = probabilities.indexOf(max);

    return charset[index];
  })
  .join('');

Para qué hacerlo

Quería ver si era posible ejecutar en el edge un modelo grandecillo. Y quería ver cómo rendía.

Es poco probable que este modelo llegue a producción, pero es un buen ejemplo de lo que se puede hacer con Cloudflare Workers y Constellation. Los modelos de OCR más grandes suelen rendir mejor, aunque puedes usar modelos más pequeños si aceptas perder precisión. Para eso, siempre tienes Tesseract.js.

Rendimiento

El modelo tarda alrededor de 1,5 segundos en ejecutarse en el edge. Es mucho más lento que ejecutarlo en local, pero aún es lo bastante rápido para producción. En un caso real, además, puedes recurrir a los servicios de Google y de Amazon Web Services, mucho más precisos. Eso sí, también son mucho más caros.

Precisión

El modelo es bastante preciso, pero no perfecto. Le cuestan algunas tipografías y no reconoce bien el texto en imágenes con mucho ruido. Además, solo maneja imágenes pequeñas, aunque eso se puede mejorar con una ventana deslizante.

Pruébalo tú mismo

Cloudflare está desarrollando Constellation de forma activa y por ahora está en beta. Si quieres probarlo, puedes apuntarte a la beta aquí.

También hay un canal de Discord muy útil donde puedes hacer preguntas y pedir ayuda al equipo de Cloudflare.

Discusión en Hacker News