facemask_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\ProgramData\\Anaconda3\\envs\\cvproject\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "import os\n",
    "import pathlib\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from util.image import unnormalize\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn\n",
    "from torch import functional as F\n",
    "\n",
    "from torchvision import transforms\n",
    "from torchinfo import torchinfo\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import albumentations as A\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "import torchmetrics\n",
    "from torchvision.utils import save_image, make_grid\n",
    "import cv2\n",
    "from util.io import load_ckpt\n",
    "\n",
    "from util.loss import  InpaintingLoss\n",
    "import os, glob\n",
    "\n",
    "import efficientunet\n",
    "import random\n",
    "\n",
    "device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = 'Datasets/Mask detection/archive/Face Mask Dataset/'\n",
    "train_dir = dataset_path+'Train/'\n",
    "val_dir = dataset_path+'Validation/'\n",
    "test_dir = dataset_path+'Test/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sizes = (64, 64)\n",
    "\n",
    "rescale_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Resize(sizes, antialias= False)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 1., 1., 1., 1., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.hstack([torch.ones(5), torch.zeros(4)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_files(path):\n",
    "    mask_files = glob.glob(path + 'WithMask/' +'*.png')\n",
    "    nomask_files = glob.glob(path + 'WithoutMask/' +'*.png')\n",
    "\n",
    "    mask_images = [rescale_transform(Image.open(x)) for x in tqdm(mask_files)]\n",
    "    unmasked_images = [rescale_transform(Image.open(x)) for x in tqdm(nomask_files)]\n",
    "\n",
    "    mask_labels = torch.ones(len(mask_images))\n",
    "    unmask_labels = torch.zeros(len(unmasked_images))\n",
    "\n",
    "    mask_images = torch.stack(mask_images)\n",
    "    unmasked_images = torch.stack(unmasked_images)\n",
    "    images = torch.vstack([mask_images, unmasked_images])\n",
    "    labels = torch.hstack([mask_labels, unmask_labels])\n",
    "\n",
    "    return images, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [00:17<00:00, 280.50it/s]\n",
      "100%|██████████| 5000/5000 [00:14<00:00, 341.31it/s]\n"
     ]
    }
   ],
   "source": [
    "train_images, train_labels = get_files(train_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 400/400 [00:01<00:00, 311.74it/s]\n",
      "100%|██████████| 400/400 [00:01<00:00, 390.79it/s]\n"
     ]
    }
   ],
   "source": [
    "val_images, val_labels = get_files(val_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalisation = transforms.Normalize(0.5, 0.5)\n",
    "\n",
    "augmentation = transforms.RandomAffine(\n",
    "    30, (0.15, 0.15), (0.8, 1.2), 10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskDataset(Dataset):\n",
    "    def __init__(self, images, labels, augmentation = None):\n",
    "        self.images = images\n",
    "        self.augmentation = augmentation\n",
    "        self.labels = labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        img = normalisation(self.images[idx])\n",
    "        if(self.augmentation is not None):\n",
    "            img = self.augmentation(img)\n",
    "        return img, self.labels[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = MaskDataset(train_images, train_labels, augmentation = augmentation)\n",
    "val_dataset = MaskDataset(val_images, val_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 250\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, BATCH_SIZE, shuffle = True)\n",
    "val_dataloader   = DataLoader(val_dataset, BATCH_SIZE, shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(3, 5, 3)        # out = 62\n",
    "        self.conv2 = nn.Conv2d(5, 10, 3)       # in = 31, out = 14\n",
    "        self.conv3 = nn.Conv2d(10, 10, 3)       # in = 12, out = 6\n",
    "\n",
    "        self.pool = nn.MaxPool2d(2)\n",
    "\n",
    "        self.fc1 = nn.Linear(10 * 6 * 6, 64)\n",
    "        self.o_n = nn.Linear(64, 1)\n",
    "\n",
    "\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, inpt):\n",
    "        out = self.activation(self.conv1(inpt))\n",
    "        out = self.pool(out)\n",
    "        \n",
    "        out = self.activation(self.conv2(out))\n",
    "        out = self.pool(out)\n",
    "\n",
    "        out = self.activation(self.conv3(out))\n",
    "        out = self.pool(out)\n",
    "\n",
    "        out = self.flatten(out)\n",
    "\n",
    "        out = self.activation(self.fc1(out))\n",
    "        out = self.o_n(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EarlyStopping:\n",
    "\n",
    "    def __init__(self, patience=5, verbose=False, delta=0, path='checkpoint.pt', trace_func=print):\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.counter = 0\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.val_loss_min = np.Inf\n",
    "        self.delta = delta\n",
    "        self.path = path\n",
    "        self.trace_func = trace_func\n",
    "        \n",
    "    def __call__(self, val_loss, model):\n",
    "\n",
    "        score = -val_loss\n",
    "\n",
    "        if self.best_score is None:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "\n",
    "        elif score < self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            self.trace_func(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "\n",
    "        else:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "            self.counter = 0\n",
    "\n",
    "    def save_checkpoint(self, val_loss, model):\n",
    "        if self.verbose:\n",
    "            self.trace_func(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')\n",
    "        torch.save(model.state_dict(), self.path)\n",
    "        self.val_loss_min = val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "CNN                                      [5, 1]                    --\n",
       "├─Conv2d: 1-1                            [5, 5, 62, 62]            140\n",
       "├─ReLU: 1-2                              [5, 5, 62, 62]            --\n",
       "├─MaxPool2d: 1-3                         [5, 5, 31, 31]            --\n",
       "├─Conv2d: 1-4                            [5, 10, 29, 29]           460\n",
       "├─ReLU: 1-5                              [5, 10, 29, 29]           --\n",
       "├─MaxPool2d: 1-6                         [5, 10, 14, 14]           --\n",
       "├─Conv2d: 1-7                            [5, 10, 12, 12]           910\n",
       "├─ReLU: 1-8                              [5, 10, 12, 12]           --\n",
       "├─MaxPool2d: 1-9                         [5, 10, 6, 6]             --\n",
       "├─Flatten: 1-10                          [5, 360]                  --\n",
       "├─Linear: 1-11                           [5, 64]                   23,104\n",
       "├─ReLU: 1-12                             [5, 64]                   --\n",
       "├─Linear: 1-13                           [5, 1]                    65\n",
       "==========================================================================================\n",
       "Total params: 24,679\n",
       "Trainable params: 24,679\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 5.40\n",
       "==========================================================================================\n",
       "Input size (MB): 0.25\n",
       "Forward/backward pass size (MB): 1.17\n",
       "Params size (MB): 0.10\n",
       "Estimated Total Size (MB): 1.51\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchinfo.summary(CNN(), (5, 3, 64, 64))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.4149, accuracy = 0.7985: 100%|██████████| 40/40 [00:08<00:00,  4.79it/s]\n",
      "Validation loss: 0.2145, accuracy = 0.9225: 100%|██████████| 4/4 [00:00<00:00, 52.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/25], Train Loss: 0.4149, Train Accuracy: 0.7985\n",
      "Epoch [1/25], Val Loss: 0.2145, Val Accuracy: 0.9225\n",
      "Validation loss decreased (inf --> 0.214532).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.2544, accuracy = 0.9038: 100%|██████████| 40/40 [00:08<00:00,  4.96it/s]\n",
      "Validation loss: 0.1746, accuracy = 0.9237: 100%|██████████| 4/4 [00:00<00:00, 50.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [2/25], Train Loss: 0.2544, Train Accuracy: 0.9038\n",
      "Epoch [2/25], Val Loss: 0.1746, Val Accuracy: 0.9237\n",
      "Validation loss decreased (0.214532 --> 0.174557).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.2170, accuracy = 0.9187: 100%|██████████| 40/40 [00:08<00:00,  4.84it/s]\n",
      "Validation loss: 0.1281, accuracy = 0.9463: 100%|██████████| 4/4 [00:00<00:00, 52.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [3/25], Train Loss: 0.2170, Train Accuracy: 0.9187\n",
      "Epoch [3/25], Val Loss: 0.1281, Val Accuracy: 0.9463\n",
      "Validation loss decreased (0.174557 --> 0.128084).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1803, accuracy = 0.9322: 100%|██████████| 40/40 [00:08<00:00,  4.88it/s]\n",
      "Validation loss: 0.1316, accuracy = 0.9650: 100%|██████████| 4/4 [00:00<00:00, 50.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [4/25], Train Loss: 0.1803, Train Accuracy: 0.9322\n",
      "Epoch [4/25], Val Loss: 0.1316, Val Accuracy: 0.9650\n",
      "EarlyStopping counter: 1 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1550, accuracy = 0.9392: 100%|██████████| 40/40 [00:08<00:00,  4.91it/s]\n",
      "Validation loss: 0.1015, accuracy = 0.9725: 100%|██████████| 4/4 [00:00<00:00, 43.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [5/25], Train Loss: 0.1550, Train Accuracy: 0.9392\n",
      "Epoch [5/25], Val Loss: 0.1015, Val Accuracy: 0.9725\n",
      "Validation loss decreased (0.128084 --> 0.101484).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1321, accuracy = 0.9503: 100%|██████████| 40/40 [00:08<00:00,  4.84it/s]\n",
      "Validation loss: 0.0774, accuracy = 0.9750: 100%|██████████| 4/4 [00:00<00:00, 53.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [6/25], Train Loss: 0.1321, Train Accuracy: 0.9503\n",
      "Epoch [6/25], Val Loss: 0.0774, Val Accuracy: 0.9750\n",
      "Validation loss decreased (0.101484 --> 0.077367).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1189, accuracy = 0.9553: 100%|██████████| 40/40 [00:08<00:00,  4.85it/s]\n",
      "Validation loss: 0.0674, accuracy = 0.9800: 100%|██████████| 4/4 [00:00<00:00, 51.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [7/25], Train Loss: 0.1189, Train Accuracy: 0.9553\n",
      "Epoch [7/25], Val Loss: 0.0674, Val Accuracy: 0.9800\n",
      "Validation loss decreased (0.077367 --> 0.067354).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1167, accuracy = 0.9569: 100%|██████████| 40/40 [00:08<00:00,  4.93it/s]\n",
      "Validation loss: 0.0759, accuracy = 0.9800: 100%|██████████| 4/4 [00:00<00:00, 49.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [8/25], Train Loss: 0.1167, Train Accuracy: 0.9569\n",
      "Epoch [8/25], Val Loss: 0.0759, Val Accuracy: 0.9800\n",
      "EarlyStopping counter: 1 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1097, accuracy = 0.9606: 100%|██████████| 40/40 [00:08<00:00,  4.85it/s]\n",
      "Validation loss: 0.0551, accuracy = 0.9837: 100%|██████████| 4/4 [00:00<00:00, 50.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [9/25], Train Loss: 0.1097, Train Accuracy: 0.9606\n",
      "Epoch [9/25], Val Loss: 0.0551, Val Accuracy: 0.9837\n",
      "Validation loss decreased (0.067354 --> 0.055099).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.1042, accuracy = 0.9619: 100%|██████████| 40/40 [00:08<00:00,  4.89it/s]\n",
      "Validation loss: 0.0469, accuracy = 0.9800: 100%|██████████| 4/4 [00:00<00:00, 26.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [10/25], Train Loss: 0.1042, Train Accuracy: 0.9619\n",
      "Epoch [10/25], Val Loss: 0.0469, Val Accuracy: 0.9800\n",
      "Validation loss decreased (0.055099 --> 0.046921).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0927, accuracy = 0.9680: 100%|██████████| 40/40 [00:07<00:00,  5.01it/s]\n",
      "Validation loss: 0.0484, accuracy = 0.9837: 100%|██████████| 4/4 [00:00<00:00, 50.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [11/25], Train Loss: 0.0927, Train Accuracy: 0.9680\n",
      "Epoch [11/25], Val Loss: 0.0484, Val Accuracy: 0.9837\n",
      "EarlyStopping counter: 1 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0953, accuracy = 0.9653: 100%|██████████| 40/40 [00:08<00:00,  4.94it/s]\n",
      "Validation loss: 0.0446, accuracy = 0.9825: 100%|██████████| 4/4 [00:00<00:00, 51.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [12/25], Train Loss: 0.0953, Train Accuracy: 0.9653\n",
      "Epoch [12/25], Val Loss: 0.0446, Val Accuracy: 0.9825\n",
      "Validation loss decreased (0.046921 --> 0.044623).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0869, accuracy = 0.9692: 100%|██████████| 40/40 [00:08<00:00,  4.79it/s]\n",
      "Validation loss: 0.0438, accuracy = 0.9850: 100%|██████████| 4/4 [00:00<00:00, 28.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [13/25], Train Loss: 0.0869, Train Accuracy: 0.9692\n",
      "Epoch [13/25], Val Loss: 0.0438, Val Accuracy: 0.9850\n",
      "Validation loss decreased (0.044623 --> 0.043768).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0880, accuracy = 0.9686: 100%|██████████| 40/40 [00:08<00:00,  4.98it/s]\n",
      "Validation loss: 0.0413, accuracy = 0.9850: 100%|██████████| 4/4 [00:00<00:00, 53.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [14/25], Train Loss: 0.0880, Train Accuracy: 0.9686\n",
      "Epoch [14/25], Val Loss: 0.0413, Val Accuracy: 0.9850\n",
      "Validation loss decreased (0.043768 --> 0.041348).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0805, accuracy = 0.9706: 100%|██████████| 40/40 [00:08<00:00,  4.94it/s]\n",
      "Validation loss: 0.0466, accuracy = 0.9800: 100%|██████████| 4/4 [00:00<00:00, 54.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [15/25], Train Loss: 0.0805, Train Accuracy: 0.9706\n",
      "Epoch [15/25], Val Loss: 0.0466, Val Accuracy: 0.9800\n",
      "EarlyStopping counter: 1 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0920, accuracy = 0.9665: 100%|██████████| 40/40 [00:08<00:00,  4.91it/s]\n",
      "Validation loss: 0.0468, accuracy = 0.9787: 100%|██████████| 4/4 [00:00<00:00, 52.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [16/25], Train Loss: 0.0920, Train Accuracy: 0.9665\n",
      "Epoch [16/25], Val Loss: 0.0468, Val Accuracy: 0.9787\n",
      "EarlyStopping counter: 2 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0801, accuracy = 0.9682: 100%|██████████| 40/40 [00:07<00:00,  5.00it/s]\n",
      "Validation loss: 0.0333, accuracy = 0.9850: 100%|██████████| 4/4 [00:00<00:00, 46.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [17/25], Train Loss: 0.0801, Train Accuracy: 0.9682\n",
      "Epoch [17/25], Val Loss: 0.0333, Val Accuracy: 0.9850\n",
      "Validation loss decreased (0.041348 --> 0.033337).  Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0804, accuracy = 0.9720: 100%|██████████| 40/40 [00:08<00:00,  4.91it/s]\n",
      "Validation loss: 0.0517, accuracy = 0.9850: 100%|██████████| 4/4 [00:00<00:00, 53.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [18/25], Train Loss: 0.0804, Train Accuracy: 0.9720\n",
      "Epoch [18/25], Val Loss: 0.0517, Val Accuracy: 0.9850\n",
      "EarlyStopping counter: 1 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0810, accuracy = 0.9713: 100%|██████████| 40/40 [00:08<00:00,  4.80it/s]\n",
      "Validation loss: 0.0460, accuracy = 0.9825: 100%|██████████| 4/4 [00:00<00:00, 51.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [19/25], Train Loss: 0.0810, Train Accuracy: 0.9713\n",
      "Epoch [19/25], Val Loss: 0.0460, Val Accuracy: 0.9825\n",
      "EarlyStopping counter: 2 out of 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training loss: 0.0755, accuracy = 0.9735: 100%|██████████| 40/40 [00:08<00:00,  4.85it/s]\n",
      "Validation loss: 0.0407, accuracy = 0.9837: 100%|██████████| 4/4 [00:00<00:00, 43.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [20/25], Train Loss: 0.0755, Train Accuracy: 0.9735\n",
      "Epoch [20/25], Val Loss: 0.0407, Val Accuracy: 0.9837\n",
      "EarlyStopping counter: 3 out of 3\n",
      "Early stopping\n",
      "------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 25\n",
    "\n",
    "train_acc = torchmetrics.classification.BinaryAccuracy().to(device)\n",
    "val_acc = torchmetrics.classification.BinaryAccuracy().to(device)\n",
    "\n",
    "model = CNN().to(device)\n",
    "optim = torch.optim.Adam(model.parameters(), lr = 3e-3)\n",
    "criterion = torch.nn.BCELoss()\n",
    "early_stopping = EarlyStopping(patience=3, verbose=True, path = 'mask_model.pth')\n",
    "\n",
    "\n",
    "for epoch_num in range(EPOCHS):\n",
    "    train_loss = 0\n",
    "    i = 0\n",
    "\n",
    "    bar = tqdm(train_dataloader)\n",
    "    for img, label in bar:\n",
    "        i+=1\n",
    "        optim.zero_grad()\n",
    "\n",
    "        img = img.to(device)\n",
    "        label = label.to(device).unsqueeze(1)\n",
    "        predictions = F.sigmoid(model(img))\n",
    "\n",
    "        batch_loss = criterion(predictions, label)\n",
    "\n",
    "        train_acc(predictions, label)\n",
    "        batch_loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "        train_loss+= batch_loss.item()\n",
    "        bar.set_description_str(\"Training loss: {:.4f}, accuracy = {:.4f}\".format(train_loss/i, train_acc.compute()))\n",
    "\n",
    "    train_loss/=i\n",
    "\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        val_loss = 0\n",
    "        i = 0\n",
    "        bar = tqdm(val_dataloader)\n",
    "        for img, label in bar:\n",
    "            i+=1\n",
    "            optim.zero_grad()\n",
    "\n",
    "            img = img.to(device)\n",
    "            label = label.to(device).unsqueeze(1)\n",
    "            predictions = F.sigmoid(model(img))\n",
    "\n",
    "            batch_loss = criterion(predictions, label)\n",
    "\n",
    "            \n",
    "            val_acc(predictions, label)\n",
    "            val_loss+= batch_loss.item()\n",
    "            bar.set_description_str(\"Validation loss: {:.4f}, accuracy = {:.4f}\".format(val_loss/i, val_acc.compute()))\n",
    "\n",
    "        val_loss/=i\n",
    "\n",
    "\n",
    "    print(\"Epoch [{}/{}], Train Loss: {:.4f}, Train Accuracy: {:.4f}\".format(epoch_num+1, EPOCHS, train_loss, train_acc.compute()))\n",
    "    print(\"Epoch [{}/{}], Val Loss: {:.4f}, Val Accuracy: {:.4f}\".format(epoch_num+1, EPOCHS, val_loss, val_acc.compute()))\n",
    "    early_stopping(val_loss, model)\n",
    "\n",
    "    train_acc.reset()\n",
    "    val_acc.reset()\n",
    "\n",
    "    if early_stopping.early_stop:\n",
    "        print(\"Early stopping\")\n",
    "        print('-'*60)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.0 ('cvproject')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8e9b273093b79eabda493c170a15e7ff42f4736405ada6f2b68f061601f943a3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}