pretrainedmodel / app.py
rararara9999's picture
Update app.py
63dd4a4 verified
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)