Spaces:
Sleeping
Sleeping
import subprocess | |
# Install necessary packages | |
subprocess.run(["pip", "install", "-U", "git+https://github.com/huggingface/transformers.git"]) | |
subprocess.run(["pip", "install", "-U", "git+https://github.com/huggingface/accelerate.git"]) | |
subprocess.run(["pip", "install", "datasets"]) | |
subprocess.run(["pip", "install", "evaluate"]) | |
subprocess.run(["pip", "install", "torchvision"]) | |
subprocess.run(["pip", "install", "scikit-learn"]) | |
# Load the necessary libraries | |
from datasets import load_dataset | |
from transformers import AutoModelForImageClassification, AutoImageProcessor, TrainingArguments, Trainer | |
import torch | |
import torchvision.transforms as transforms | |
import numpy as np | |
from evaluate import load | |
# Load the dataset from Hugging Face Hub | |
dataset = load_dataset("DamarJati/Face-Mask-Detection") | |
# Define the labels | |
labels = dataset["train"].features["label"].names | |
label2id, id2label = dict(), dict() | |
for i, label in enumerate(labels): | |
label2id[label] = i | |
id2label[i] = label | |
# Load the pre-trained model and processor | |
model_checkpoint = "microsoft/resnet-50" | |
model = AutoModelForImageClassification.from_pretrained(model_checkpoint, ignore_mismatched_sizes=True, num_labels=len(labels)) | |
image_processor = AutoImageProcessor.from_pretrained(model_checkpoint) | |
# Define the image transformations | |
transform = transforms.Compose([ | |
transforms.Resize((224, 224)), | |
transforms.ToTensor(), | |
transforms.Normalize(mean=image_processor.image_mean, std=image_processor.image_std) | |
]) | |
# Preprocess the dataset | |
def preprocess(example_batch): | |
example_batch["pixel_values"] = [transform(image.convert("RGB")) for image in example_batch["image"]] | |
return example_batch | |
dataset = dataset.with_transform(preprocess) | |
# Split the dataset into training and validation sets | |
splits = dataset["train"].train_test_split(test_size=0.3) | |
train_ds = splits['train'] | |
val_ds = splits['test'] | |
# Define the evaluation metric | |
metric = load("accuracy") | |
def compute_metrics(eval_pred): | |
predictions = np.argmax(eval_pred.predictions, axis=1) | |
return metric.compute(predictions=predictions, references=eval_pred.label_ids) | |
# Define the data collator | |
def collate_fn(examples): | |
pixel_values = torch.stack([example["pixel_values"] for example in examples]) | |
labels = torch.tensor([example["label"] for example in examples]) | |
return {"pixel_values": pixel_values, "labels": labels} | |
# Define the training arguments | |
args = TrainingArguments( | |
output_dir="./results", | |
per_device_eval_batch_size=128, | |
remove_unused_columns=False, | |
) | |
# Initialize the Trainer | |
trainer = Trainer( | |
model=model, | |
args=args, | |
eval_dataset=val_ds, | |
compute_metrics=compute_metrics, | |
data_collator=collate_fn, | |
) | |
# Evaluate the pre-trained model | |
metrics = trainer.evaluate() | |
print(metrics) | |