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)