Técnica de organização de dados em pastas para análise de imagens

Uma jornada de pré-processamento de dados em Pytorch

Durante uma etapa do processo de desenvolvimento de um modelo de Redes Neurais Artificiais Convolucionais (CNN), enfrentei o desafio de entender como seria feita a transformação da base de dados em tensores, que seriam o conjunto de dados responsáveis por realizar o treino e teste do modelo conexionista. Isto posto, decidi compreender e registrar todo o método que utilizei em Pytorch para definir minhas imagens e seus respectivos rótulos.

A finalidade deste artigo é explicar de forma simplificada o processo de transformação de dados. É importante mencionar que não abordarei definições sobre o que é uma classe ou função e apenas descreverei os processos que utilizei em meu trabalho de classificação de imagens de pessoas usando, não usando ou usando máscaras de forma incorreta, especificamente no contexto pandêmico.

Embasamento teórico

Imagine-se em um corredor com várias portas, e ao abrir cada uma delas, encontrará informações que deseja obter. Essa analogia se aplica especificamente ao processo de navegar em pastas em busca dos nossos dados.

Geralmente, ao trabalhar com imagens, o nome das subpastas em nosso conjunto de dados valida a classe ou rótulo que cada dado possui. Existem várias maneiras de trabalhar e agrupar essas informações corretamente, como utilizar um DataFrame, o Keras ou o Tensorflow. Além disso, podemos usar o ImageDataGenerator, que lê um diretório e já separa automaticamente os dados em treinamento e teste.

Leitura de imagens em Pytorch

Porém, utilizando o Pytorch e avaliando que não temos separado em diretórios diferentes o treino e teste, devemos buscar alternativas que sejam capazes de facilitar este processo e criar um out que nos dê os dados tratados que precisamos, seguem as etapas:

class ImageFolderDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory
        self.transform = transform

        self.class_to_idx = {}
        self.imgs, self.labels = [], []
        idx = 0
        for label in os.listdir(directory):
            if label not in self.class_to_idx:
                self.class_to_idx[label] = idx
                idx += 1
            for filename in os.listdir(os.path.join(directory, label)):
                file_path = os.path.join(directory, label, filename)
                if os.path.isfile(file_path):
                    self.imgs.append(file_path)
                    self.labels.append(self.class_to_idx[label])

De acordo com o código acima, definimos duas entradas, de forma respectiva, uma obrigatória, que é o diretório da nossa pasta de imagens, e outra opcional, que se refere a um possível pré-processamento de dados. Essa etapa pode ser responsável por realizar operações como augmentation ou resize nas nossas imagens, dependendo da situação.

Além disso, é possível observar que foram criadas duas variáveis, denominadas imgs e labels, que são responsáveis por armazenar todas as imagens contidas no diretório especificado, juntamente com os seus respectivos rótulos.

Consequentemente, é sabido que a leitura da base de dados pelo modelo conexionista é realizada por meio de tensores que possuem valores numéricos. Em algumas situações, como no meu caso, é necessário realizar uma espécie de rename, que tem como objetivo transformar o nome das pastas, que representam os rótulos, em valores numéricos. Isso é feito para atender à obrigatoriedade de entrada da base de dados no modelo convolucional.

class ImageFolderDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory
        self.transform = transform

        self.class_to_idx = {}
        self.imgs, self.labels = [], []
        idx = 0
        for label in os.listdir(directory):
            if label not in self.class_to_idx:
                self.class_to_idx[label] = idx
                idx += 1
            for filename in os.listdir(os.path.join(directory, label)):
                file_path = os.path.join(directory, label, filename)
                if os.path.isfile(file_path):
                    self.imgs.append(file_path)
                    self.labels.append(self.class_to_idx[label])

    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = self.labels[index] + 1

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, label

Portanto, adicionamos uma função à classe inicializada que irá percorrer cada imagem e seu respectivo rótulo, utilizando o índice definido pelo parâmetro da função para acessar cada elemento individualmente, como se estivéssemos recuperando uma informação específica.

Em seguida, cada imagem será convertida em três canais de cores, na ordem, vermelho, verde e azul (RGB). Caso seja fornecida uma transformação ou método augmentation, o “if” na função descrita executará essa transformação. Caso contrário, a imagem original será retornada sem nenhum processamento adicional.

class ImageFolderDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory
        self.transform = transform

        self.class_to_idx = {}
        self.imgs, self.labels = [], []
        idx = 0
        for label in os.listdir(directory):
            if label not in self.class_to_idx:
                self.class_to_idx[label] = idx
                idx += 1
            for filename in os.listdir(os.path.join(directory, label)):
                file_path = os.path.join(directory, label, filename)
                if os.path.isfile(file_path):
                    self.imgs.append(file_path)
                    self.labels.append(self.class_to_idx[label])

    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = self.labels[index] + 1

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, label

    def __len__(self):
        return len(self.imgs)

Por fim, criamos uma outra função capaz de nos fornecer a quantidade de imagens que estamos passando por meio da classe.

Validação de subpastas

Como podemos perceber, as etapas descritas acima são o que temos disponível na documentação do Pytorch como etapa de leitura a partir de pastas que irão conter nossos dados e rótulos. Porém, ao utilizar apenas esta classe, certos arquivos não serão lidos (os mais comuns sim) devido a diversos problemas, por exemplo, subpastas com mesmo nome, a imagem abaixo retrata o problema que ocorreu no desenvolvimento de um dos meus projetos:

Caso você utilize uma API para baixar a base de dados diretamente do Kaggle e decida renomear ou não manualmente, a função pode acabar engolindo o rótulo da imagem e retornar apenas duas classes, além de ser custoso renomear a cada leitura da API. Nesse sentido, foi necessário adicionar nos códigos descritos neste artigo um processo capaz de distinguir subpastas com nomes iguais, como podemos visualizar a seguir.

#Criamos um dict capaz de armazenar as subpastas de cada classe
        subfolders_dict = {}

        for label in os.listdir(directory):
            if label not in self.class_to_idx:
                self.class_to_idx[label] = idx
                idx += 1

            #Verificamos se alguma subpasta está presente em mais de um rótulo
            for subfolder in os.listdir(os.path.join(directory, label)):
                if subfolder not in subfolders_dict:
                    subfolders_dict[subfolder] = [label]
                else:
                    subfolders_dict[subfolder].append(label)

            #Adicionamos as imagens e rótulos à lista de dados que declaramos
            for filename in os.listdir(os.path.join(directory, label)):
                file_path = os.path.join(directory, label, filename)
                if os.path.isfile(file_path):
                    self.imgs.append(file_path)
                    self.labels.append(self.class_to_idx[label])

        #Verificamos se existe possibilidade de conflito e adiciona o rótulo extra
        extra_label = idx
        self.class_to_idx["extra"] = extra_label
        for subfolder, labels in subfolders_dict.items():
            if len(labels) > 1:
                for i in range(len(self.labels)):
                    if self.class_to_idx[labels[0]] == self.labels[i]:
                        if subfolder in self.imgs[i]:
                            self.labels[i] = extra_label

O que foi feito foi uma organização das pastas e subpastas usando um dicionário. De maneira simplificada, o processo itera sobre as pastas e verifica quais existem ou não no dicionário. Se uma subpasta com o mesmo nome já tiver sido armazenada, o processo adicionará o rótulo atual ao nome da subpasta para diferenciá-la. Essa organização resolveu o problema encontrado durante o treinamento do modelo, quando os dados foram carregados no DataLoader do PyTorch. Segue o código completo da função de leitura de pastas de imagens em Pytorch:

class ImageFolderDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory
        self.transform = transform

        self.class_to_idx = {}
        self.imgs, self.labels = [], []
        idx = 0
        #Criamos um dict capaz de armazenar as subpastas de cada classe
        subfolders_dict = {}

        for label in os.listdir(directory):
            if label not in self.class_to_idx:
                self.class_to_idx[label] = idx
                idx += 1

            #Verificamos se alguma subpasta está presente em mais de um rótulo
            for subfolder in os.listdir(os.path.join(directory, label)):
                if subfolder not in subfolders_dict:
                    subfolders_dict[subfolder] = [label]
                else:
                    subfolders_dict[subfolder].append(label)

            #Adicionamos as imagens e rótulos à lista de dados que declaramos
            for filename in os.listdir(os.path.join(directory, label)):
                file_path = os.path.join(directory, label, filename)
                if os.path.isfile(file_path):
                    self.imgs.append(file_path)
                    self.labels.append(self.class_to_idx[label])

        #Verificamos se existe possibilidade de conflito e adiciona o rótulo extra
        extra_label = idx
        self.class_to_idx["extra"] = extra_label
        for subfolder, labels in subfolders_dict.items():
            if len(labels) > 1:
                for i in range(len(self.labels)):
                    if self.class_to_idx[labels[0]] == self.labels[i]:
                        if subfolder in self.imgs[i]:
                            self.labels[i] = extra_label

    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = self.labels[index] + 1

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, label

    def __len__(self):
        return len(self.imgs)

Conclusão

Vale salientar que existem diversas formas de contornar a mesma problemática que eu tive, mas decidi documentar aqui o processo que aprendi durante uma semana de entendimento de como pode funcionar a alocação de classes por meio da leitura de pastas.

Se você enfrentou o mesmo problema descrito neste artigo, lembre-se de incluir suas transformações (augmentações) como parâmetro na função de leitura pertencente à classe que foi desenvolvida. Abaixo, apresento como ficou a continuação e implementação da leitura e pré-processamento das imagens que utilizei no classificador de uso de máscaras no combate à Covid-19.

#Inicialmente crio uma varável que irá transformar nossas imagens
transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

#Já separo em dados de treino e teste conforme o desejado
train_data = []
test_data = []
train_size = 0.8

#Realiza e chama nossa função que trabalhamos
for folder in subfolders:
    dataset = ImageFolderDataset(folder, transform=transforms)
    
    #Dividimos o conjunto de dados em treino e teste
    n = len(dataset)
    n_train = int(train_size * n)
    n_test = n - n_train
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [n_train, n_test])
    
    train_data.append(train_dataset)
    test_data.append(test_dataset)

#Muito provavel, que ainda não tenhamos um pré-processamento capaz de linkar ao modelo
#Por isso, iremos utilizar uma função de concatenação e posteriormente criar nossos tensores Pytorch

train_data = torch.utils.data.ConcatDataset(train_data)
test_data = torch.utils.data.ConcatDataset(test_data)

#Tensores DataLoader
train_dataloader = DataLoader(train_data, batch_size=32, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=32, shuffle=False)

#Multiplicamos por 32 para sabermos o total geral não particionado de nossa base de dados
print(len(train_dataloader)*32, len(test_dataloader)*32)