app / app.py
rararara9999's picture
Create app.py
55fe921 verified
raw
history blame
4.67 kB
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"])
# Rest of your code
model_checkpoint = "microsoft/resnet-50"
batch_size = 128
from datasets import load_dataset
# Load the dataset from Hugging Face Hub
dataset = load_dataset("DamarJati/Face-Mask-Detection")
# Continue with the rest of your script...
labels = dataset["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
label2id[label] = i
id2label[i] = label
from transformers import AutoImageProcessor
image_processor = AutoImageProcessor.from_pretrained(model_checkpoint)
from torchvision.transforms import (
CenterCrop, Compose, Normalize, RandomHorizontalFlip, RandomResizedCrop, Resize, ToTensor, ColorJitter, RandomRotation
)
normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
if "height" in image_processor.size:
size = (image_processor.size["height"], image_processor.size["width"])
crop_size = size
max_size = None
elif "shortest_edge" in image_processor.size:
size = image_processor.size["shortest_edge"]
crop_size = (size, size)
max_size = image_processor.size.get("longest_edge")
train_transforms = Compose(
[
RandomResizedCrop(crop_size),
RandomHorizontalFlip(),
RandomRotation(degrees=15),
ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
ToTensor(),
normalize,
]
)
val_transforms = Compose(
[
Resize(size),
CenterCrop(crop_size),
RandomRotation(degrees=15),
ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
ToTensor(),
normalize,
]
)
def preprocess_train(example_batch):
example_batch["pixel_values"] = [
train_transforms(image.convert("RGB")) for image in example_batch["image"]
]
return example_batch
def preprocess_val(example_batch):
example_batch["pixel_values"] = [val_transforms(image.convert("RGB")) for image in example_batch["image"]]
return example_batch
splits = dataset["train"].train_test_split(test_size=0.3)
train_ds = splits['train']
val_ds = splits['test']
train_ds.set_transform(preprocess_train)
val_ds.set_transform(preprocess_val)
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
model = AutoModelForImageClassification.from_pretrained(model_checkpoint,
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes = True)
model_name = model_checkpoint.split("/")[-1]
args = TrainingArguments(
f"{model_name}-finetuned",
remove_unused_columns=False,
evaluation_strategy = "epoch",
save_strategy = "epoch",
save_total_limit = 5,
learning_rate=1e-3,
per_device_train_batch_size=batch_size,
gradient_accumulation_steps=2,
per_device_eval_batch_size=batch_size,
num_train_epochs=2,
warmup_ratio=0.1,
weight_decay=0.01,
lr_scheduler_type="cosine",
logging_steps=10,
load_best_model_at_end=True,
metric_for_best_model="accuracy",)
import numpy as np
def compute_metrics(eval_pred):
"""Computes accuracy on a batch of predictions"""
predictions = np.argmax(eval_pred.predictions, axis=1)
return metric.compute(predictions=predictions, references=eval_pred.label_ids)
import torch
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}
trainer = Trainer(model,
args,
train_dataset=train_ds,
eval_dataset=val_ds,
tokenizer=image_processor,
compute_metrics=compute_metrics,
data_collator=collate_fn,)
train_results = trainer.train()
# Save the model
trainer.save_model()
trainer.log_metrics("train", train_results.metrics)
trainer.save_metrics("train", train_results.metrics)
trainer.save_state()
metrics = trainer.evaluate()
# Log and save evaluation metrics
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)