Perfeziona un modello di segmentazione semantica con un dataset personalizzato

Ottimizza la segmentazione semantica con un dataset personalizzato.

Questa guida mostra come è possibile ottimizzare Segformer, un modello di segmentazione semantica all’avanguardia. Il nostro obiettivo è costruire un modello per un robot di consegna delle pizze, in modo che possa vedere dove guidare e riconoscere gli ostacoli 🍕🤖. Prima etichetteremo un insieme di immagini di marciapiede su Segments.ai. Poi ottimizzeremo un modello pre-addestrato SegFormer utilizzando 🤗 transformers, una libreria open-source che offre implementazioni facili da usare di modelli all’avanguardia. Lungo il percorso, imparerai come lavorare con Hugging Face Hub, il catalogo open-source più grande di modelli e set di dati.

La segmentazione semantica è il compito di classificare ogni pixel in un’immagine. Puoi vederlo come un modo più preciso di classificare un’immagine. Ha una vasta gamma di casi d’uso in campi come l’imaging medico e la guida autonoma. Ad esempio, per il nostro robot di consegna delle pizze, è importante sapere esattamente dove si trova il marciapiede in un’immagine, non solo se c’è o meno un marciapiede.

Poiché la segmentazione semantica è un tipo di classificazione, le architetture di rete utilizzate per la classificazione delle immagini e la segmentazione semantica sono molto simili. Nel 2014, un documento fondamentale di Long et al. ha utilizzato reti neurali convoluzionali per la segmentazione semantica. Più recentemente, i Transformers sono stati utilizzati per la classificazione delle immagini (ad esempio, ViT), e ora vengono utilizzati anche per la segmentazione semantica, spingendo ulteriormente gli sviluppi all’avanguardia.

• Annuncio del Programma di Residenza di Ricerca sull’Intelligenza Artificiale 🤗
• Esperti di Machine Learning – Margaret Mitchell
• Presentazione dei Decision Transformers su Hugging Face 🤗

SegFormer è un modello per la segmentazione semantica introdotto da Xie et al. nel 2021. Ha un codificatore Transformer gerarchico che non utilizza codifiche posizionali (a differenza di ViT) e un semplice decodificatore a perceptron multistrato. SegFormer raggiunge prestazioni all’avanguardia su più set di dati comuni. Vediamo come si comporta il nostro robot di consegna delle pizze per le immagini di marciapiede.

Iniziamo installando le dipendenze necessarie. Poiché andremo a caricare il nostro set di dati e il modello su Hugging Face Hub, è necessario installare Git LFS e accedere a Hugging Face.

L’installazione di git-lfs potrebbe essere diversa nel tuo sistema. Si noti che Google Colab ha Git LFS pre-installato.

pip install -q transformers datasets evaluate segments-ai
apt-get install git-lfs
git lfs install
huggingface-cli login

Il primo passo in qualsiasi progetto di ML è assemblare un buon set di dati. Per addestrare un modello di segmentazione semantica, abbiamo bisogno di un set di dati con etichette di segmentazione semantica. Possiamo utilizzare un set di dati esistente da Hugging Face Hub, come ADE20k, o creare il nostro set di dati.

Per il nostro robot di consegna delle pizze, potremmo utilizzare un set di dati di guida autonoma esistente come CityScapes o BDD100K. Tuttavia, questi set di dati sono stati acquisiti da auto che guidano sulla strada. Poiché il nostro robot di consegna guiderà sul marciapiede, ci sarà una discrepanza tra le immagini in questi set di dati e i dati che il nostro robot vedrà nel mondo reale.

Non vogliamo che il nostro robot di consegna si confonda, quindi creeremo il nostro set di dati di segmentazione semantica utilizzando immagini acquisite sui marciapiedi. Ti mostreremo come etichettare le immagini che abbiamo acquisito nei prossimi passaggi. Se vuoi solo utilizzare il nostro set di dati finito e etichettato, puoi saltare la sezione “Crea il tuo set di dati” e continuare da “Utilizza un set di dati dal Hub”.

Crea il tuo set di dati

Per creare il tuo set di dati di segmentazione semantica, avrai bisogno di due cose:

1. immagini che coprano le situazioni che il tuo modello incontrerà nel mondo reale
2. etichette di segmentazione, ovvero immagini in cui ogni pixel rappresenta una classe/categoria.

Siamo andati avanti e abbiamo acquisito mille immagini di marciapiede in Belgio. Raccogliere e etichettare un tale set di dati può richiedere molto tempo, quindi puoi iniziare con un set di dati più piccolo e espanderlo se il modello non funziona abbastanza bene.

Alcuni esempi delle immagini grezze nel set di dati del marciapiede.

Per ottenere le etichette di segmentazione, dobbiamo indicare le classi di tutte le regioni/oggetti in queste immagini. Questo può essere un impegno che richiede tempo, ma utilizzando gli strumenti giusti è possibile accelerare significativamente il compito. Per l’etichettatura, utilizzeremo Segments.ai, poiché dispone di strumenti intelligenti per l’etichettatura della segmentazione delle immagini e di un SDK Python facile da usare.

Configura il compito di etichettatura su Segments.ai

Prima, crea un account su https://segments.ai/join. Successivamente, crea un nuovo set di dati e carica le tue immagini. Puoi farlo tramite l’interfaccia web o tramite l’SDK Python (vedi il notebook).

Etichettare le immagini

Ora che i dati grezzi sono stati caricati, vai su segments.ai/home e apri il dataset appena creato. Clicca su “Inizia l’etichettatura” e crea delle maschere di segmentazione. Puoi utilizzare gli strumenti di superpixel e autosegmentazione potenziati da ML per etichettare più velocemente.

Suggerimento: quando utilizzi lo strumento di superpixel, scorri per cambiare la dimensione dei superpixel e clicca e trascina per selezionare i segmenti.

Carica il risultato su Hugging Face Hub

Una volta completata l’etichettatura, crea un nuovo rilascio del dataset contenente i dati etichettati. Puoi farlo nella scheda “Rilasci” su Segments.ai, o programmattivamente tramite l’SDK come mostrato nel notebook.

Nota che la creazione del rilascio potrebbe richiedere alcuni secondi. Puoi controllare la scheda “Rilasci” su Segments.ai per verificare se il tuo rilascio è ancora in fase di creazione.

Ora convertiremo il rilascio in un dataset di Hugging Face tramite l’SDK di Segments.ai Python. Se non hai ancora configurato il client Python di Segments, segui le istruzioni nella sezione “Configura il compito di etichettatura su Segments.ai” del notebook.

Nota che la conversione potrebbe richiedere del tempo, a seconda delle dimensioni del tuo dataset.

from segments.huggingface import release2dataset

release = segments_client.get_release(dataset_identifier, release_name)
hf_dataset = release2dataset(release)

Se ispezioniamo le caratteristiche del nuovo dataset, possiamo vedere la colonna delle immagini e l’etichetta corrispondente. L’etichetta è composta da due parti: una lista di annotazioni e una mappa di segmentazione. L’annotazione corrisponde agli oggetti diversi presenti nell’immagine. Per ogni oggetto, l’annotazione contiene un id e una category_id. La mappa di segmentazione è un’immagine in cui ogni pixel contiene l’id dell’oggetto in quel pixel. Ulteriori informazioni possono essere trovate nella documentazione pertinente.

Per la segmentazione semantica, abbiamo bisogno di una mappa semantica che contenga una category_id per ogni pixel. Utilizzeremo la funzione get_semantic_bitmap dell’SDK di Segments.ai per convertire le mappe in mappe semantiche. Per applicare questa funzione a tutte le righe del nostro dataset, utilizzeremo dataset.map.

from segments.utils import get_semantic_bitmap

def convert_segmentation_bitmap(example):
    return {
        "label.segmentation_bitmap":
            get_semantic_bitmap(
                example["label.segmentation_bitmap"],
                example["label.annotations"],
                id_increment=0,
            )
    }


semantic_dataset = hf_dataset.map(
    convert_segmentation_bitmap,
)

Puoi anche riscrivere la funzione convert_segmentation_bitmap per utilizzare i batch e passare batched=True a dataset.map. Questo velocizzerà significativamente la mappatura, ma potresti dover regolare la dimensione del batch (batch_size) per evitare problemi di memoria.

Il modello SegFormer che andremo a perfezionare in seguito richiede nomi specifici per le caratteristiche. Per comodità, adatteremo il formato ora. Pertanto, rinomineremo la caratteristica image in pixel_values e label.segmentation_bitmap in label e elimineremo le altre caratteristiche.

semantic_dataset = semantic_dataset.rename_column('image', 'pixel_values')
semantic_dataset = semantic_dataset.rename_column('label.segmentation_bitmap', 'label')
semantic_dataset = semantic_dataset.remove_columns(['name', 'uuid', 'status', 'label.annotations'])

Ora possiamo caricare il dataset trasformato su Hugging Face Hub. In questo modo, il tuo team e la community di Hugging Face potranno utilizzarlo. Nella prossima sezione vedremo come caricare il dataset da Hub.

hf_dataset_identifier = f"{hf_username}/{dataset_name}"

semantic_dataset.push_to_hub(hf_dataset_identifier)

Utilizzare un dataset da Hub

Se non desideri creare il tuo dataset, ma hai trovato un dataset adatto al tuo caso d’uso su Hugging Face Hub, puoi definire l’identificatore qui.

Ad esempio, puoi utilizzare il dataset completo delle marciapiedi etichettate. Nota che puoi visualizzare gli esempi direttamente nel tuo browser.

hf_dataset_identifier = "segments/sidewalk-semantic"

Ora che abbiamo creato un nuovo set di dati e lo abbiamo caricato su Hugging Face Hub, possiamo caricare il set di dati in una singola riga.

from datasets import load_dataset

ds = load_dataset(hf_dataset_identifier)

Mescoliamo il set di dati e suddividiamolo in un set di addestramento e un set di test.

ds = ds.shuffle(seed=1)
ds = ds["train"].train_test_split(test_size=0.2)
train_ds = ds["train"]
test_ds = ds["test"]

Estrarremo il numero di etichette e gli ID leggibili dall’utente, in modo da poter configurare correttamente il modello di segmentazione in seguito.

import json
from huggingface_hub import hf_hub_download

repo_id = f"datasets/{hf_dataset_identifier}"
filename = "id2label.json"
id2label = json.load(open(hf_hub_download(repo_id=hf_dataset_identifier, filename=filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
label2id = {v: k for k, v in id2label.items()}

num_labels = len(id2label)

Estrattore di caratteristiche e aumento dei dati

Un modello SegFormer si aspetta che l’input abbia una certa forma. Per trasformare i nostri dati di addestramento in modo che corrispondano alla forma prevista, possiamo utilizzare SegFormerFeatureExtractor. Potremmo utilizzare la funzione ds.map per applicare l’estrattore di caratteristiche all’intero set di dati di addestramento in anticipo, ma questo può occupare molto spazio su disco. Invece, utilizzeremo una trasformazione, che preparerà solo un batch di dati quando quei dati vengono effettivamente utilizzati (on-the-fly). In questo modo, possiamo iniziare l’addestramento senza attendere ulteriori preprocessamenti dei dati.

Nella nostra trasformazione, definiremo anche alcune tecniche di aumento dei dati per rendere il nostro modello più resiliente a diverse condizioni di illuminazione. Utilizzeremo la funzione ColorJitter da torchvision per cambiare casualmente la luminosità, il contrasto, la saturazione e la tonalità delle immagini nel batch.

from torchvision.transforms import ColorJitter
from transformers import SegformerFeatureExtractor

feature_extractor = SegformerFeatureExtractor()
jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1) 

def train_transforms(example_batch):
    images = [jitter(x) for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = feature_extractor(images, labels)
    return inputs


def val_transforms(example_batch):
    images = [x for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = feature_extractor(images, labels)
    return inputs


# Imposta le trasformazioni
train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)

Carica il modello per il fine-tuning

Gli autori di SegFormer definiscono 5 modelli di dimensioni crescenti: B0 a B5. Il seguente grafico (tratto dall’articolo originale) mostra le prestazioni di questi diversi modelli sul dataset ADE20K, confrontate con altri modelli.

Fonte

Qui caricheremo il modello SegFormer più piccolo (B0), preaddestrato su ImageNet-1k. Ha solo circa 14 MB di dimensione! Utilizzare un modello piccolo garantirà che il nostro modello possa funzionare senza problemi sul nostro robot per la consegna di pizza.

from transformers import SegformerForSemanticSegmentation

pretrained_model_name = "nvidia/mit-b0" 
model = SegformerForSemanticSegmentation.from_pretrained(
    pretrained_model_name,
    id2label=id2label,
    label2id=label2id
)

Configurazione del Trainer

Per il fine-tuning del modello sui nostri dati, utilizzeremo l’API del Trainer di Hugging Face. Dobbiamo configurare la configurazione di addestramento e una metrica di valutazione da utilizzare con il Trainer.

Innanzitutto, configureremo TrainingArguments. Questo definisce tutti gli iperparametri di addestramento, come il tasso di apprendimento e il numero di epoche, la frequenza di salvataggio del modello e così via. Specificaremo anche di caricare il modello su hub dopo l’addestramento ( push_to_hub=True ) e specificare un nome del modello ( hub_model_id ).

from transformers import TrainingArguments

epochs = 50
lr = 0.00006
batch_size = 2

hub_model_id = "segformer-b0-finetuned-segments-sidewalk-2"

training_args = TrainingArguments(
    "segformer-b0-finetuned-segments-sidewalk-outputs",
    learning_rate=lr,
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    save_total_limit=3,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=20,
    eval_steps=20,
    logging_steps=1,
    eval_accumulation_steps=5,
    load_best_model_at_end=True,
    push_to_hub=True,
    hub_model_id=hub_model_id,
    hub_strategy="end",
)

Successivamente, definiremo una funzione che calcola la metrica di valutazione con cui vogliamo lavorare. Poiché stiamo facendo segmentazione semantica, utilizzeremo il Intersection over Union medio (mIoU), accessibile direttamente nella libreria evaluate. IoU rappresenta l’overlap delle maschere di segmentazione. L’IoU medio è la media dell’IoU di tutte le classi semantiche. Dai un’occhiata a questo articolo per una panoramica delle metriche di valutazione per la segmentazione delle immagini.

Poiché il nostro modello restituisce i logit con dimensioni altezza/4 e larghezza/4, dobbiamo ingrandirli prima di poter calcolare il mIoU.

import torch
from torch import nn
import evaluate

metric = evaluate.load("mean_iou")

def compute_metrics(eval_pred):
  with torch.no_grad():
    logits, labels = eval_pred
    logits_tensor = torch.from_numpy(logits)
    # scala i logit alle dimensioni dell'etichetta
    logits_tensor = nn.functional.interpolate(
        logits_tensor,
        size=labels.shape[-2:],
        mode="bilinear",
        align_corners=False,
    ).argmax(dim=1)

    pred_labels = logits_tensor.detach().cpu().numpy()
    # attualmente utilizzando _compute invece di compute
    # vedi questo problema per ulteriori informazioni: https://github.com/huggingface/evaluate/pull/328#issuecomment-1286866576
    metrics = metric._compute(
            predictions=pred_labels,
            references=labels,
            num_labels=len(id2label),
            ignore_index=0,
            reduce_labels=feature_extractor.do_reduce_labels,
        )
    
    # aggiungi metriche per categoria come coppie chiave-valore individuali
    per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
    per_category_iou = metrics.pop("per_category_iou").tolist()

    metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
    metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})
    
    return metrics

Infine, possiamo istanziare un oggetto Trainer.

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    compute_metrics=compute_metrics,
)

Ora che il nostro trainer è configurato, l’addestramento è semplice come chiamare la funzione train. Non dobbiamo preoccuparci di gestire la nostra/e GPU, il trainer si occuperà di questo.

trainer.train()

Quando abbiamo finito con l’addestramento, possiamo caricare il nostro modello fine-tuned e il feature extractor sul Hub.

Questo creerà automaticamente anche una scheda del modello con i nostri risultati. Forniremo alcune informazioni extra in kwargs per rendere la scheda del modello più completa.

kwargs = {
    "tags": ["visione", "segmentazione-immagini"],
    "fine-tuned_from": pretrained_model_name,
    "dataset": hf_dataset_identifier,
}

feature_extractor.push_to_hub(hub_model_id)
trainer.push_to_hub(**kwargs)

Ora viene la parte eccitante, utilizzare il nostro modello fine-tuned! In questa sezione, mostreremo come è possibile caricare il proprio modello dal Hub e usarlo per l’inferring.

Tuttavia, è anche possibile provare il proprio modello direttamente sul Hugging Face Hub, grazie ai widget interessanti forniti dall’API di inference ospitata. Se hai caricato il tuo modello sul Hub nel passaggio precedente, dovresti vedere un widget di inference sulla pagina del tuo modello. Puoi aggiungere esempi predefiniti al widget definendo gli URL delle immagini di esempio nella scheda del tuo modello. Vedi questa scheda del modello come esempio.

Usare il modello dal Hub

Prima caricheremo il modello dal Hub utilizzando SegformerForSemanticSegmentation.from_pretrained().

from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation

feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained(f"{hf_username}/{hub_model_id}")

Successivamente, caricheremo un’immagine dal nostro dataset di test.

image = test_ds[0]['pixel_values']
gt_seg = test_ds[0]['label']
image

Per segmentare questa immagine di test, prima dobbiamo preparare l’immagine utilizzando l’estrattore di caratteristiche. Poi lo inoltriamo attraverso il modello.

Dobbiamo anche ricordarci di ridimensionare i logit di output alle dimensioni dell’immagine originale. Per ottenere le previsioni effettive delle categorie, dobbiamo semplicemente applicare un argmax sui logit.

from torch import nn

inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits  # forma (batch_size, num_labels, altezza/4, larghezza/4)

# Prima, ridimensiona i logit alle dimensioni dell'immagine originale
upsampled_logits = nn.functional.interpolate(
    logits,
    size=image.size[::-1], # (altezza, larghezza)
    mode='bilinear',
    align_corners=False
)

# Secondo, applica argmax sulla dimensione delle classi
pred_seg = upsampled_logits.argmax(dim=1)[0]

Ora è il momento di visualizzare il risultato. Visualizzeremo il risultato accanto alla maschera ground truth.

Cosa ne pensi? Invieresti il nostro robot di consegna delle pizze sulla strada con queste informazioni di segmentazione?

Il risultato potrebbe non essere ancora perfetto, ma possiamo sempre espandere il nostro dataset per rendere il modello più robusto. Possiamo ora anche addestrare un modello SegFormer più grande e vedere come si comporta.

Ecco fatto! Ora sai come creare il tuo dataset di segmentazione delle immagini e come usarlo per affinare un modello di segmentazione semantica.

Ti abbiamo presentato alcuni strumenti utili lungo il percorso, come:

• Segments.ai per etichettare i tuoi dati
• 🤗 datasets per creare e condividere un dataset
• 🤗 transformers per affinare facilmente un modello di segmentazione all’avanguardia
• Hugging Face Hub per condividere il nostro dataset e il modello, e per creare un widget di inferenza per il nostro modello

Speriamo che tu abbia apprezzato questo post e imparato qualcosa. Sentiti libero di condividere il tuo modello con noi su Twitter (@TobiasCornille, @NielsRogge e @huggingface).