Treinar um classificador de imagens

A estrutura de pastas do conjunto de dados#

Os conjuntos de dados de classificação são diretórios, não arquivos YAML. Aponte data= para a pasta raiz:

parts/
├── train/
│   ├── good/      img_0001.jpg ...
│   ├── scratch/   img_0412.jpg ...
│   └── dent/      img_0733.jpg ...
└── val/
    ├── good/
    ├── scratch/
    └── dent/

O que o carregador faz com ela, de acordo com o código-fonte instalado e o guia de conjuntos de dados de classificação:

Na Ultralytics Platform, um conjunto de dados de classificação contém um rótulo por imagem, em vez de formas:

Treinar com MNIST160#

MNIST160 contém 160 dígitos manuscritos, 80 em train/ e 80 em test/, cada conjunto dividido em dez pastas de dígitos — pequeno o bastante para um teste rápido na CPU. Não há uma pasta val/, então observe o aviso do carregador e a utilização de test/ como alternativa:

from ultralytics import YOLO

model = YOLO("yolo26n-cls.pt")
model.train(data="mnist160", epochs=1, imgsz=64, device="cpu")
metrics = model.val()
print(f"top-1: {metrics.top1:.3f}  top-5: {metrics.top5:.3f}")

WARNING ⚠️ Dataset 'split=val' not found, using 'split=test' instead.
val: .../datasets/mnist160/test... found 80 images in 10 classes ✅
                   all      0.162      0.463
top-1: 0.162  top-5: 0.463

Uma época mal altera a nova cabeça de dez classes. Na nossa execução, epochs=10 alcançou top-1 de 0.55 e top-5 de 0.95 em cerca de dez segundos na CPU de um laptop — o ciclo funciona; só precisa de um cronograma de treinamento. O guia da tarefa Classify lista os tamanhos pré-treinados.

Top-1 é a proporção de imagens cuja classe com maior probabilidade está correta. Top-5 conta como acerto se a classe correta estiver entre as cinco mais prováveis. Top-5 é informativo para o ImageNet de 1.000 classes, cujas classes vão das mais abrangentes às muito específicas; com cinco classes ou menos, o valor é sempre 1.0, e com dez é generoso. Use a acurácia top-1 e uma matriz de confusão por classe como critérios.

Ler as probabilidades no código#

from ultralytics import YOLO

model = YOLO("yolo26n-cls.pt")  # ImageNet, 1,000 classes
result = model("https://ultralytics.com/images/bus.jpg")[0]

probs = result.probs
print(f"top-1: {result.names[probs.top1]} ({float(probs.top1conf):.2f})")
for i, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"  {result.names[i]:20s} {conf:.3f}")

top-1: minibus (0.52)
  minibus              0.516
  police_van           0.369
  trolleybus           0.096
  minivan              0.006
  ambulance            0.003

probs.top1 é um índice, probs.top5 é uma lista de índices e probs.data é o vetor completo de probabilidades. Observe a confiança top-1: 0.52 é só um pouco mais do que jogar uma moeda para decidir entre minibus e police_van. É esse valor que você usa para encaminhar as previsões.

Encaminhar previsões incertas para uma pessoa#

Um classificador em produção raramente precisa acertar todas as imagens; precisa acertar as imagens que classifica sem ajuda. Este script treina o modelo por 10 épocas e, em seguida, mede o efeito de diferentes limites de confiança nas 80 imagens de dígitos reservadas para teste:

from pathlib import Path

from ultralytics import YOLO

model = YOLO("yolo26n-cls.pt")
model.train(data="mnist160", epochs=10, imgsz=64, device="cpu", plots=False)  # model now holds best.pt
print(model.names)  # folder names, sorted, became class indices

test = Path(model.trainer.data["test"])  # mnist160/test, the split the loader validated on
images = sorted(test.rglob("*.png"))
results = model(images, imgsz=64, verbose=False)
truth = [p.parent.name for p in images]

for threshold in (0.0, 0.3, 0.6):
    kept = [(r, t) for r, t in zip(results, truth) if r.probs.top1conf >= threshold]
    correct = sum(r.names[r.probs.top1] == t for r, t in kept)
    print(f"threshold {threshold:.1f}: auto-accept {len(kept):2d}/{len(images)}, accuracy on accepted {correct / max(len(kept), 1):.0%}")

{0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9'}
threshold 0.0: auto-accept 80/80, accuracy on accepted 55%
threshold 0.3: auto-accept 41/80, accuracy on accepted 68%
threshold 0.6: auto-accept 10/80, accuracy on accepted 80%

Cada linha representa uma troca: aumentar o limite torna as previsões aceitas mais precisas e encaminha mais imagens para uma pessoa. Com 0.6, este modelo decide sozinho apenas 10 dos 80 dígitos — ainda não está pronto para automatizar nada, e a solução é treinar mais, não escolher um limite mais engenhoso. Gere a mesma tabela novamente após cada rodada de treinamento: ela mostra tanto a acurácia que você terá quanto o volume exato de trabalho de revisão que cada limite gera.

Armadilhas#

Os classificadores são pré-treinados em imgsz=224, não em 640. Para objetos pequenos em um quadro grande, é melhor detectar e depois classificar (próxima lição) do que usar um imgsz maior.