Hola amigos, hoy les traigo este tutorial de como podemos crear un clasificador de imágenes implementado una red neural convolucional conocida como CNN del inglés Convolutional Neural Network.

Una red neuronal convolucional es un tipo de red neuronal artificial diseñada específicamente para procesar datos con estructura de cuadrícula, como imágenes. Utiliza capas de convolución para detectar patrones y características en los datos de entrada, seguidas de capas de agrupación que reducen la dimensionalidad. Las CNN son ampliamente utilizadas en tareas de visión por computadora, como reconocimiento de imágenes y clasificación, debido a su capacidad para aprender características jerárquicas y invariantes a la traslación.

Para lograr esto usaremos como lenguaje principal Python, y otras librerías como:

La implementación completa la puedes encontrar en mi cuenta de GitHub en el siguiente repositorio:

https://github.com/emmendoza2794/basic-image-classifier

Este proyecto se encuentra dividido en 3 partes:

La interfaz gráfica realizada con Streamlit
El predictor de imágenes a partir de un modelo ya entrenado
El entrenador de un nuevo modelo a partir de nuevas imágenes

1. Interfaz gráfica

Para este proyecto tenemos 2 secciones, una en donde probamos nuestros modelos ya entrenados y otra en donde entrenamos un nuevo modelo.

Interfaz de prueba de modelos:

Interfaz de entrenamiento de nuevo modelo:

Nuestro código en Python seria este:

import streamlit as st  

from src.predictor import Predictor  
from src.train import Train  
from src.utils import Utils  

if 'prediction_result' not in st.session_state:  
    st.session_state.prediction_result = None  

if 'classes_list' not in st.session_state:  
    st.session_state.classes_list = None  

if 'train_result' not in st.session_state:  
    st.session_state.train_result = None  


def predict_image():  
    if st.session_state.image_file is None:  
        st.error("No image file")  
        return  

  st.session_state.prediction_result = Predictor().predict(  
        name_model=st.session_state.model,  
        image=st.session_state.image_file  
    )  


def get_classes():  
    if st.session_state.classes_list is not None:  
        classes = Predictor().load_classes(  
            name_model=st.session_state.model  
        )  

        st.session_state.classes_list = ', '.join(classes)  

    else:  
        st.session_state.classes_list = "There are no models to test, download the test models"  

def train_model():  

    st.session_state.train_result = Train().train_model(  
        epochs=st.session_state.epochs,  
        model_name=st.session_state.model_name  
    )  


st.set_page_config(  
    page_title="Basic Image Classifier",  
    page_icon="👋",  
    layout="wide",  
)  

st.header('Test model', divider='rainbow')  

col1, col2, col3 = st.columns([0.3, 0.5, 0.2])  

with col1:  
    with st.container(border=True):  
        st.subheader("Image Classifier")  

        st.selectbox(  
            label='Select model',  
            key="model",  
            options=Predictor().load_list_models()  
        )  

        get_classes()  

        st.markdown(f'**Classes:** {st.session_state.classes_list}')  

        st.file_uploader(  
            label="image file",  
            key="image_file",  
            type=['jpg', 'jpeg', 'png', 'bmp', 'gif', 'webp']  
        )  

        st.button(  
            label="Predict category",  
            type="primary",  
            on_click=predict_image,  
            use_container_width=True,  
        )  

    with st.container(border=True):  
        st.write("download example models for testing")  
        st.button(  
            label="Download models",  
            type="primary",  
            on_click=Utils().download_models,  
            use_container_width=True,  
        )  

with col2:  
    with st.container(border=True, height=600):  
        st.subheader("Image preview")  

        if st.session_state.image_file is not None:  
            st.image(  
                image=st.session_state.image_file  
            )  

with col3:  
    with st.container(border=True, height=600):  
        st.subheader("Prediction result")  

        if st.session_state.prediction_result is not None:  

            data_result = st.session_state.prediction_result  
            first_result = next(iter(data_result.items()))  

            st.subheader(f":green[{first_result[0]} {round(first_result[1], 2)}%]")  
            st.write(data_result)  


st.header('Train model', divider='rainbow')  


col4, col5 = st.columns([0.3, 0.7])  

with col4:  
    with st.container(border=True):  
        st.subheader("Configuration")  

        st.markdown('''  
 :red[**Important:** to train the model with your own data or another dataset, first copy the images separated by        folders into the root folder "image_files". For example:]   
        ''')  

        st.image("assets/image_files.png")  

        st.text_input(  
            label="Model name",  
            value="model_test",  
            key="model_name"  
  )  

        st.slider(  
            label="Epochs",  
            min_value=10,  
            max_value=100,  
            value=40,  
            key="epochs",  
            step=1  
  )  

        st.markdown("*At least 40 epochs are recommended to have an accuracy > 80%*")  

        st.button(  
            label="Train model",  
            type="primary",  
            on_click=train_model,  
            use_container_width=True,  
        )  


with col5:  
    with st.container(border=True, height=600):  
        st.subheader("Training result")  

        if st.session_state.train_result is not None:  
            st.write(st.session_state.train_result)

2. Predictor de imágenes a partir de un modelo entrenado

En esta parte usamos Pytorch para cargar el modelo y transformar la imagen a predecir de tal manera que nuestro modelo pueda usarla.

Para esto primero debemos cargar un modelo ya entrenado, en la interfaz tenemos un botón para descargar estos modelos de prueba, estos modelos se descargan en la carpeta raíz “models”, por cada modelo abran 2 archivos, uno con las clases posibles que el modelo puede predecir que será un archivo .json y otro con el formato .pth que es un archivo que contiene los parámetros entrenados del modelo.

La clase principal en donde definimos nuestro modelo es la siguiente:

from torch import nn  
import torch.nn.functional as F  


class Net(nn.Module):  
    def __init__(self, num_classes: int):  
        super().__init__()  
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)  
        self.pool = nn.MaxPool2d(2, 2)  
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  
        self.dropout = nn.Dropout(0.5)  
        self.fc1 = nn.Linear(64 * 64 * 64, 128)  
        self.fc2 = nn.Linear(128, num_classes)  

    def forward(self, x):  
        x = self.pool(F.relu(self.conv1(x)))  
        x = self.pool(F.relu(self.conv2(x)))  
        x = x.view(-1, 64 * 64 * 64)  
        x = F.relu(self.fc1(x))  
        x = self.dropout(x)  
        x = self.fc2(x)  
        return F.log_softmax(x, dim=1)

Esta clase es la misma que usamos para entrenar como para hacer las predicciones, es como la estructura que sigue nuestro modelo cuando entrena y cuando va a realizar una predicción.

Ahora bien, para realizar la predicción tenemos que instanciar nuestro modelo y cargarle el archivo .pth con los parámetros del modelo entrenado, esto lo hacemos con el siguiente código:

import json  
import os  

import streamlit as st  
import torch  
from PIL import Image  
from torchvision import transforms  
import torch.nn.functional as F  
from src.model import Net  


@st.cache_resource  
def load_model(num_classes: int):  
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  
    net = Net(num_classes=num_classes)  
    net.to(device=device)  

    return net  


class Predictor:  

    def __init__(self):  
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  


    def load_list_models(self):  
        list_models = []  
        for file_name in os.listdir('models'):  
            if file_name.endswith(".pth"):  
                list_models.append(file_name)  

        return list_models  


    def load_classes(self, name_model):  

        name_model = name_model.replace(".pth", "")  

        with open(f"models/classes_{name_model}.json", "r") as file:  
            classes = json.load(file)  
        return classes  

    def predict(self, name_model, image):  

        if image is None:  
            print("No image file")  
            return  

  image = Image.open(image)  

        transform = transforms.Compose([  
            transforms.Resize((256, 256)),  
            transforms.RandomHorizontalFlip(p=0.5),  
            transforms.RandomVerticalFlip(p=0.5),  
            transforms.RandomRotation(degrees=15),  
            transforms.ToTensor(),  
            transforms.Normalize((0.5,), (0.5,))  
        ])  

        image = transform(image)  
        image = image.unsqueeze(0)  

        image = image.to(device=self.device)  

        classes = self.load_classes(name_model=name_model)  

        model = load_model(num_classes=len(classes))  

        model.load_state_dict(torch.load(f"models/{name_model}"))  

        model.eval()  

        with torch.no_grad():  
            output = model(image)  
            probabilities = F.softmax(output, dim=1)  
            percentage = probabilities * 100  

  values = percentage[0].tolist()  

        results = dict(zip(classes, values))  
        results = dict(sorted(results.items(), key=lambda item: item[1], reverse=True))  

        return results

Básicamente, lo que hacemos es en la función predict, recibimos la imagen y le aplicamos unas transformaciones para que sea compatible con nuestro modelo, luego cargamos el modelo que tenemos seleccionado en la interfaz, cargamos las clases posibles que puede predecir el modelo y finalmente retornamos los resultados de la predicción.

2. Entrenar un modelo nuevo desde cero

Para poder entrenar un nuevo modelo desde cero primero debemos tener las imágenes separadas en carpetas y que estas carpetas sean las clases, podemos encontrar muchos datasets con imágenes ya clasificadas en páginas como https://www.kaggle.com/

Una vez tengamos nuestras imágenes, debemos copiarlas a la carpeta que está en la raíz del proyecto llamada “images_files”, por ejemplo:

Antes de cargar las imágenes a nuestro modelo tenemos una función que busca y elimina archivos corruptos, esto es importante porque si hay archivos dañados nuestro modelo no podrá entrenarse, el código que hace eso es el siguiente:

def clean_corrupt_images(self, root_dir: str):  
    for subdir, dirs, files in os.walk(root_dir):  
        for file in files:  
            file_path = os.path.join(subdir, file)  
            try:  
                with Image.open(file_path) as img:  
                    img.verify()  
            except (IOError, SyntaxError):  
                print(f'corrupt image: {file_path}')  
                os.remove(file_path)

Ahora si tenemos todo listo para empezar nuestro entrenamiento, desde la interfaz podemos darle un nombre al nuevo modelo y la cantidad de épocas que queremos que nuestro modelo se entrene, nuestro código para el entrenamiento es el siguiente:

import json  
import streamlit as st  
from src.model import Net  
from src.utils import Utils  
import torch  
from torchvision.transforms import transforms  
from torchvision.datasets import ImageFolder  
from torch.utils.data import DataLoader, random_split  
import torch.nn as nn  
import torch.optim as optim  


class Train:  
    def __init__(self):  
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  

    def _evaluate_model(self, model, data_loader):  
        model.eval()  
        correct = 0  
  total = 0  
  with torch.no_grad():  
            for data in data_loader:  
                images, labels = data[0].to(self.device), data[1].to(self.device)  
                outputs = model(images)  
                _, predicted = torch.max(outputs.data, 1)  
                total += labels.size(0)  
                correct += (predicted == labels).sum().item()  
        accuracy = 100 * correct / total  
        return accuracy  

    def train_model(self, epochs: int, model_name: str):  
        transform = transforms.Compose([  
            transforms.Resize((256, 256)),  
            transforms.RandomHorizontalFlip(p=0.5),  
            transforms.RandomVerticalFlip(p=0.5),  
            transforms.RandomRotation(degrees=15),  
            transforms.ToTensor(),  
            transforms.Normalize((0.5,), (0.5,))  
        ])  

        Utils().clean_corrupt_images(root_dir="images_files")  

        data_files = ImageFolder(root="images_files", transform=transform)  

        classes = data_files.classes  

        with open(f'models/classes_{model_name}.json', 'w') as file:  
            file.write(json.dumps(classes, indent=4))  

        train_size = int(0.85 * len(data_files))  
        test_size = len(data_files) - train_size  
        train_set, test_set = random_split(data_files, [train_size, test_size])  

        train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)  
        test_loader = DataLoader(test_set, batch_size=32, shuffle=True, num_workers=2)  

        model = Net(num_classes=len(classes))  
        model.to(device=self.device)  

        criterion = nn.CrossEntropyLoss()  
        optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)  

        train_result = []  

        progress_bar = st.progress(0, "Initializing training...")  

        for epoch in range(epochs):  

            model.train()  
            running_loss = 0.0  
  for i, data in enumerate(train_loader, 0):  
                inputs, labels = data[0].to(self.device), data[1].to(self.device)  

                optimizer.zero_grad()  

                outputs = model(inputs)  
                loss = criterion(outputs, labels)  
                loss.backward()  
                optimizer.step()  

                running_loss += loss.item()  

            accuracy = self._evaluate_model(model, test_loader)  

            info = f'Epoch {epoch + 1} -> Loss: {round(running_loss / len(train_loader), 2)} - Accuracy: {round(accuracy, 2)}%'  

  print(info)  

            percentage = (epoch + 1) / epochs  

  progress_bar.progress(percentage, info)  

            train_result.append({  
                'epoch': epoch + 1,  
                'loss': running_loss / len(train_loader),  
                'accuracy': accuracy  
            })  

        print('Finished Training')  

        progress_bar.empty()  

        torch.save(model.state_dict(), f'models/{model_name}.pth')  

        return train_result

En nuestra función tran_model seguimos los siguientes pasos:

Definir las transformaciones que le vamos a hacer a nuestro dataset de imágenes, que son las mismas transformaciones que usamos para la predicción
Luego limpiamos los archivos corruptos
Usamos los nombres de las carpetas para definir las clases posibles del modelo y lo guardamos como json
Definimos los DataLoaders de prueba y entrenamiento, usamos un 85% para entrenamiento y un 15% para pruebas
Definimos nuestra función de perdida para el entrenamiento del modelo y el optimizador que se utilizara para ajustar los parámetros del modelo durante el entrenamiento
Finalmente, empezamos a entrenar nuestro modelo con la cantidad de épocas que definimos en la interfaz, durante el entrenamiento, al finalizar cada época evaluamos nuestro modelo, esto con el fin de poder visualizar como avanza el entrenamiento en cada ciclo, cuando termina guardamos los parámetros del entrenamiento en un archivo .pth para luego usarlo en la vista de pruebas de nuestros modelos

Conclusiones

Gracias a las CNN podemos entrenar modelos para la clasificación de imágenes de una manera muy rápida y efectiva, logrando en nuestro caso una predicción de las del 85% sin importar el tamaño de nuestra imagen ni otros factores que gracias a los transformers de Pytorch podemos obviar, esto hace que tengamos muchos escenarios en donde podamos implementar este tipo de tecnologías.

Muchas gracias por leer mi post, hasta la próxima.

Blog

Clasificador de imágenes con una red neuronal convolucional (CNN)

Edinson Mauricio Mendoza