Clasificador de gatos y perros usando TensorFlow

El otro día estaba buscando información acerca de redes de convolución y encontré un estupendo video¹ en youtube, el cual explicaba como realizar un clasificador de imágenes en 5 min, así que dije ¿Por que no hago un clasificador de gatos y perros, pero a mi estilo? - why not?.

Configurando todo

Básicamente es utilizar el ejemplo de TensorFlow llamado retrain.py que es un script de entrenamiento de una red de convolución Inception-v3² y entrenar nuestra red con las imágenes que queremos clasificar, en este caso gatos y perros.

Vamos a necesitar muchas pero muchas imágenes de gatos y perros así que para obtener las imágenes de nuestro entrenamiento tenemos dos posibles ideas.

* Ir a Imagenet y obtener el archivo con todas las url de las imágenes y usando wget descargarlas.
* Descargar las imágenes de google images utilizando un plugin para chrome o firefox.

Una vez obtenidas las imágenes para nuestro entrenamiento vamos a crear nuestro proyecto, primero tenemos que crear un entorno con virtualenv e instalar TensorFlow y pillow.

  virtualenv clasificador-mascota
  cd clasificador-mascota
  source bin/activate
  pip install tensorflow
  pip install pillow

Creamos la estructura de nuestro proyecto.

  mkdir src
  mkdir -p data/mascota
  mkdir img

src contiene toda la talacha (scripts)
data/mascota es donde estarán los gatos y perros (imágenes)
img están las imágenes de prueba (test)

Entrenamiento

Dentro de la carpeta data/mascota vamos a crear las carpeta con las categorías a clasificar e.g. gato y perro y dentro de esta colocaremos todas las imágenes relacionadas a esta.

Lo siguiente es descargar el script de retrain.py que esta en el repositorio de TensorFlow en la seccion de example. Después lo colocamos en nuestra carpeta src y hacemos la magia.

  python src/retrain.py --bottleneck_dir=src/bottlenecks \
        --how_many_training_steps 500 \
        --model_dir=src/inception \
        --output_graph=src/retrained_graph.pb \
        --output_labels=src/retrained_labels.txt \
        --image_dir data/mascota

Indica donde se va a guardar la información de las imágenes a entrenar.
Numero de iteraciones para entrenar la red.
Indica donde se va a guardar el modelo de la red.
Indica donde se va a guardar el archivo que contiene las categorías a clasificar
Dirección de la carpeta que contiene las imágenes a entrenar.

Clasificación

Ya que tenemos nuestro entrenamiento es hora de hacer nuestro clasificador.

La idea es que ingresemos una imagen y nos indique si es un gato, perro o desconocido y devuelva una nueva imagen con la categoría a la que pertenece.

Let's do it!!

Creamos un script en src llamado classfier-2.1.py e importamos las bibliotecas que necesitaremos y cargamos la imagen de prueba y las categorías. threshold es un umbral que al no superarlo indica que la imagen no pertenece a una categoría.

import tensorflow as tf
import sys
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw 

# Direccion de la imagen
image_path = sys.argv[1]

# Se lee la imagen y se guarda en image_data 
image_data = tf.gfile.FastGFile(image_path, 'rb').read()

# carga el archivo de etiquetas
label_lines = [line.rstrip() for line in tf.gfile.GFile("src/retrained_labels.txt")]

# Umbral de predicción
threshold = 0.9

Cargamos nuestro modelo el cual es nuestro entrenamiento.

with tf.gfile.FastGFile("src/retrained_graph.pb", 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    _ = tf.import_graph_def(graph_def, name='')

Es hora de ingresar nuestra imagen en el modelo y que esta nos devuelva a que categoría pertenece.

with tf.Session() as sess:
    # Se ingresa image_data como entrada de graph para obtener una prediccion
    softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
    
    predictions = sess.run(softmax_tensor, \
            {'DecodeJpeg/contents:0': image_data})
    
    top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
    
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)

    # Si no supera el umbral no tiene categoria
    if predictions[0][top_k[0]] <= threshold:
        # Que mal, no era un gato o perro
        print('\x1b[0;31;40m' + 'Sin categoria' + '\x1b[0m')
        categorie = 'unknow'
    else:
        # Tenemos un gato o un perro :)
        print('\x1b[0;36;40m' + 'Categoria: ' + label_lines[top_k[0]] + '\x1b[0m')
        categorie = label_lines[top_k[0]]
    
    # Escribo en la imagen a que categoría pertenece
    draw.text((0, 0), categorie, (255, 255, 255), font=font)
    
    # Guardo la nueva imagen
    img.save(categorie + '-' + str(predictions[0][top_k[0]])[2:7] + '.jpg')

Y finalmente hacemos toda la magia con:

  python src/classfier-2.1.py img/1.jpg

img/1.jpg es nuestra imagen de prueba.

Ya tenemos nuestro clasificador de gatos y perros.

Codigo GitHub

[1] [Siraj Raval]. (2016, Sep 9). Build a TensorFlow Image Classifier in 5 Min. Retrieved from https://youtu.be/QfNvhPx5Px8
[2] Train your own image classifier with Inception in TensorFlow. (2016, March 09). Retrieved April 23, 2017, from https://research.googleblog.com/2016/03/train-your-own-image-classifier-with.html

Askatl

Buscar este blog