Salta el contingut

Classificació d'imatges amb YOLO

La classificació d'imatges és la tasca més senzilla de les tres i consisteix a classificar una imatge completa en una de les classes predefinides.

El resultat d'un classificador d'imatges és una etiqueta de classe i una puntuació de confiança. La classificació d'imatges és útil quan només necessiteu saber a quina classe pertany una imatge i no necessiteu saber on es troben els objectes d'aquesta classe ni quina és la seva forma exacta.

De manera molt similar a la detecció d'objectes, amb YOLO podem realitzar classificació d'imatges utilitzant models específics dissenyats per a aquesta tasca. Els models responen a un nom del tipus yolov8n-cls.pt, on el sufix -cls indica que és un model de classificació.

Python
1
2
3
4
5
6
7
8
from ultralytics import YOLO

# Load a model
model = YOLO("yolo26n-cls.pt")  # load an official model

# Predict with the model
results = model("https://ultralytics.com/images/bus.jpg")  # predict on an image
print(results)  # print results

Anàlisi dels resultats

De nou tenim un objecte Results que conté tota la informació sobre la predicció. Dins dels distints Result tenim un atribut probs que conté les probabilitats de cada classe per a la imatge d'entrada en un array probs.data

Cada element indica la probabilitat de classificar dita imatge en una classe concreta. Per exemple, si tenim 100 classes, probs serà un array de 100 elements on cada element representa la probabilitat que la imatge pertanyi a aquesta classe.

Per a facilitat la classificació, també tenim probs.top5 i probs.top5conf que ens retorna les 5 classes més probables juntament amb les seves probabilitats associades.

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
for result in results:
    probs = result.probs

    print(f"All class probabilities:")
    for idx, prob in enumerate(probs.data):
        class_name = result.names[idx]
        percentage = prob.item() * 100
        print(f"  {class_name}: {percentage:.4f}%")

    top5_indices = probs.top5
    top5_confidences = probs.top5conf

    print("Top 5 predictions:")
    for idx, conf in zip(top5_indices, top5_confidences):
        class_name = result.names[idx]
        percentage = conf.item() * 100
        print(f"  {class_name}: {percentage:.2f}%")

En aquest exemple, iterem sobre les probabilitats de cada classe i imprimim el nom de la classe juntament amb la seva probabilitat en format percentual. A continuació, obtenim les 5 classes més probables i les seves confidències associades, i les imprimim també en format percentual.