diff --git a/mask_detector/dataset.py b/mask_detector/dataset.py
deleted file mode 100644
index 607cab1..0000000
--- a/mask_detector/dataset.py
+++ /dev/null
@@ -1,31 +0,0 @@
-""" Dataset module
-"""
-import cv2
-import numpy as np
-import pandas as pd
-from torch import float, tensor
-from torch.utils.data.dataset import Dataset
-from torchvision.transforms import Compose, Resize, ToPILImage, ToTensor
-
-
-class MaskDataset(Dataset):
- def __init__(self, csv_file, image_size=100):
- self.dataFrame = pd.read_csv(csv_file)
-
- self.transform = Compose([
- ToPILImage(),
- Resize((image_size, image_size)),
- ToTensor(), # [0, 1] | [no_mask, mask]
- ])
-
- def __getitem__(self, key):
- row = self.dataFrame.iloc[key]
-
- mask = tensor([row['mask']], dtype=float)
- image = cv2.imdecode(np.fromfile(
- row['image'], dtype=np.uint8), cv2.IMREAD_UNCHANGED)
-
- return self.transform(image), mask
-
- def __len__(self):
- return len(self.dataFrame.index)
diff --git a/mask_detector/datasets/__init__.py b/mask_detector/datasets/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/mask_detector/datasets/masked_face_net.py b/mask_detector/datasets/masked_face_net.py
new file mode 100644
index 0000000..bdcd960
--- /dev/null
+++ b/mask_detector/datasets/masked_face_net.py
@@ -0,0 +1,33 @@
+""" Dataset module
+"""
+import cv2
+import torch
+import numpy as np
+import pandas as pd
+
+from PIL import Image
+from torch.utils.data.dataset import Dataset
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToPILImage, ToTensor
+
+
+class MaskedFaceNetDataset(Dataset):
+ def __init__(self, csv_file, image_size):
+ self.data_frame = pd.read_csv(csv_file)
+
+ self.transform = Compose([
+ Resize(image_size), # for MobileNetV2 - set image size to 256
+ CenterCrop(224),
+ ToTensor(),
+ Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+ ])
+
+ def __getitem__(self, key):
+ row = self.data_frame.iloc[key]
+
+ mask = torch.tensor([row['mask']], dtype=torch.float)
+ image = Image.open(row['image'])
+
+ return self.transform(image), mask
+
+ def __len__(self):
+ return len(self.data_frame)
diff --git a/mask_detector/mask_classifier.py b/mask_detector/mask_classifier.py
index ede8224..76c603f 100644
--- a/mask_detector/mask_classifier.py
+++ b/mask_detector/mask_classifier.py
@@ -5,19 +5,21 @@
from torch.utils.data import DataLoader, random_split
from torch.optim import Adam
from pytorch_lightning import LightningModule, Trainer, seed_everything
+from pytorch_lightning.metrics import Recall
from pytorch_lightning.callbacks import ModelCheckpoint
-from sklearn.metrics import accuracy_score
-from model import BasicCNN
-from dataset import MaskDataset
from utils import train_val_test_split
+from models.basic_cnn import BasicCNN
+from models.mobile_net_v2 import MobileNetV2
+from datasets.masked_face_net import MaskedFaceNetDataset
class MaskClassifier(LightningModule):
- def __init__(self, net, learning_rate=1e-3):
+ def __init__(self, net, learning_rate=0.001):
super().__init__()
- self.save_hyperparameters()
self.net = net
+ self.learning_rate = learning_rate
+ self.recall = Recall()
def forward(self, x):
return self.net(x)
@@ -25,10 +27,11 @@ def forward(self, x):
def training_step(self, batch, batch_idx):
x, y = batch
out = self.net(x)
-
loss = binary_cross_entropy(out, y)
+ recall = self.recall(out, y)
- self.log('train_loss', loss, on_epoch=True)
+ self.log('train_loss', loss, on_step=False, on_epoch=True)
+ self.log('train_recall', recall, on_step=False, on_epoch=True)
return loss
@@ -36,55 +39,73 @@ def validation_step(self, batch, batch_idx):
x, y = batch
out = self.net(x)
loss = binary_cross_entropy(out, y)
+ recall = self.recall(out, y)
+
+ self.log('val_loss', loss, on_step=False, on_epoch=True)
+ self.log('val_recall', recall, on_step=False, on_epoch=True)
- self.log('valid_loss', loss, on_step=True)
+ return loss
def test_step(self, batch, batch_idx):
x, y = batch
out = self.net(x)
loss = binary_cross_entropy(out, y)
+ recall = self.recall(out, y)
- _, out = torch.max(out, dim=1)
- val_acc = accuracy_score(out.cpu(), y.cpu())
- val_acc = torch.tensor(val_acc)
+ self.log('test_loss', loss, on_step=False, on_epoch=True)
+ self.log('test_recall', recall, on_step=False, on_epoch=True)
- return {'test_loss': loss, 'test_acc': val_acc}
+ return loss
def configure_optimizers(self):
# self.hparams available because we called self.save_hyperparameters()
- return Adam(self.parameters(), lr=self.hparams.learning_rate)
+ return Adam(self.parameters(), lr=self.learning_rate)
+
+ @staticmethod
+ def add_model_specific_args(parent_parser):
+ parser = ArgumentParser(parents=[parent_parser], add_help=False)
+ parser.add_argument('--learning_rate', type=float, default=0.001)
+ return parser
def cli_main():
seed_everything(1234)
+ # ------------
+ # args
+ # ------------
+ parser = ArgumentParser()
+
+ parser.add_argument('--batch_size', default=32, type=int)
+ parser.add_argument('--image_size', default=120, type=int)
+
+ parser = Trainer.add_argparse_args(parser)
+ parser = MaskClassifier.add_model_specific_args(parser)
+ args = parser.parse_args()
+
# ------------
# data
# ------------
- dataset = MaskDataset(csv_file='data/dataframe/mask_df.csv')
+ dataset = MaskedFaceNetDataset(
+ csv_file='data/dataframe/mask_df.csv', image_size=args.image_size)
ds_train, ds_validate, ds_test = train_val_test_split(
dataset, train_ratio=0.8, validate_ratio=0.1, test_ratio=0.1)
- train_loader = DataLoader(ds_train, batch_size=128)
- val_loader = DataLoader(ds_validate, batch_size=128)
- test_loader = DataLoader(ds_test, batch_size=128)
+ train_loader = DataLoader(ds_train, batch_size=args.batch_size)
+ val_loader = DataLoader(ds_validate, batch_size=args.batch_size)
+ test_loader = DataLoader(ds_test, batch_size=args.batch_size)
# ------------
# model
# ------------
- net = BasicCNN()
- model = MaskClassifier(net, learning_rate=0.0001)
+ net = MobileNetV2()
+ model = MaskClassifier(net, learning_rate=args.learning_rate)
# ------------
# training
# ------------
- checkpoint_callback = ModelCheckpoint(
- verbose=True,
- monitor='test_acc',
- mode='max'
- )
- trainer = Trainer(max_epochs=1, checkpoint_callback=checkpoint_callback)
- trainer.fit(model, train_loader, val_loader)
+ trainer = Trainer.from_argparse_args(args)
+ result = trainer.fit(model, train_loader, val_loader)
# ------------
# testing
diff --git a/mask_detector/model.py b/mask_detector/model.py
deleted file mode 100644
index 01bd4d8..0000000
--- a/mask_detector/model.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from torch.nn import Module, Conv2d, Linear, MaxPool2d, ReLU, Sequential, Softmax, Flatten
-
-
-class BasicCNN(Module):
- def __init__(self):
- super().__init__()
-
- self.convLayers1 = Sequential(
- Conv2d(3, 32, kernel_size=(3, 3), padding=(1, 1)),
- ReLU(),
- MaxPool2d(kernel_size=(2, 2))
- )
-
- self.convLayers2 = Sequential(
- Conv2d(32, 64, kernel_size=(3, 3), padding=(1, 1)),
- ReLU(),
- MaxPool2d(kernel_size=(2, 2))
- )
-
- self.convLayers3 = Sequential(
- Conv2d(64, 128, kernel_size=(3, 3), padding=(1, 1)),
- ReLU(),
- MaxPool2d(kernel_size=(2, 2))
- )
-
- self.linearLayers = Sequential(
- Linear(in_features=128*25*25, out_features=2048),
- ReLU(),
- Linear(in_features=2048, out_features=1024),
- ReLU(),
- Linear(in_features=1024, out_features=1),
- Softmax()
- )
-
- def forward(self, x):
- x = self.convLayers1(x)
- x = self.convLayers2(x)
- x = self.convLayers3(x)
- x = x.view(-1, 128*25*25)
- x = self.linearLayers(x)
-
- return x
diff --git a/mask_detector/models/README.md b/mask_detector/models/README.md
new file mode 100644
index 0000000..d9839a3
--- /dev/null
+++ b/mask_detector/models/README.md
@@ -0,0 +1,7 @@
+# Models
+
+## Current performance measures
+
+| Model | Params | Loss | Recall |
+| ----------- | ----------------------------------------------------- | ------ | ------ |
+| MobileNetV2 | max_epochs=10
batch_size=32
learning_rate=0.001 | 0.1449 | 0.9794 |
diff --git a/mask_detector/models/__init__.py b/mask_detector/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/mask_detector/models/basic_cnn.py b/mask_detector/models/basic_cnn.py
new file mode 100644
index 0000000..e2eff73
--- /dev/null
+++ b/mask_detector/models/basic_cnn.py
@@ -0,0 +1,43 @@
+from torch.nn import Dropout, Module, Conv2d, Flatten, Linear, MaxPool2d, ReLU, Sequential, Sigmoid, Softmax, Flatten
+
+
+class BasicCNN(Module):
+ def __init__(self):
+ super().__init__()
+
+ self.convLayers1 = Sequential(
+ Conv2d(3, 32, kernel_size=(3, 3), padding=(1, 1)),
+ ReLU(),
+ MaxPool2d(kernel_size=(2, 2)),
+ Dropout(p=0.3)
+ )
+
+ self.convLayers2 = Sequential(
+ Conv2d(32, 64, kernel_size=(3, 3), padding=(1, 1)),
+ ReLU(),
+ MaxPool2d(kernel_size=(2, 2)),
+ Dropout(p=0.3)
+ )
+
+ self.convLayers3 = Sequential(
+ Conv2d(64, 128, kernel_size=(3, 3), padding=(1, 1)),
+ ReLU(),
+ MaxPool2d(kernel_size=(2, 2)),
+ Dropout(p=0.3)
+ )
+
+ self.linearLayers = Sequential(
+ Flatten(),
+ Linear(in_features=128*15*15, out_features=1024),
+ ReLU(),
+ Linear(in_features=1024, out_features=1),
+ Sigmoid()
+ )
+
+ def forward(self, output):
+ output = self.convLayers1(output)
+ output = self.convLayers2(output)
+ output = self.convLayers3(output)
+ output = self.linearLayers(output)
+
+ return output
diff --git a/mask_detector/models/mobile_net_v2.py b/mask_detector/models/mobile_net_v2.py
new file mode 100644
index 0000000..3a25b98
--- /dev/null
+++ b/mask_detector/models/mobile_net_v2.py
@@ -0,0 +1,30 @@
+import torch
+from torch.nn import AvgPool2d, Dropout, Flatten, Module, Conv2d, Linear, MaxPool2d, ReLU, Sequential, Sigmoid
+
+
+class MobileNetV2(Module):
+ def __init__(self):
+ super().__init__()
+
+ net = torch.hub.load('pytorch/vision:v0.6.0',
+ 'mobilenet_v2', pretrained=True)
+
+ self.features = net.features
+
+ self.classifier = Sequential(
+ AvgPool2d(kernel_size=(7, 7)),
+ Flatten(),
+ Linear(in_features=1280, out_features=128),
+ Dropout(p=0.5),
+ ReLU(),
+ Linear(in_features=128, out_features=1),
+ Sigmoid()
+ )
+
+ def forward(self, output):
+ with torch.no_grad():
+ output = self.features(output)
+
+ output = self.classifier(output)
+
+ return output