Carlos012/vit_papas

Clasificador de calidad de papas con Vision Transformer

Descripción

Este modelo clasifica la calidad de tubérculos de papa mediante visión por computadora. Está basado en la arquitectura Vision Transformer (ViT), específicamente en google/vit-base-patch16-224, y fue entrenado con un conjunto híbrido de 36 000 imágenes de papa chaucha y papa chola.

El objetivo del modelo es apoyar la clasificación automática de papas en dos categorías:

Clase	Descripción
Buen estado	Tubérculos de papa sin defectos visibles.
Defectuoso	Tubérculos de papa con cortes, brotes o signos de pudrición visibles.

Datos y arquitectura

• Tarea: clasificación binaria de imágenes.
• Modelo base: Vision Transformer base con parches de 16 x 16 píxeles.
• Resolución de entrada: 224 x 224 píxeles.
• Conjunto de datos: conjunto híbrido de imágenes de papas.
• División de datos: entrenamiento, validación y prueba.
• Métricas evaluadas: exactitud, precisión, sensibilidad, puntaje F1, coeficiente de correlación de Matthews (MCC) y pérdida.

Estrategias de entrenamiento

Se evaluaron dos estrategias de entrenamiento: Transfer Learning y Fine-Tuning parcial.

Transfer Learning

En esta etapa se congelaron todas las capas del modelo base y se entrenó únicamente el clasificador final.

Hiperparámetro	Valor
Épocas	30
Tamaño de lote	32
Tasa de aprendizaje	1e-4
Optimizador	AdamW
Scheduler	Cosine Annealing
Weight decay	0.01
Betas	(0.9, 0.999)
Fracción de calentamiento	0.1
Precisión mixta	Sí
Capas congeladas	Todas las capas

Fine-Tuning parcial

En esta etapa se descongeló parte del modelo para ajustar mejor las representaciones visuales al dominio de las papas.

Hiperparámetro	Valor
Épocas	20
Tamaño de lote	32
Tasa de aprendizaje	1e-5
Optimizador	AdamW
Scheduler	Cosine Annealing
Weight decay	0.01
Betas	(0.9, 0.999)
Fracción de calentamiento	0.1
Precisión mixta	Sí
Capas congeladas	6 primeras capas

Resultados

Los resultados muestran que el Fine-Tuning parcial obtuvo el mejor desempeño general. Esta estrategia alcanzó métricas cercanas a 1.0 en validación y prueba, con una pérdida considerablemente menor que la obtenida con Transfer Learning.

Validación final

Técnica de entrenamiento	Accuracy	Precision	Recall	F1-score	MCC	Pérdida
Transfer Learning	0.9628	0.9631	0.9628	0.9628	0.9260	0.1319
Fine-Tuning parcial	0.9995	0.9995	0.9995	0.9995	0.9990	0.0026

Evaluación en prueba

Técnica de entrenamiento	Accuracy	Precision	Recall	F1-score	MCC	Pérdida
Transfer Learning	0.9600	0.9602	0.9600	0.9600	0.9202	0.1353
Fine-Tuning parcial	0.9999	0.9999	0.9999	0.9999	0.9997	0.0015

Matriz de confusión

La matriz de confusión compara el desempeño de ambas estrategias. El modelo entrenado con Fine-Tuning parcial presenta una separación casi perfecta entre las clases, mientras que el modelo con Transfer Learning registra una mayor cantidad de errores de clasificación.

Archivos principales

Recurso	Ruta
Pesos de Transfer Learning	modelo_v3_hibrido/hibrido_v2/transfer_learning/vit_papas_pesos_tl.pt
Pesos de Fine-Tuning parcial	modelo_v3_hibrido/hibrido_v2/fine_tuning_parcial/vit_papas_pesos_ftp.pt
Métricas de prueba de Transfer Learning	modelo_v3_hibrido/hibrido_v2/resultados/metricas_test_tl.csv
Métricas de prueba de Fine-Tuning parcial	modelo_v3_hibrido/hibrido_v2/resultados/metricas_test_ftp.csv
Curvas y gráficas de evaluación	modelo_v3_hibrido/hibrido_v2/resultados/

Uso del modelo

Para utilizar el modelo, se debe cargar la arquitectura ViT base, aplicar el mismo preprocesamiento usado durante el entrenamiento y cargar los pesos correspondientes a la estrategia seleccionada.

import torch
from PIL import Image
from transformers import ViTForImageClassification, ViTImageProcessor

RUTA_PESOS = "modelo_v3_hibrido/hibrido_v2/fine_tuning_parcial/vit_papas_pesos_ftp.pt"
RUTA_IMAGEN = "ruta/a/la/imagen.jpg"

id2label = {
    0: "Buen estado",
    1: "Defectuoso",
}

processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
modelo = ViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224",
    num_labels=2,
    id2label=id2label,
    label2id={label: idx for idx, label in id2label.items()},
    ignore_mismatched_sizes=True,
)

checkpoint = torch.load(RUTA_PESOS, map_location="cpu")
estado = checkpoint.get("model_state_dict", checkpoint)
estado = estado.get("state_dict", estado)

modelo.load_state_dict(estado)
modelo.eval()

imagen = Image.open(RUTA_IMAGEN).convert("RGB")
entradas = processor(images=imagen, return_tensors="pt")

with torch.no_grad():
    salidas = modelo(**entradas)
    clase_predicha = salidas.logits.argmax(dim=1).item()

print(f"Predicción: {id2label[clase_predicha]}")

Recomendaciones

• Usar preferentemente los pesos de Fine-Tuning parcial, ya que presentan el mejor rendimiento en prueba.
• Mantener el mismo tamaño de entrada y normalización del modelo base para evitar diferencias entre entrenamiento e inferencia.
• Validar el modelo con imágenes tomadas en condiciones reales de iluminación, fondo y distancia antes de integrarlo en un flujo de clasificación real.