En este notebook se muestra cómo crear un modelo de detección de objetos usando la arquitectura Faster R-CNN. Para crear nuestro modelo vamos a utilizar la librería IceVision que es una librería para crear modelos de detección usando FastAI.

En esta práctica vamos a hacer un uso intensivo de la GPU, así que es importante activar su uso desde la opción Configuración del cuaderno del menú Editar (esta opción debería estar habilitada por defecto, pero es recomendable que lo compruebes).

Librerías

Comenzamos descargando la librería IceVision. Al finalizar la instalación deberás reiniciar el kernel (menú Entorno de ejecución -> Reiniciar Entorno de ejecución).

!pip install icevision[all] -Uq

A continuación, cargamos aquellas librerías que son necesarias.

from icevision.all import *

Dataset

Para esta práctica vamos a usar como ejemplo el Fruit Images for Object Detection dataset. Este dataset consta de 240 imágenes de entrenamiento y 60 de test con tres categorías: manzanas, plátanos y naranjas. Los siguientes comandos descargan y descomprimen dicho dataset.

%%capture
!wget https://www.dropbox.com/s/1dsfd5rrmg3riqj/fruits.zip?dl=1 -O fruits.zip
!unzip fruits.zip

Vamos a explorar el contenido de este dataset. Para ello vamos a crear un objeto Path que apunta al directorio que acabamos de crear.

path=Path('fruits')

Como en la práctica anterior, podemos ver el contenido de este directorio usando el comando ls().

path.ls()

Si exploráis el directorio podréis ver que hay dos carpetas (llamadas train y test), y que cada una de ellas contiene dos carpetas, llamadas images y labels. La carpeta images contiene las imágenes del dataset, y la carpeta labels contiene las anotaciones en formato Pascal VOC. Para cada imagen, hay un fichero xml con el mismo nombre que contiene su extensión.

Icevision

El proceso para crear y evaluar un modelo de IceVision consta de los siguientes pasos:

  1. Crear un parser para leer las imágenes y las anotaciones.
  2. Construir objetos record a partir de los parser.
  3. Crear los datasets a partir de los records y los aumentos que queramos aplicar.
  4. Crear un dataloader a partir de los datasets.
  5. Definir un modelo.
  6. Entrenar el modelo.
  7. Guardar el modelo.
  8. Usar el modelo para inferencia

Vamos a ver en detalle cada uno de estos pasos.

Parser

IceVision proporciona una serie de parsers definidos por defecto para leer las anotaciones en distintos formatos entre ellos Pascal VOC y COCO. También es posible crear parsers propios. A pesar de que existe un parser para el formato de nuestra anotación, vamos a ver cómo crear un parser desde cero.

El primer paso es crear una plantilla para nuestro dataset.

template_record = ObjectDetectionRecord()

A continuación IceVision proporciona el método generate_template que nos proporciona los métodos que debemos implementar.

Parser.generate_template(template_record)

A continuación vamos a implementar nuestra clase con cada uno de esos métodos.

import xml.etree.ElementTree as ET

class MyParser(Parser):
    """Definimos el constructor de nuestra clase que va a recibir cuatro parámetros:
       - La plantilla definida previamente.
       - El path al directorio donde se encuentran las imágenes.
       - El path al directorio donde se encuentran las anotaciones.
       - Un objeto class_map con las clases que tiene nuestro dataset.
    """
    def __init__(self, template_record,path_img,path_anotaciones,class_map):
        super().__init__(template_record=template_record)
        self.path_img = path_img
        self.path_anotaciones= path_anotaciones
        self.class_map = class_map

    """El método iter escanea el directorio de anotaciones y nos devuelve el nombre 
    de cada fichero. Dicho nombre será utilizado por el resto de método"""    
    def __iter__(self):
        with os.scandir(self.path_anotaciones) as ficheros:
            for fichero in ficheros:
                yield fichero.name
                
    """El método len nos indica el número de elementos de los que consta nuestro 
    dataset"""
    def __len__(self):
        return len(self.path_anotaciones)

    """A partir del nombre del fichero de anotación, record_id debe devolver el identificador
    (o nombre) de la imagen asociada"""
    def record_id(self, o) -> Hashable: #o --> nombre de la anotación
        return o[:o.find('.')]

    """A continuación deberíamos definir el método parse_fields, pero vamos a definir una serie
    de definiciones previas que nos serán útiles"""

    """El método prepare recibe el nombre de un fichero de anotación como parámetro y realiza
    una serie de labores de preprocesamiento sobre dicho fichero de anotación. En este caso lo procesa
    usando la funcionalidad de la librería para trabajar con xml"""
    def prepare(self, o):
        tree = ET.parse(str(self.path_anotaciones)+'/'+str(o))
        self._root = tree.getroot()

    """El método filepath a partir del nombre del fichero de anotación devuelve el path de 
    la imagen asociada"""
    def filepath(self, o) -> Union[str, Path]:
        path=Path(f"{o[:o.find('.')]}.jpg")
        return self.path_img / path

    """La función image_width_height devuelve el ancho y el alto de una imagen a partir del nombre
    del fichero de anotación"""
    def image_width_height(self, o) -> Tuple[int, int]:
        return get_img_size(str(self.path_img)+'/'+f"{o[:o.find('.')]}.jpg")

    """La función labels recibe el nombre del fichero de anotación y debe devolver una lista 
    con los identificadores de las clases contenidas en dicho fichero."""
    def labels(self, o) -> List[int]:
        labels = []
        for object in self._root.iter("object"):
            label = object.find("name").text
            label_id = self.class_map.get_by_name(label)
            labels.append(label)

        return labels

    """La función bboxes recibe el nombre del fichero de anotación y debe devolver una lista 
    de bboxes que son las anotaciones contenidas en dicho fichero. El formato de cada BBOX es
    xmin, ymin, xmax, ymax."""
    def bboxes(self, o) -> List[BBox]:
        def to_int(x):
            return int(float(x))

        bboxes = []
        for object in self._root.iter("object"):
            xml_bbox = object.find("bndbox")
            xmin = to_int(xml_bbox.find("xmin").text)
            ymin = to_int(xml_bbox.find("ymin").text)
            xmax = to_int(xml_bbox.find("xmax").text)
            ymax = to_int(xml_bbox.find("ymax").text)

            bbox = BBox.from_xyxy(xmin, ymin, xmax, ymax)
            bboxes.append(bbox)

        return bboxes


    """Definimos a continuación el método parse_fields para cada elemento de nuestro 
    dataset proporcionamos:
       - El path a la imagen.
       - El tamaño de la imagen.
       - El mapa de clases.
       - Los rectángulos que indican cada uno de los objetos de la imagen.
       - Las etiquetas de cada uno de los objetos de la imagen."""
    def parse_fields(self, o: Any, record: BaseRecord, is_new: bool):
        self.prepare(o)
        if is_new:
            record.set_filepath(self.filepath(o))
            record.set_img_size(self.image_width_height(o))
            record.detection.set_class_map(self.class_map)
        record.detection.add_bboxes(self.bboxes(o))
        record.detection.add_labels(self.labels(o))

Una vez que hemos definido nuestra clase para parsear las anotaciones de nuestro dataset, vamos a construir los objetos correspondientes.

Lo primero que tenemos que hacer es construir nuestro class_map que es un objeto de la clase ClassMap y que contiene las clases de objetos de nuestro dataset.

class_map = ClassMap(['apple','banana','orange'])

A continuación definimos nuestros parsers. Uno para leer el conjunto de entrenamiento, y otro para el de test.

trainPath = Path('fruits')/'train'
parserTrain = MyParser(template_record, trainPath/'images', trainPath/'labels', class_map)

testPath = Path('fruits')/'test'
parserTest = MyParser(template_record,testPath/'images', testPath/'labels', class_map)

Records

Un record es un diccionario que contiene todos los campos parseados definidos en el proceso anterior. Mientras que cada parser es específico para cada anotación concreta, los objetos record tienen una estructura común. Para construir los records a partir de nuestros objetos parser debemos llamar al método parse e indicarle cómo se van a repartir los datos que se lean.

Como siempre, vamos a dividir nuestro dataset en tres partes: un conjunto de entrenamiento, uno de validación y uno de test. Por lo tanto tendremos que construir tres records llamados train_records, valid_records y test_records. Los records de entrenamiento y validación los construiremos a partir de los datos de entrenamiento usando una partición 90/10. Mientras que el record de test se construye a partir del conjunto de test usándolo completamente.

train_records, valid_records = parserTrain.parse(RandomSplitter((0.9, 0.1)))
test_records,_= parserTest.parse(RandomSplitter((1.0, 0.0)))

Transforms

Las transformaciones o aumentos son una parte fundamental cuando estamos construyendo modelos en visión por computador. IceVision incluye por defecto la librería Albumentations que nos proporciona una gran cantidad de transformaciones. Además es capaz de gestionar los cambios en la anotación que son necesarios cuando se trabaja en detección de objetos.

IceVision proporciona una función muy útil que es tfms.A.aug_tfms con una gran cantidad de transformaciones. Además podemos añadirle cualquier otra transformación de Albumentations.

Para nuestro ejemplo vamos a usar las transformaciones sugeridas por defecto por IceVision y aplicar la técnica de presizing vista para clasificación de imágenes; además será necesario normalizar las imágenes al rango de las imágenes de ImageNet. Notar que las transformaciones se aplicarán solo al conjunto de entrenamiento. Para los conjuntos de validación y test únicamente tendremos que escalar la imagen al tamaño adecuado y normalizarla.

presize = 512
size = 384
train_tfms = tfms.A.Adapter(
    [*tfms.A.aug_tfms(size=size, presize=presize), tfms.A.Normalize()]
)
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size), tfms.A.Normalize()])

Dataset

La clase Dataset sirve para combinar los records y las transformaciones. Debemos crear un dataset para nuestro conjunto de entrenamiento, otro para el conjunto de validación y otro para el de test.

train_ds = Dataset(train_records, train_tfms)
valid_ds = Dataset(valid_records, valid_tfms)
test_ds = Dataset(test_records, valid_tfms)

Una vez creados dichos datasets podemos mostrar imágenes de los mismos. En concreto la siguiente instrucción muestra imágenes del conjunto de entrenamiento a las cuáles se les han aplicado una serie de transformaciones. Es conveniente ejecutar esta visualización para comprobar que todo está correcto.

samples = [train_ds[0] for _ in range(3)]
show_samples(samples, ncols=3, class_map=class_map, denormalize_fn=denormalize_imagenet)

DataLoaders

Al igual que vimos para los modelos de clasificación de FastAI, el último paso es crear nuestros DataLoaders a partir de los datasets construidos anteriormente. Notar que cada modelo tiene su propio DataLoader. En este caso como vamos a crear un modelo de Faster RCNN debemos usar las siguientes instrucciones.

train_dl = models.torchvision.faster_rcnn.train_dl(train_ds, batch_size=8, num_workers=0, shuffle=True)
valid_dl = models.torchvision.faster_rcnn.valid_dl(valid_ds, batch_size=8, num_workers=0, shuffle=False)
test_dl = models.torchvision.faster_rcnn.valid_dl(test_ds, batch_size=8, num_workers=0, shuffle=False)

Entrenando el modelo

Para crear y entrenar nuestro modelo debemos crear un objeto Learner de FastAI. Para crear dicho objeto, lo primero que debemos hacer es construir un modelo con la arquitectura que queremos usar, en este caso Faster RCNN y con un backbone que es Resnet18.

model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn(pretrained=True),
                                       num_classes=len(class_map))

A continuación debemos proporcionar las métricas que queremos utilizar para evaluar el modelo. Por el momento la única métrica soportada por IceVision es el mAP de COCO, por lo tanto utilizaremos dicha métrica.

metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]

Ya estamos listos para construir nuestro Learner. Notar que dicho objeto se construye de manera distinta dependiendo de la arquitectura que queramos utilizar.

learn = models.torchvision.faster_rcnn.fastai.learner(dls=[train_dl, valid_dl], model=model, metrics=metrics)

Ahora podemos entrenar nuestro modelo utilizando la técnica de fine tuning que vimos en clase.

learn.fine_tune(10,freeze_epochs=2)

Una vez finalizado el entrenamiento podemos guardar nuestro modelo del siguiente modo.

torch.save(model.state_dict(),'fasterRCNNFruits.pth')

Evaluando el modelo

Al igual que vimos para los modelos de clasificación, la métrica mostrada durante el proceso de entrenamiento se refiere al conjunto de validación, mientras que nos interesa saber el resultado obtenido para el conjunto de test.

Para ello, lo primero que debemos hacer es construir un nuevo dataloader del siguiente modo, indicando que el conjunto de validación es el de test.

newdl = fastai.DataLoaders(models.torchvision.faster_rcnn.fastai.convert_dataloader_to_fastai(train_dl),
                           models.torchvision.faster_rcnn.fastai.convert_dataloader_to_fastai(test_dl)).to('cuda')

A continuación modificamos el dataloader del objeto Learn que hemos entrenado anteriormente.

learn.dls = newdl

Por último evaluamos nuestro modelo usando el método validate(). Al igual que en el caso de los modelos de clasificación el método validate() devuelve dos valores, el valor de la pérdida y el valor de la métrica asociada al conjunto de validación, que en este caso es el de test.

learn.validate()

Inferencia

Vamos a ver cómo usar el modelo ante una nueva imagen. Para ello lo primero que vamos a hacer es cargar dicho modelo. Para ello debemos crear un modelo con la arquitectura que utilizamos (Faster RCNN en nuestro caso), y posteriormente cargar el modelo.

model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn,
                                             num_classes=len(class_map))
state_dict = torch.load('fasterRCNNFruits.pth')
model.load_state_dict(state_dict)

El siguiente paso es cargar la imagen, para lo que usaremos la librería PIL.

import PIL
img = PIL.Image.open('fruits/test/images/mixed_25.jpg')

La siguiente instrucción permite mostrar la imagen que acabamos de cargar.

img

Ya estaríamos listos para relizar las predicciones sobre la imagen. Sin embargo, cabe recordar que primero debemos reescalar las imágenes y normalizarlas.

infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size),tfms.A.Normalize()])

Ya podemos realizar las predicciones mediante el método end2end_detect. Este método, que depende de la arquitectura que hayamos utilizado, recibe como parámetros la imagen sobre la que queremos realizar las predicciones, las transformaciones a aplicar, el modelo (movido a la CPU), el mapa de clases, y el nivel de confianza mínimo para realizar la predicción.

pred_dict  = models.torchvision.faster_rcnn.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=class_map, detection_threshold=0.5)
pred_dict['img']