Add CLI params support + MobileNetV2 model + small refactoring

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.dataFrame = 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.dataFrame.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.dataFrame.index)

mask_detector/mask_classifier.py

@@ -5,85 +5,95 @@
 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 Accuracy
 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):
         super().__init__()
-        self.save_hyperparameters()
         self.net = net
+        self.learning_rate = learning_rate
+        self.accuracy = Accuracy()
 
     def forward(self, x):
         return self.net(x)
 
     def training_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
-
         loss = binary_cross_entropy(out, y)
 
-        self.log('train_loss', loss, on_epoch=True)
-
-        return loss
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     def validation_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
         loss = binary_cross_entropy(out, y)
 
-        self.log('valid_loss', loss, on_step=True)
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     def test_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
         loss = binary_cross_entropy(out, y)
 
-        _, out = torch.max(out, dim=1)
-        val_acc = accuracy_score(out.cpu(), y.cpu())
-        val_acc = torch.tensor(val_acc)
-
-        return {'test_loss': loss, 'test_acc': val_acc}
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     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=1e-3)
+        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 = Trainer.from_argparse_args(args)
     trainer.fit(model, train_loader, val_loader)
 
     # ------------

mask_detector/model.py → mask_detector/models/basic_cnn.py

@@ -1,4 +1,4 @@
-from torch.nn import Module, Conv2d, Linear, MaxPool2d, ReLU, Sequential, Softmax, Flatten
+from torch.nn import Dropout, Module, Conv2d, Flatten, Linear, MaxPool2d, ReLU, Sequential, Sigmoid, Softmax, Flatten
 
 
 class BasicCNN(Module):
@@ -8,35 +8,35 @@ def __init__(self):
         self.convLayers1 = Sequential(
             Conv2d(3, 32, kernel_size=(3, 3), padding=(1, 1)),
             ReLU(),
-            MaxPool2d(kernel_size=(2, 2))
+            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))
+            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))
+            MaxPool2d(kernel_size=(2, 2)),
+            Dropout(p=0.3)
         )
 
         self.linearLayers = Sequential(
-            Linear(in_features=128*25*25, out_features=2048),
-            ReLU(),
-            Linear(in_features=2048, out_features=1024),
+            Flatten(),
+            Linear(in_features=128*15*15, out_features=1024),
             ReLU(),
             Linear(in_features=1024, out_features=1),
-            Softmax()
+            Sigmoid()
         )
 
     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

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, x):
+        with torch.no_grad():
+            x = self.features(x)
+
+        x = self.classifier(x)
+
+        return x