rararara9999 commited on
Commit
3aabd17
1 Parent(s): 62d494a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +158 -0
app.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import subprocess
2
+
3
+ # Install the required packages
4
+ subprocess.check_call(["pip", "install", "-U", "git+https://github.com/huggingface/transformers.git"])
5
+ subprocess.check_call(["pip", "install", "-U", "git+https://github.com/huggingface/accelerate.git"])
6
+ subprocess.check_call(["pip", "install", "datasets"])
7
+ subprocess.check_call(["pip", "install", "evaluate"])
8
+ subprocess.check_call(["pip", "install", "scikit-learn"])
9
+ subprocess.check_call(["pip", "install", "torchvision"])
10
+
11
+ model_checkpoint = "Heem2/Facemask-detection"
12
+ batch_size = 128
13
+
14
+ from datasets import load_dataset
15
+ from evaluate import load
16
+
17
+ metric = load("accuracy")
18
+
19
+ # Load the dataset directly from Hugging Face
20
+ dataset = load_dataset("DamarJati/Face-Mask-Detection")
21
+ labels = dataset["train"].features["label"].names
22
+ label2id, id2label = dict(), dict()
23
+ for i, label in enumerate(labels):
24
+ label2id[label] = i
25
+ id2label[i] = label
26
+
27
+ from transformers import AutoImageProcessor
28
+ image_processor = AutoImageProcessor.from_pretrained(model_checkpoint)
29
+ image_processor
30
+
31
+ from torchvision.transforms import (
32
+ CenterCrop,
33
+ Compose,
34
+ Normalize,
35
+ RandomHorizontalFlip,
36
+ RandomResizedCrop,
37
+ Resize,
38
+ ToTensor,
39
+ ColorJitter,
40
+ RandomRotation
41
+ )
42
+
43
+ normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
44
+ size = (image_processor.size["height"], image_processor.size["width"])
45
+
46
+ train_transforms = Compose(
47
+ [
48
+ RandomResizedCrop(size),
49
+ RandomHorizontalFlip(),
50
+ RandomRotation(degrees=15),
51
+ ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
52
+ ToTensor(),
53
+ normalize,
54
+ ]
55
+ )
56
+
57
+ val_transforms = Compose(
58
+ [
59
+ Resize(size),
60
+ CenterCrop(size),
61
+ RandomRotation(degrees=15),
62
+ ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
63
+ ToTensor(),
64
+ normalize,
65
+ ]
66
+ )
67
+
68
+ def preprocess_train(example_batch):
69
+ example_batch["pixel_values"] = [
70
+ train_transforms(image.convert("RGB")) for image in example_batch["image"]
71
+ ]
72
+ return example_batch
73
+
74
+ def preprocess_val(example_batch):
75
+ example_batch["pixel_values"] = [val_transforms(image.convert("RGB")) for image in example_batch["image"]]
76
+ return example_batch
77
+
78
+ splits = dataset["train"].train_test_split(test_size=0.3)
79
+ train_ds = splits['train']
80
+ val_ds = splits['test']
81
+
82
+ train_ds.set_transform(preprocess_train)
83
+ val_ds.set_transform(preprocess_val)
84
+
85
+ from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
86
+
87
+ model = AutoModelForImageClassification.from_pretrained(model_checkpoint,
88
+ label2id=label2id,
89
+ id2label=id2label,
90
+ ignore_mismatched_sizes=True)
91
+
92
+ model_name = model_checkpoint.split("/")[-1]
93
+
94
+ args = TrainingArguments(
95
+ f"{model_name}-finetuned",
96
+ remove_unused_columns=False,
97
+ eval_strategy="epoch", # Updated parameter
98
+ save_strategy="epoch",
99
+ save_total_limit=5,
100
+ learning_rate=1e-3,
101
+ per_device_train_batch_size=batch_size,
102
+ gradient_accumulation_steps=2,
103
+ per_device_eval_batch_size=batch_size,
104
+ num_train_epochs=2,
105
+ warmup_ratio=0.1,
106
+ weight_decay=0.01,
107
+ lr_scheduler_type="cosine",
108
+ logging_steps=10,
109
+ load_best_model_at_end=True,
110
+ metric_for_best_model="accuracy",
111
+ )
112
+
113
+ import numpy as np
114
+
115
+ def compute_metrics(eval_pred):
116
+ """Computes accuracy on a batch of predictions"""
117
+ predictions = np.argmax(eval_pred.predictions, axis=1)
118
+ return metric.compute(predictions=predictions, references=eval_pred.label_ids)
119
+
120
+ import torch
121
+
122
+ def collate_fn(examples):
123
+ pixel_values = torch.stack([example["pixel_values"] for example in examples])
124
+ labels = torch.tensor([example["label"] for example in examples])
125
+ return {"pixel_values": pixel_values, "labels": labels}
126
+
127
+ trainer = Trainer(
128
+ model=model,
129
+ args=args,
130
+ train_dataset=train_ds,
131
+ eval_dataset=val_ds,
132
+ tokenizer=image_processor,
133
+ compute_metrics=compute_metrics,
134
+ data_collator=collate_fn,
135
+ )
136
+
137
+ print("Starting training...")
138
+ train_results = trainer.train()
139
+ print("Training completed.")
140
+
141
+ # Save model
142
+ trainer.save_model()
143
+ trainer.log_metrics("train", train_results.metrics)
144
+ trainer.save_metrics("train", train_results.metrics)
145
+ trainer.save_state()
146
+
147
+ print("Starting evaluation...")
148
+ metrics = trainer.evaluate()
149
+ print("Evaluation completed.")
150
+
151
+ # Log and save metrics
152
+ trainer.log_metrics("eval", metrics)
153
+ trainer.save_metrics("eval", metrics)
154
+
155
+ # Print evaluation metrics
156
+ print("Evaluation Metrics:")
157
+ for key, value in metrics.items():
158
+ print(f"{key}: {value}")