ddompe/vit-waste-classification

Vision Transformer (ViT) para Clasificación de Residuos

Descripción

Modelo Vision Transformer (ViT) afinado sobre el dataset de segregación de residuos para clasificar imágenes en 9 categorías. Desarrollado como parte de un proyecto de aprendizaje automático en el Instituto Tecnológico de Costa Rica (ITCR).

Basado en google/vit-base-patch16-224-in21k con ajuste fino mediante transfer learning.

Resultados de Entrenamiento

Configuración

• Modelo base: google/vit-base-patch16-224-in21k
• Épocas: 3
• Tasa de aprendizaje: 0.0002
• Tamaño de lote: 16
• Optimizador: AdamW con warmup
• Regularización (weight decay): 0.01

Datos

• Entrenamiento: 3044 imágenes
• Validación: 762 imágenes
• Clases: 9

Métricas

• Pérdida de entrenamiento: 1.6462
• Evaluación: por época
• Criterio de guardado: mejor accuracy en validación

Metodología

El afinamiento se realizó en las siguientes etapas:

1. Preprocesamiento: normalización con ViTImageProcessor (224×224 px)
2. Transfer learning: pesos preentrenados en ImageNet-21k
3. Afinamiento: tasa de aprendizaje baja (2e-4) para preservar representaciones aprendidas
4. Regularización: weight decay (0.01) y warmup lineal (10% de pasos)
5. Parada anticipada: monitoreo de accuracy en validación

Clases

ID	Clase
0	Cardboard
1	Food Organics
2	Glass
3	Metal
4	Miscellaneous Trash
5	Paper
6	Plastic
7	Textile Trash
8	Vegetation

Uso

from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image

processor = ViTImageProcessor.from_pretrained("ddompe/vit-waste-classification")
model = ViTForImageClassification.from_pretrained("ddompe/vit-waste-classification")

imagen = Image.open("residuo.jpg").convert("RGB")
entradas = processor(images=imagen, return_tensors="pt")
salidas = model(**entradas)
clase_pred = salidas.logits.argmax(-1).item()
print(model.config.id2label[clase_pred])