FiftyOne Finetune DinoV2 for ChestXray Dataset

# !pip install datasets transformers --upgrade

# !pip install fiftyone 

import os
import tqdm
import fiftyone as fo

from datasets import load_dataset

dataset_root = "./datasets"
dataset_splits = ["train", "validation", "test"]

dataset_name = "trpakov/chest-xray-classification"

dataset = load_dataset(dataset_name, name="full")

class_names = dataset["train"].features["labels"].names
labels = class_names

id2label = {id: str(label) for id, label in enumerate(labels)}
label2id = {str(label): id for id, label in enumerate(labels)}

print(f"id2label: {id2label}")

for dataset_split in dataset_splits:
    dataset_save_folder = os.path.join(dataset_root, dataset_split)
    if not os.path.exists(dataset_save_folder):
        os.makedirs(dataset_save_folder)
        
        for image_idx, image_cl_dict in tqdm.tqdm(enumerate(dataset[dataset_split]),
                                                  total=len(dataset[dataset_split]),
                                                  desc=f"Saving {dataset_split} images"):
            image = image_cl_dict["image"]
            cls = id2label[int(image_cl_dict["labels"])]
            
            cls_folder = os.path.join(dataset_save_folder, cls)
            if not os.path.exists(cls_folder):
                os.makedirs(cls_folder)
                
            image.save(f"{cls_folder}/image_{image_idx}.jpg")

print(dataset)

example = dataset["train"][0]

print(f"label: {example['labels']}")
example["image"]

try:
    fo.delete_dataset(dataset_name)
except:
    pass
fiftyone_dataset = fo.Dataset(dataset_name)

dataset_type = fo.types.ImageClassificationDirectoryTree

for dataset_split in dataset_splits:
    dataset_dir = os.path.join(dataset_root, dataset_split)
    fiftyone_dataset.add_dir(
        dataset_dir=dataset_dir,
        dataset_type=dataset_type,
        tags=dataset_split,
    )

session = fo.launch_app(fiftyone_dataset)

import torch

from transformers import AutoImageProcessor, AutoModelForImageClassification
from datasets import load_dataset

model_name = "facebook/dinov2-base"
processor = AutoImageProcessor.from_pretrained(model_name)

print(processor)

model = AutoModelForImageClassification.from_pretrained(
    model_name,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id
)

from torchvision.transforms import Compose, RandomResizedCrop, RandomHorizontalFlip, ColorJitter, ToTensor, Normalize
import torch

# Extarct image mean, std and interpolation from the inference processor
mean = processor.image_mean
std = processor.image_std
interpolation = processor.resample

train_transform = Compose([
    RandomResizedCrop(size=(224, 224), scale=(0.08, 1.0), ratio=(0.75, 1.3333), interpolation=interpolation),
    RandomHorizontalFlip(p=0.5),
    ColorJitter(brightness=(0.6, 1.4), contrast=(0.6, 1.4), saturation=(0.6, 1.4)),
    ToTensor(),
    Normalize(mean=mean, std=std),
])

def prepare_data(batch, mode="train"):
  images = batch["image"]

  if mode == "train":
    images = [train_transform(image.convert("RGB")) for image in images]
    pixel_values = torch.stack(images)
  elif mode == "test":
    pixel_values = processor(images, return_tensors="pt").pixel_values
  else:
    raise ValueError(f"Mode {mode} not supported")

  inputs = {}
  inputs["pixel_values"] = pixel_values
  inputs["labels"] = torch.tensor(batch["labels"])

  return inputs

# set num_proc equal to the number of CPU cores on your machine
train_dataset = dataset["train"].map(prepare_data, num_proc=1, batched=True, batch_size=20, fn_kwargs={"mode":"train"})
eval_dataset = dataset["test"].map(prepare_data, num_proc=1, batched=True, batch_size=20, fn_kwargs={"mode":"test"})

train_dataset.set_format("torch")
eval_dataset.set_format("torch")

import torch

def collate_fn(examples):
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    labels = torch.tensor([example["labels"] for example in examples])
    return {"pixel_values": pixel_values, "labels": labels}

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir=f"{model_name}",
    per_device_train_batch_size=32,
    evaluation_strategy="steps",
    num_train_epochs=4,
    fp16=True,
    save_steps=100,
    eval_steps=100,
    logging_steps=100,
    learning_rate=5e-5,
    save_total_limit=2,
    remove_unused_columns=False,
    push_to_hub=False,
    report_to='tensorboard',
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
)

from sklearn.metrics import accuracy_score
import numpy as np

def compute_metrics(eval_pred):
    """Computes accuracy on a batch of predictions"""
    predictions = np.argmax(eval_pred.predictions, axis=1)
    accuracy = accuracy_score(y_pred=predictions, y_true=eval_pred.label_ids)
    return {"accuracy": accuracy}

from transformers import Trainer

trainer = Trainer(
    model,
    training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=processor,
    compute_metrics=compute_metrics,
    data_collator=collate_fn
)

train_results = trainer.train()

trainer.save_model()
trainer.log_metrics("train", train_results.metrics)
trainer.save_metrics("train", train_results.metrics)
trainer.save_state()

metrics = trainer.evaluate()

trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)

from transformers import AutoModelForImageClassification, AutoImageProcessor

finetuned_model_folder = "./facebook/dinov2-base"

image_processor = AutoImageProcessor.from_pretrained(finetuned_model_folder)
model = AutoModelForImageClassification.from_pretrained(finetuned_model_folder)

from PIL import Image
import torch

def infer(image):
    # prepare image for the model
    encoding = image_processor(image.convert("RGB"), return_tensors="pt")
    
    # forward pass
    with torch.no_grad():
        outputs = model(**encoding)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        predicted_class = id2label[predicted_class_idx]
        return predicted_class

fiftyone_test_dataset = fiftyone_dataset.match_tags("test")
for sample in fiftyone_test_dataset.iter_samples(autosave=True, progress=True):
    image_path = sample.filepath
    image = Image.open(image_path)
    predicted_class = infer(image)
    sample["dinov2_prediction"] = fo.Classification(label=predicted_class)

session.view = fiftyone_test_dataset

results = fiftyone_test_dataset.evaluate_classifications(
    "dinov2_prediction",
    gt_field="ground_truth",
    eval_key="eval",
    method="binary",
    classes=labels,
)

results.print_report()

plot = results.plot_confusion_matrix()
plot.show()

plot = results.plot_pr_curve()
plot.show()

print(fiftyone_test_dataset.count_values("eval"))

from fiftyone import ViewField as F

session.view = (
    fiftyone_test_dataset
    .match(F("dinov2_prediction.label") != F("ground_truth.label"))
)

session.freeze()