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:
- PyTorch (https://pytorch.org/)
- Streamlit (https://streamlit.io/)
- Pillow (https://python-pillow.org/)
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.
Top comments (0)