Pesquisa de site

Detecção de objetos usando máscara R-CNN com TensorFlow 1.14 e Keras

Nota do editor: este artigo foi publicado originalmente em novembro de 2020 e algumas de suas informações estão desatualizadas. A teoria central mostrada, no entanto, é apoiada por pesquisas sólidas, e o código ainda é executável.

Mask R-CNN é um modelo de detecção de objetos baseado em redes neurais convolucionais profundas (CNN) desenvolvido por um grupo de pesquisadores de IA do Facebook em 2017. O modelo pode retornar a caixa delimitadora e uma máscara para cada objeto detectado em uma imagem.

O modelo foi originalmente desenvolvido em Python usando a biblioteca de aprendizagem profunda Caffe2. O código-fonte original está disponível no GitHub. Para oferecer suporte ao modelo Mask R-CNN com bibliotecas mais populares, como TensorFlow, existe um projeto popular de código aberto chamado Mask_RCNN que oferece uma implementação baseada em Keras e TensorFlow 1.14.

O Google lançou oficialmente o TensorFlow 2.0 em setembro de 2020. O TensorFlow 2.0 é mais bem organizado e muito mais fácil de aprender em comparação com o TensorFlow $\geq$1.0. Infelizmente, o projeto Mask_RCNN ainda não oferece suporte ao TensorFlow 2.0.

Este tutorial usa a versão TensorFlow 1.14 do projeto Mask_RCNN para fazer previsões e treinar o modelo Mask R-CNN usando um conjunto de dados personalizado. Em outro tutorial, o projeto será modificado para tornar o Mask R-CNN compatível com o TensorFlow 2.0.

Este tutorial cobre o seguinte:

  • Visão geral do projeto Mask_RCNN
  • Detecção de objetos com TensorFlow 1.14
  • Preparando os parâmetros de configuração do modelo
  • Construindo a arquitetura do modelo Mask R-CNN
  • Carregando os pesos do modelo
  • Lendo uma imagem de entrada
  • Detectando Objetos
  • Visualizando os Resultados
  • Código completo para previsão
  • Baixando o conjunto de dados de treinamento
  • Preparando o conjunto de dados de treinamento
  • Preparando a configuração do modelo
  • Máscara de treinamento R-CNN no TensorFlow 1.14
  • Conclusão

Pré-requisitos

É necessário um conhecimento básico de código Python e redes neurais para acompanhar este tutorial. Recomendamos este artigo para codificadores intermediários a avançados com experiência no desenvolvimento de novas arquiteturas.

O código neste artigo pode ser executado em um PC doméstico normal ou no DigitalOcean Droplet.

Visão geral do projeto Mask_RCNN

O projeto Mask_RCNN é de código aberto e está disponível no GitHub sob a licença MIT, que permite a qualquer pessoa usar, modificar ou distribuir o código gratuitamente.

A contribuição deste projeto é o suporte do modelo de detecção de objetos Mask R-CNN no TensorFlow $\geq$1.0, construindo todas as camadas do modelo Mask R-CNN e oferecendo uma API simples para treiná-lo e testá-lo.

O modelo Mask R-CNN prevê o rótulo da classe, a caixa delimitadora e a máscara para os objetos em uma imagem. Aqui está um exemplo do que o modelo pode detectar.

A primeira versão do projeto (Mask_RCNN 1.0) foi publicada em 3 de novembro de 2017. A última versão (Mask_RCNN 2.1) foi publicada em 20 de março de 2019. Desde esta data, nenhum novo lançamento foi publicado.

Para obter o projeto no seu PC, basta cloná-lo conforme o próximo comando:

git clone https://github.com/matterport/Mask_RCNN.git

Também é possível baixar o projeto em arquivo ZIP neste link. Vamos dar uma olhada rápida no conteúdo do projeto assim que estiver disponível localmente.

No momento em que este tutorial foi escrito, o projeto tinha 4 diretórios:

  1. mrcnn: Este é o diretório principal que contém o código Python do projeto.
  2. amostras: notebooks Jupyter fornecendo alguns exemplos de uso do projeto.
  3. imagens: uma coleção de imagens de teste.
  4. ativos: algumas imagens anotadas.

O diretório mais importante é mrcnn, pois contém o código-fonte do projeto. Possui os seguintes arquivos Python:

  • __init__.py: marca a pasta mrcnn como uma biblioteca Python.
  • model.py: Possui as funções e classes para construção das camadas e do modelo.
  • config.py: Contém uma classe chamada Config que contém alguns parâmetros de configuração sobre o modelo.
  • utils.py: Inclui algumas funções e classes auxiliares.
  • visualize.py: visualiza os resultados do modelo.
  • parallel_model.py: Para suporte a múltiplas GPUs.

Alguns dos arquivos na raiz do projeto são:

  • setup.py: Usado para instalar o projeto usando pip.
  • README.md: um arquivo Markdown que documenta o projeto.
  • LICENÇA: A licença do MIT.
  • requisitos.txt: Bibliotecas necessárias para usar o projeto.

Com base no arquivo requirements.txt, a versão do TensorFlow deve ser pelo menos 1.3.0. Para Keras, deve ser 2.0.8 ou superior.

Existem duas maneiras de usar o projeto:

  1. Instale-o usando pip.
  2. Copie a pasta mrcnn para onde você usará o projeto. Nesse caso, certifique-se de que todas as bibliotecas necessárias no arquivo requirements.txt estejam instaladas.

Para instalar o projeto, basta emitir o seguinte comando no prompt de comando ou terminal. Para plataformas diferentes do Windows, substitua “python” por “python3”.

python setup.py install

Uma forma alternativa de utilizar o projeto é copiar a pasta mrcnn para onde o projeto será utilizado. Suponha que exista um diretório chamado “Object Detection” dentro do qual existe um arquivo Python chamado object_detection.py que usa o código na pasta mrcnn. Em seguida, basta copiar a pasta mrcnn dentro do diretório “Object Detection”.

Aqui está a estrutura do diretório:

Object Detection
  mrcnn
  object_detection.py

Agora estamos prontos para usar o projeto Mask_RCNN. A próxima seção discute como usar o projeto com TensorFlow $\geq$1.0.

Detecção de objetos no TensorFlow 1

Antes de iniciar esta seção, certifique-se de que o TensorFlow 1 ($\geq$1.3.0) esteja instalado. Você pode verificar a versão usando o seguinte código:

import tensorflow

print(tensorflow.__version__)

As etapas para usar o projeto Mask_RCNN para detectar objetos em uma imagem são:

  1. Prepare os parâmetros de configuração do modelo.
  2. Construa a arquitetura do modelo Mask R-CNN.
  3. Carregue os pesos do modelo.
  4. Leia uma imagem de entrada.
  5. Detecte objetos na imagem.
  6. Visualize os resultados.

Esta seção cria um exemplo que usa uma máscara R-CNN pré-treinada para detectar os objetos no conjunto de dados COCO. As próximas subseções discutem cada uma das etapas listadas acima.

1. Prepare os parâmetros de configuração do modelo

Para construir o modelo Mask R-CNN, vários parâmetros devem ser especificados. Esses parâmetros controlam supressão não máxima (NMS), interseção sobre união (IoU), tamanho da imagem, número de ROIs por imagem, camada de pooling de ROI e muito mais.

A pasta mrcnn possui um script chamado config.py que possui uma única classe chamada Config. Esta classe possui alguns valores padrão para os parâmetros. Você pode estender esta classe e substituir alguns dos parâmetros padrão. O código a seguir cria uma nova classe chamada SimpleConfig que estende a classe mrcnn.config.Config.

import mrcnn.config

class SimpleConfig(mrcnn.config.Config):
  ...

Um dos parâmetros críticos que devem ser substituídos é o número de classes, cujo padrão é 1.

NUM_CLASSES = 1

Neste exemplo, o modelo detecta os objetos em uma imagem do conjunto de dados COCO. Este conjunto de dados possui 80 classes. Lembre-se que o antecedente deve ser considerado como uma classe adicional. Como resultado, o número total de turmas é 81.

NUM_CLASSES = 81

Os outros 2 parâmetros devem ser atribuídos com cuidado, que são GPU_COUNT e IMAGES_PER_GPU. O padrão é 1 e 2, respectivamente.

Estas 2 variáveis são usadas para calcular o tamanho do lote:

BATCH_SIZE = IMAGES_PER_GPU * GPU_COUNT

Supondo que os valores padrão sejam usados, o tamanho do lote será 2*1=2. Isso significa que 2 imagens são alimentadas no modelo de uma vez. Como resultado, o usuário deve alimentar 2 imagens de uma vez.

Em alguns casos o usuário está interessado apenas em detectar os objetos em uma única imagem. Assim, a propriedade IMAGES_PER_GPU deve ser definida como 1.

GPU_COUNT = 1
IMAGES_PER_GPU = 1

Aqui está o código completo para a classe de configuração. A propriedade NAME é um nome exclusivo para a configuração.

import mrcnn.config

class SimpleConfig(mrcnn.config.Config):
  NAME = "coco_inference"

  GPU_COUNT = 1
  IMAGES_PER_GPU = 1

  NUM_CLASSES = 81

2. Construir a arquitetura do modelo Mask R-CNN

Para construir a arquitetura do modelo Mask R-CNN, o script mrcnn.model possui uma classe chamada MaskRCNN. O construtor desta classe aceita 3 parâmetros:

  1. modo: "treinamento" ou "inferência".
  2. config: uma instância da classe de configuração.
  3. model_dir: Diretório para salvar logs de treinamento e pesos treinados.

O próximo exemplo cria uma instância da classe mrcnn.model.MaskRCNN. A instância criada é salva na variável model.

import mrcnn.model

model = mrcnn.model.MaskRCNN(mode="inference", 
                             config=SimpleConfig(),
                             model_dir=os.getcwd())

O modelo Keras é salvo no atributo keras_model da instância. Utilizando este atributo, o resumo do modelo pode ser impresso.

model.keras_model.summary()

A arquitetura do modo é grande; apenas 4 camadas da parte superior e inferior estão listadas abaixo. A camada final chamada mrcnn_mask retorna apenas as máscaras dos 100 principais ROIs.

___________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to
===========================================================================
input_image (InputLayer)        (None, None, None, 3 0         
___________________________________________________________________________
zero_padding2d_1 (ZeroPadding2D (None, None, None, 3 0           input_image[0][0]
___________________________________________________________________________
conv1 (Conv2D)                  (None, None, None, 6 9472        zero_padding2d_1[0][0]
___________________________________________________________________________
bn_conv1 (BatchNorm)            (None, None, None, 6 256         conv1[0][0]
___________________________________________________________________________
...
___________________________________________________________________________
mrcnn_mask_bn4 (TimeDistributed (None, 100, 14, 14,  1024        mrcnn_mask_conv4[0][0]
___________________________________________________________________________
activation_74 (Activation)      (None, 100, 14, 14,  0           mrcnn_mask_bn4[0][0]
___________________________________________________________________________
mrcnn_mask_deconv (TimeDistribu (None, 100, 28, 28,  262400      activation_74[0][0]
___________________________________________________________________________
mrcnn_mask (TimeDistributed)    (None, 100, 28, 28,  20817       mrcnn_mask_deconv[0][0]
===========================================================================
Total params: 64,158,584
Trainable params: 64,047,096
Non-trainable params: 111,488

3. Carregue os pesos do modelo

A última subseção criou a arquitetura do modelo. Esta subseção carrega os pesos no modelo criado usando o método load_weights(). É uma versão modificada do método Keras load_weights() que suporta o uso de multi-GPU, além da capacidade de excluir algumas camadas.

Os 2 parâmetros usados são:

  1. filepath: Aceita o caminho do arquivo de pesos.
  2. by_name: Se True, então cada camada recebe os pesos de acordo com seu nome.

O próximo código chama o método load_weights() enquanto passa o caminho do arquivo de pesos mask_rcnn_coco.h5. Este arquivo pode ser baixado neste link.

model.load_weights(filepath="mask_rcnn_coco.h5", 
                   by_name=True)

4. Leia uma imagem de entrada

Depois que o modelo for criado e seus pesos carregados, em seguida precisamos ler uma imagem e alimentá-la no modelo.

O próximo trecho de código usa OpenCV para ler uma imagem e reordenar seus canais de cores para RGB, em vez de BGR.

import cv2

image = cv2.imread("3627527276_6fe8cd9bfe_z.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

A próxima figura mostra a imagem que acabamos de ler. A imagem está disponível neste link.

5. Detectar objetos

Dado o modelo e a imagem de entrada, os objetos na imagem podem ser detectados usando o método detect(). Aceita 2 parâmetros:

  1. imagens: uma lista de imagens.
  2. verbose: Quando definido como 1, algumas mensagens de log são impressas.

O código a seguir chama o método detect(). Observe que o comprimento da lista atribuída ao argumento images deve ser igual ao tamanho do lote. Com base nas propriedades de configuração GPU_COUNT e IMAGES_PER_GPU que definimos, o tamanho do lote é 1. Assim, a lista deve ter uma única imagem. O resultado da detecção é retornado na variável r.

r = model.detect(images=[image], 
                 verbose=0)

Se mais de uma imagem for passada (por exemplo, images=[image, image, image]), a seguinte exceção será gerada indicando que o comprimento do argumento images deve ser igual para a propriedade de configuração BATCH_SIZE.

...
File "D:\Object Detection\Pure Keras\mrcnn\model.py", in detect
  assert len(images) == self.config.BATCH_SIZE, "len(images) must be equal to BATCH_SIZE"

AssertionError: len(images) must be equal to BATCH_SIZE

Para cada imagem de entrada, o método detect() retorna um dicionário que contém informações sobre os objetos detectados. Para retornar as informações sobre a primeira imagem alimentada ao modelo, então o índice 0 é utilizado com a variável r.

r = r[0]

O código a seguir imprime as chaves no dicionário. Existem 4 elementos no dicionário, com as seguintes chaves:

  1. rois: As caixas ao redor de cada objeto detectado.
  2. class_ids: os IDs de classe dos objetos.
  3. pontuações: as pontuações da classe para cada objeto.
  4. máscaras: As máscaras.
print(r.keys())

dict_keys(['rois', 'class_ids', 'scores', 'masks'])

6. Visualize os resultados

Assim que o método detect() for concluído, é hora de visualizar os objetos detectados. O script mrcnn.visualize é usado para esse propósito. A função mrcnn.visualize.display_instances() é usada para exibir as caixas de detecção, máscaras, nomes de classes e pontuações.

Dentre os parâmetros aceitos por esta função, são utilizados os seguintes:

  • image: A imagem na qual as caixas e máscaras de detecção são desenhadas.
  • caixas: As caixas de detecção.
  • masks: as máscaras detectadas.
  • class_ids: IDs de classe detectados.
  • class_names: uma lista de nomes de classes no conjunto de dados.
  • pontuações: pontuações de previsão para cada objeto.

No código a seguir, os nomes das classes são preparados na lista CLASS_NAMES. Observe que o rótulo da primeira classe é BG para o plano de fundo. A função mrcnn.visualize.display_instances() é chamada para exibir a imagem anotada.

import mrcnn.visualize

CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

r = r[0]

mrcnn.visualize.display_instances(image=image, 
                                  boxes=r['rois'], 
                                  masks=r['masks'], 
                                  class_ids=r['class_ids'], 
                                  class_names=CLASS_NAMES, 
                                  scores=r['scores'])

Após a execução da função, é mostrada uma figura (conforme abaixo) na qual são desenhadas as caixas, máscaras, notas das aulas e rótulos.

Até este ponto, são discutidas todas as etapas necessárias para usar o projeto Mask_RCNN para detectar objetos.

Código completo para previsão

O código completo para usar o projeto Mask_RCNN para detectar objetos em uma imagem está listado abaixo.

import mrcnn
import mrcnn.config
import mrcnn.model
import mrcnn.visualize
import cv2
import os

CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

class SimpleConfig(mrcnn.config.Config):
    NAME = "coco_inference"

    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    NUM_CLASSES = len(CLASS_NAMES)

model = mrcnn.model.MaskRCNN(mode="inference", 
                             config=SimpleConfig(),
                             model_dir=os.getcwd())

model.load_weights(filepath="mask_rcnn_coco.h5", 
                   by_name=True)

image = cv2.imread("sample2.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

r = model.detect([image], verbose=0)

r = r[0]

mrcnn.visualize.display_instances(image=image, 
                                  boxes=r['rois'], 
                                  masks=r['masks'], 
                                  class_ids=r['class_ids'], 
                                  class_names=CLASS_NAMES, 
                                  scores=r['scores'])

Até este ponto, o código para fazer previsões está completo. O restante do tutorial discute como treinar o modelo Mask R-CNN usando um conjunto de dados de treinamento personalizado. A próxima seção baixa o conjunto de dados.

Baixe o conjunto de dados de treinamento

Você deve ter um conjunto de dados para treinar um modelo de aprendizado de máquina ou aprendizado profundo. Para cada amostra nos dados de treinamento, pode haver dados verdadeiros. Esses dados podem ser simples, como um rótulo de classe, ou complexos, como os usados para um modelo de detecção de objetos.

Geralmente, os dados verdadeiros para modelos de detecção de objetos incluem uma caixa delimitadora e um rótulo de classe para cada objeto na imagem. Especificamente para o modelo Mask R-CNN, existe uma máscara adicional que marca os pixels pertencentes ao objeto.

Cada imagem pode ter mais de um objeto e, portanto, preparar os dados verdadeiros para um conjunto de dados inteiro é cansativo.

Nesta seção, um conjunto de dados existente de imagens Kangaroo é usado para treinar Mask R-CNN usando o projeto Mask_RCNN. O conjunto de dados Kangaroo pode ser baixado aqui. Ele vem com dados de anotação (ou seja, dados verdadeiros) e, portanto, está pronto para uso.

A próxima figura mostra uma imagem do conjunto de dados onde a caixa delimitadora prevista, a máscara e a pontuação são desenhadas. Observe que a máscara não é precisa, pois o modelo foi treinado para apenas uma época.

O conjunto de dados vem com 2 pastas:

  1. imagens: as imagens no conjunto de dados.
  2. anotações: as anotações para cada imagem como um arquivo XML separado.

A próxima seção prepara o conjunto de dados para uso posterior para treinar e validar o modelo Mask R-CNN.

Prepare o conjunto de dados de treinamento

O projeto Mask_RCNN possui uma classe chamada Dataset dentro do módulo mrcnn.utils. Esta classe simplesmente armazena informações sobre todas as imagens de treinamento nas listas. Quando os detalhes de todas as imagens forem armazenados em uma única estrutura de dados será mais fácil gerenciar o conjunto de dados.

Por exemplo, existe uma lista chamada class_info que contém informações sobre cada classe do conjunto de dados. Da mesma forma, uma lista chamada image_info contém informações sobre cada imagem. Para treinar o modelo Mask R-CNN, a lista image_info é usada para recuperar as imagens de treinamento e suas anotações. As anotações incluem caixas delimitadoras e rótulos de classe para todos os objetos dentro de cada imagem.

A classe mrcnn.utils.Dataset possui vários métodos úteis, que incluem:

  • add_class(): Adiciona uma nova classe.
  • add_image(): Adiciona uma nova imagem ao conjunto de dados.
  • image_reference(): A referência (por exemplo, caminho ou link) pela qual a imagem é recuperada.
  • prepare(): Depois de adicionar todas as classes e imagens ao conjunto de dados, este método prepara o conjunto de dados para uso.
  • source_image_link(): Retorna o caminho ou link da imagem.
  • load_image(): Lê e retorna uma imagem.
  • load_mask(): Carrega as máscaras dos objetos em uma imagem.

O próximo bloco de código cria uma instância vazia do mrcnn.utils.Dataset chamada KangaroDataset.

import mrcnn.utils

class KangarooDataset(mrcnn.utils.Dataset):
    pass

Dentro da nova classe, sinta-se à vontade para substituir qualquer um dos métodos mencionados anteriormente se precisar de personalização. Além disso, adicione qualquer novo método que possa ajudar.

De todos os métodos listados anteriormente, o método load_mask() deve ser substituído. A razão é que a recuperação das máscaras dos objetos difere com base no formato do arquivo de anotação e, portanto, não existe uma maneira única de carregar as máscaras. Como resultado, carregar as máscaras é uma tarefa que o desenvolvedor deve realizar.

No próximo bloco de código abaixo, construiremos 3 métodos:

  1. load_dataset(): Aceita o diretório onde existem as pastas images e annots, além de um parâmetro booleano que representa se o diretório se refere a os dados de treinamento ou validação.
  2. load_mask(): Este método carrega as máscaras do conjunto de dados Kangaroo. Aceita o ID da imagem no parâmetro image_id. O ID da imagem é apenas um valor único para cada imagem. Sinta-se à vontade para atribuir os IDs de sua escolha. O método retorna as máscaras e os IDs de classe de cada objeto. Observe que este conjunto de dados possui uma única classe que representa cangurus.
  3. extract_boxes: O método load_mask() chama o método extract_boxes() que é responsável por retornar as coordenadas de cada caixa delimitadora, além do largura e altura de cada imagem.

A implementação de cada um desses três métodos está listada no próximo bloco de código.

A primeira linha do método load_dataset() chama o método add_class() para criar uma classe chamada kangaroo com ID 1. Existe outra classe com ID 0, que é o background com o rótulo BG. Não precisamos adicioná-lo explicitamente porque ele existe por padrão. A última linha chama o método add_image() para adicionar uma imagem ao conjunto de dados. O ID da imagem, o caminho e o caminho do arquivo de anotação são passados para este método.

O método load_dataset() divide o conjunto de dados para que 150 imagens sejam usadas para treinamento, enquanto o restante é usado para teste.

import mrcnn.utils

class KangarooDataset(mrcnn.utils.Dataset):

    def load_dataset(self, dataset_dir, is_train=True):
    self.add_class("dataset", 1, "kangaroo")

    images_dir = dataset_dir + '/images/'
    annotations_dir = dataset_dir + '/annots/'

    for filename in os.listdir(images_dir):
      image_id = filename[:-4]

      if image_id in ['00090']:
        continue

      if is_train and int(image_id) >= 150:
        continue

      if not is_train and int(image_id) < 150:
        continue

      img_path = images_dir + filename
      ann_path = annotations_dir + image_id + '.xml'

      self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)

    def extract_boxes(self, filename):
    tree = xml.etree.ElementTree.parse(filename)

    root = tree.getroot()

    boxes = list()
    for box in root.findall('.//bndbox'):
      xmin = int(box.find('xmin').text)
      ymin = int(box.find('ymin').text)
      xmax = int(box.find('xmax').text)
      ymax = int(box.find('ymax').text)
      coors = [xmin, ymin, xmax, ymax]
      boxes.append(coors)

    width = int(root.find('.//size/width').text)
    height = int(root.find('.//size/height').text)
    return boxes, width, height

    def load_mask(self, image_id):
    info = self.image_info[image_id]
    path = info['annotation']
    boxes, w, h = self.extract_boxes(path)
    masks = zeros([h, w, len(boxes)], dtype='uint8')

    class_ids = list()
    for i in range(len(boxes)):
      box = boxes[i]
      row_s, row_e = box[1], box[3]
      col_s, col_e = box[0], box[2]
      masks[row_s:row_e, col_s:col_e, i] = 1
      class_ids.append(self.class_names.index('kangaroo'))
    return masks, asarray(class_ids, dtype='int32')

Com base na classe KangarooDataset, o código a seguir prepara o conjunto de dados de treinamento. Uma nova instância da classe é simplesmente criada. Para carregar as imagens é chamado o método load_dataset(), que aceita o caminho das imagens do dataset e suas anotações no parâmetro dataset_dir. Tudo isso além do sinalizador is_train. Se este fag for True, então os dados serão considerados dados de treinamento. Caso contrário, os dados serão usados para validação ou teste.

O método prepare() é chamado para preparar o conjunto de dados para uso. Ele apenas cria mais atributos sobre os dados, como número de classes, número de imagens e muito mais.

train_set = KangarooDataset()
train_set.load_dataset(dataset_dir='D:\kangaroo', is_train=True)
train_set.prepare()

Da mesma forma, o conjunto de dados de validação é preparado de acordo com o código a seguir. A única diferença é que o sinalizador is_train está definido como False.

valid_dataset = KangarooDataset()
valid_dataset.load_dataset(dataset_dir='D:\kangaroo', is_train=False)
valid_dataset.prepare()

A próxima seção prepara alguns parâmetros de configuração do modelo.

Preparar a configuração do modelo

Uma subclasse da classe mrcnn.config.Config deve ser criada para conter os parâmetros de configuração do modelo. O próximo código cria uma nova classe chamada KangarooConfig, que estende a classe mrcnn.config.Config.

Observe que o número de classes (NUM_CLASSES) está definido como 2 porque o conjunto de dados possui apenas 2 classes, que são BG (para o plano de fundo) e kangaroo.

import mrcnn.config

class KangarooConfig(mrcnn.config.Config):
    NAME = "kangaroo_cfg"

    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    NUM_CLASSES = 2

    STEPS_PER_EPOCH = 131

Após a preparação do conjunto de dados e da configuração do modelo, a próxima seção discute o treinamento do modelo Mask R-CNN usando o TensorFlow 1.0.

Treine a máscara R-CNN no TensorFlow 1.0

Esta seção pressupõe que uma versão do TensorFlow 1.0 esteja instalada e usada para executar o código mencionado. Também é possível criar um ambiente virtual no qual o TensorFlow 1.0 esteja instalado.

O próximo código cria uma instância da classe mrcnn.model.MaskRCNN, que constrói a arquitetura do modelo Mask R-CNN. O parâmetro mode é definido como 'training' para indicar que o modelo será treinado. Quando o modelo é carregado para treinamento, há camadas de entrada extras em comparação ao carregamento do modelo apenas para inferência (ou seja, previsão). As camadas extras contêm as imagens de entrada e suas anotações (por exemplo, caixas delimitadoras).

import mrcnn.model
model = mrcnn.model.MaskRCNN(mode='training', 
                             model_dir='./', 
                             config=KangarooConfig())

Uma vez criada a arquitetura do modelo, os pesos são carregados usando o método load_weights() de acordo com o próximo código. Este método aceita os 3 parâmetros a seguir:

  1. filepath: O caminho do arquivo de pesos. O arquivo mask_rcnn_coco.h5 pode ser baixado neste link.
  2. by_name: Se deseja atribuir os pesos das camadas de acordo com seus nomes.
  3. exclude: Os nomes das camadas para as quais não carregamos seus pesos. Estas são as camadas no topo da arquitetura, que mudam com base no tipo de problema (por exemplo, número de classes).

As camadas excluídas são as responsáveis por produzir as probabilidades de classe, caixas delimitadoras e máscaras.

model.load_weights(filepath='mask_rcnn_coco.h5', 
                   by_name=True, 
                   exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",  "mrcnn_bbox", "mrcnn_mask"])

Depois de carregar os pesos nas camadas do modelo, o modelo é treinado usando o método train(). Observe que o argumento layers é definido como heads para indicar que apenas as camadas no topo da arquitetura são treinadas. Você também pode especificar os nomes das camadas a serem treinadas.

Observe que o parâmetro exclude no método load_weights() aceita os nomes das camadas que não terão seus pesos carregados, mas o parâmetro layers em o método train() aceita os nomes das camadas a serem treinadas.

model.train(train_dataset=train_set, 
            val_dataset=valid_dataset, 
            learning_rate=KangarooConfig().LEARNING_RATE, 
            epochs=10, 
            layers='heads')

O tempo de treinamento difere com base no poder de computação da máquina. Se você executar o treinamento no Gradient com uma GPU poderosa e de baixo custo, ele deverá ser relativamente rápido. Depois que o modelo for treinado, os pesos treinados podem ser salvos usando o método Keras save_weights().

model_path = 'Kangaroo_mask_rcnn.h5'
model.keras_model.save_weights(model_path)

Este é o código completo para treinar um modelo.

import os
import xml.etree
from numpy import zeros, asarray

import mrcnn.utils
import mrcnn.config
import mrcnn.model

class KangarooDataset(mrcnn.utils.Dataset):

  def load_dataset(self, dataset_dir, is_train=True):
    self.add_class("dataset", 1, "kangaroo")

    images_dir = dataset_dir + '/images/'
    annotations_dir = dataset_dir + '/annots/'

    for filename in os.listdir(images_dir):
      image_id = filename[:-4]

      if image_id in ['00090']:
        continue

      if is_train and int(image_id) >= 150:
        continue

      if not is_train and int(image_id) < 150:
        continue

      img_path = images_dir + filename
      ann_path = annotations_dir + image_id + '.xml'

      self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)

  def extract_boxes(self, filename):
    tree = xml.etree.ElementTree.parse(filename)

    root = tree.getroot()

    boxes = list()
    for box in root.findall('.//bndbox'):
      xmin = int(box.find('xmin').text)
      ymin = int(box.find('ymin').text)
      xmax = int(box.find('xmax').text)
      ymax = int(box.find('ymax').text)
      coors = [xmin, ymin, xmax, ymax]
      boxes.append(coors)

    width = int(root.find('.//size/width').text)
    height = int(root.find('.//size/height').text)
    return boxes, width, height

  def load_mask(self, image_id):
    info = self.image_info[image_id]
    path = info['annotation']
    boxes, w, h = self.extract_boxes(path)
    masks = zeros([h, w, len(boxes)], dtype='uint8')

    class_ids = list()
    for i in range(len(boxes)):
      box = boxes[i]
      row_s, row_e = box[1], box[3]
      col_s, col_e = box[0], box[2]
      masks[row_s:row_e, col_s:col_e, i] = 1
      class_ids.append(self.class_names.index('kangaroo'))
    return masks, asarray(class_ids, dtype='int32')

class KangarooConfig(mrcnn.config.Config):
    NAME = "kangaroo_cfg"

    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    NUM_CLASSES = 2

    STEPS_PER_EPOCH = 131

train_set = KangarooDataset()
train_set.load_dataset(dataset_dir='kangaroo', is_train=True)
train_set.prepare()

valid_dataset = KangarooDataset()
valid_dataset.load_dataset(dataset_dir='kangaroo', is_train=False)
valid_dataset.prepare()

kangaroo_config = KangarooConfig()

model = mrcnn.model.MaskRCNN(mode='training', 
                             model_dir='./', 
                             config=kangaroo_config)

model.load_weights(filepath='mask_rcnn_coco.h5', 
                   by_name=True, 
                   exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",  "mrcnn_bbox", "mrcnn_mask"])

model.train(train_dataset=train_set, 
            val_dataset=valid_dataset, 
            learning_rate=kangaroo_config.LEARNING_RATE, 
            epochs=1, 
            layers='heads')

model_path = 'Kangaro_mask_rcnn.h5'
model.keras_model.save_weights(model_path)

Conclusão

Este tutorial apresentou o projeto Python de código aberto Mask_RCNN, que constrói o modelo Mask R-CNN para segmentação de instâncias de objetos. O projeto oferece suporte apenas a uma versão do TensorFlow $\geq$1.0. Este tutorial abordou as etapas para fazer previsões e treinar o modelo em um conjunto de dados personalizado.

Antes do modelo ser treinado, os conjuntos de dados de treinamento e validação são preparados usando uma classe filha da classe mrcnn.utils.Dataset. Depois de preparar os parâmetros de configuração do modelo, o modelo pode ser treinado.

Artigos relacionados: