Tutorial: detectar texto en una imagen con Vision en Swift

El uso de técnicas de Machine Learning ha crecido de manera exponencial de un tiempo a esta parte, pasando de ser algo reservado a unos pocos a estar al alcance de la mayoría gracias a librerías como TensorFlow o Keras y frameworks como CoreML de Apple y ML Kit de Google. 

Una de las disciplinas con mayor número de crecimiento es el deep learning, más concretamente, el reconocimiento y/o clasificación de imágenes. En esta guía veremos como Apple nos brinda su ayuda en este menester mediante el framework Vision.

¿Qué es Vision?

Vision, aparte de ser el androide supérheroe sintético que forma parte de Los Vengadores, es el nombre que recibe uno de los frameworks desarrollados por Apple para que los programadores podamos aplicar Machine Learning a nuestras apps. Construido sobre CoreML y Core Image, Vision nos da las herramientas necesarias para trabajar con imágenes, además de:

• Encontrar caras y sus características.
• Rectángulos.
• Seguimiento de objetos (tracking).
• Detección de códigos de barras (en diferentes formatos).
• Situar la línea del horizonte.
• Alineación de imágenes.
• Análisis de imágenes con Machine Learning.
• Encontrar texto en imágenes.

En nuestro caso vamos a dedicarnos a encontrar texto dentro de las imágenes y mediante CoreML averiguar a qué imagen corresponde de las que tenemos almacenadas en la app.

BiciMAD. Un caso Práctico

Vamos a desarrollar una app que nos permita leer los números de identificación de las bicicletas del servicio BiciMAD, de la ciudad de Madrid. En este caso la app tendrá que hacer lo siguiente:

1. Cargar una imagen
2. Buscar texto
3. Extraer las imágenes correspondientes a cada uno de los cuatro números
4. Pasar esas imágenes al modelo de CoreML
5. Presentarle al usuario al número que ve en la bici

Pues manos a la obra…

Necesitamos imágenes

Para poder probar la app vamos a necesitar imágenes de los códigos de las bicis, así que lo primero es importarlas al simulador de iOS.

Como no todos vivimos en Madrid o tenemos acceso a fotos de estas bicis vamos a proporcionaros unas imágenes de números con el mismo formato que el del servicio BiciMAD. En este enlace encontraréis un pequeño conjunto de imágenes con las que podréis hacer las pruebas . 

Una vez que tengais descargadas las imágenes en vuestro ordenador, arrancad el simulador y abrid la app Photos y arrastrar los archivos de las fotos hasta el simulador. ¡Ya está!

Vale, ya tengo las imágenes. ¿Ahora qué?

Abrimos Xcode y creamos un proyecto de tipo Single View. Para seleccionar las imágenes que acabamos de importar en el simulador nos valdremos de UIImagePickerController.

private lazy var imagePicker: UIImagePickerController = {
    let picker = UIImagePickerController()
    picker.sourceType = .savedPhotosAlbum
    picker.modalPresentationStyle = .currentContext
    picker.delegate = self as! (UIImagePickerControllerDelegate & UINavigationControllerDelegate)
    return picker
}()

A nuestro view controller le añadimos un UIButton que nos servirá para presentar el UIImagePickerController. En cuanto selecciones una imagen el delegado del picker controller saltará y podremos cogerla desde la función imagePickerController(_:didFinishPickingMediaWithInfo:). Esta función recibe como uno de sus parámetros un diccionario en el que se encuentran valores como la ruta al archivo de la imagen o la imagen en sí misma.

func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [String : Any]) -> Void {
    // 1
    guard let image = info[UIImagePickerControllerOriginalImage] as? UIImage, let ciImage = image.ciImage else {
        return
    }
    
    // 2
    let imageRequest = VNImageRequestHandler(ciImage: ciImage)
    
    // 3
    do {
        try imageRequest.perform([ self.requestText ])
    } catch {
        print("No podemos procesar la imagen del Carrete de Fotos. \(error.localizedDescription)")
    }
}

Lo primero que hacemos ahora que tenemos una imagen, es convertirla a un tipo CIImage o una CGImage (no amigos, a CoreML no le gustan las UIImage) (1). Con la imagen ya convertida lanzamos una solicitud de análisis sobre la imagen (2) y le decimos a Vision que lo que esperamos encontrar es texto (3).

Supongo que te preguntarás que qué es ese self.requestText que se le pasa a la función perform. Pues eso es la petición de búsqueda de texto que le pasamos a Vision.

private func prepareTextDetection() -> Void {
    self.requestText = VNDetectTextRectanglesRequest(completionHandler: self.handleTextDetection)
    self.requestText.reportCharacterBoxes = true
}

private func handleTextDetection(request: VNRequest, error: Error?) -> Void {
    // ...
}

La clase VNDetectTextRectanglesRequest recibe como parámetro un closure de tipo VNRequestCompletionHandler, que se traduce en (request: VNRequest, error: Error?) -> Void. El hecho de incluir el código como un trailing closure se debe a que la funcionalidad es un tanto extensa , por lo que por claridad del código lo ponemos en un función que cumple con la firma de VNRequestCompletionHandler.

Fijaos en la propiedad reportCharacterBoxes, se usa para indicarle a Vision que queremos, o no, que nos informe además de cada uno de los caracteres que encuentre en el texto. Como el modelo espera como entrada un sólo dígito le asignamos un valor de true.

¿Eso es todo? ¡Qué fácil!

No amigos, esto no ha hecho más que empezar. Lo único que hemos hecho es llamar la atención de Vision sobre la imagen. Tenemos que decirle ahora qué queremos hacer cuando encuentre texto, si lo encuentra en la imagen, y para eso tenemos que programarlo dentro de la función handleTextDetection.

Lo primero que tenemos que hacer es comprobar que tenemos todo los necesario para poder convertir esa foto en números.

// 1
guard let results = request.results, let ciBikeImage = self.ciBikeImage, !results.isEmpty else {
    return
}

// 2
guard let result = results.first, let observation = result as? VNTextObservation, observation.isBikeIdentifier else {
    return
}

Cuando sepamos que tenemos resultados válidos, y que esos resultados tienen el formato esperado (4 caracteres) nos preparamos para procesar cada uno de esos caracteres.

let visionary = Visionary()
visionary.delegate = self

let imageSize = ciBikeImage.extent.size
var boxes = [CIImage]()

La clase Visionary es un wrapper que nos facilita el uso del modelo de CoreML.

Ya estamos listos para empezar a procesar cada uno de los caracteres encontrados, así que vamos a ver qué tenemos que hacer. Lo primero es traducir las coordenadas de las observación a coordenadas sobre la imagen. La app que acompaña al artículo incluye una función showMarker que muestra sobre la imagen los rectángulos tanto para la palabra completa (en color rojo) como para cada uno de los caracteres (en color azul).

let box_rect = box.boundingBox.scaled(to: imageSize)

// Rectify the detected image and reduce it to inverted grayscale for applying model.
let topLeft = box.topLeft.scaled(to: imageSize)
let topRight = box.topRight.scaled(to: imageSize)
let bottomLeft = box.bottomLeft.scaled(to: imageSize)
let bottomRight = box.bottomRight.scaled(to: imageSize)

Con las coordenadas para el carácter, cortamos la porción de imagen en la que se encuentra y aplicamos una serie de filtros sobre la imagen.

let correctedImage = ciBikeImage
    .cropped(to: box_rect)
    .applyingFilter("CIColorInvert")
    .applyingFilter("CIPerspectiveCorrection", parameters: [
        "inputTopLeft": CIVector(cgPoint: topLeft),
        "inputTopRight": CIVector(cgPoint: topRight),
        "inputBottomLeft": CIVector(cgPoint: bottomLeft),
        "inputBottomRight": CIVector(cgPoint: bottomRight)
    ])
    .applyingFilter("CIColorControls", parameters: [
        kCIInputSaturationKey: 0,
        kCIInputContrastKey: 32
    ])

La función cropped es la encargada de recortar, el filtro CIColorInvert realiza una inversión de color, CIPerspectiveCorrection corrige la perspectiva de la imagen para ponerla de frente y por último CIColorControls pone la imagen en blanco y negro. Con todo esto ya tenemos una imagen tal y como la espera el modelo.

Ya tengo las imágenes listas ¿Hemos terminado?

Nos queda muy poco. Ahora que tenemos las 4 imagenes, una para caracter, se las pasamos a la función prediction de la clase Visionary. Desde ahora esa clase va a ser la encargada de tratar con el modelo y clasificar las imágenes.

visionary.prediction(images: boxes)

Para cada una de las imágenes tenemos que hacer una solicitud de procesado de imagen (sí, otra vez) pero en este caso lo que queremos que ejecute es un VNCoreMLRequest, es decir, una llamada directa a nuestro modelo para que nos devuelva una predicción o una clasificación, dependiendo del tipo de modelo.

Dicho lisa y llanamente, lo que nos va a decir VNCoreMLRequest es si la imagen es un 3, un 5 o un 8.

let request = VNCoreMLRequest(model: model) { [weak self] (request, error) -> Void in
    self?.handleModelRequest(request: request, forIndex: index)
}

// Como escalamos la imagen para adecuarla al
// tamaño esperado por el modelo.
request.imageCropAndScaleOption = .scaleFit

let handler = VNImageRequestHandler(ciImage: image)

do {
    try handler.perform([ request ])
} catch {
    print("No se puede ejecutar la operación con Vision. \(error.localizedDescription)")
}

¡¿Pero dónde está el resultado?!

Paciencia joven Padawan, que ya casi hemos terminado. Si te fijas en el código anterior, verás que VNCoreMLRequest también tiene un closure como parámetro que se invoca cuando el modelo devuelve su resultado.

private func handleModelRequest(request: VNRequest, forIndex index: Int) -> Void {
    guard let results = request.results as? [VNClassificationObservation], let observation = results.first, let digit = Int(observation.identifier) else {
        return
    }
    
    ///
    print("La imagen es un \(digit)")
}

El parámetro request contiene una array de observaciones, almacenadas en la propiedad results, que son las probabilidades para cada una de las clases (números) del modelo.

Cuando selecciono el primer valor de results estoy cogiendo el número con más probabilidades de coincidir con la imagen según lo ha interpretado el modelo.

Una reflexión final

Hay que terner clara una cosa, Vision aplicado a la detección de texto se limita a eso, a detectar texto. No esperéis que por arte de magía una variable de tipo String tenga el contenido del texto. Para ello es necesario un modelo de Machine Learning aparte.

Pero que esto no nos haga pensar que Vision no es tan bueno como parece, todo lo contrario. Que en nuestro bolsillo podamos llevar tal capacidad de análisis de imágenes y de tipo tan variado es todo un lujo.

En este repositorio de GitHub está el proyecto Xcode junto las imágenes para pruebas. Un saludo, y ahora más que nunca, Good Apple Coding.