Upload 7 files

config.py

@@ -0,0 +1,23 @@
+import torch
+from albumentations.pytorch import ToTensorV2
+import albumentations as A
+
+# DEVICE = "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TRAIN_DIR = "images/train"
+VAL_DIR = "images/validation"
+LEARNING_RATE = 2e-3
+BATCH_SIZE = 128
+NUM_WORKERS = 2
+NUM_EPOCHS = 100
+LOAD_MODEL = True
+SAVE_MODEL = True
+CHECKPOINT = "face_emotion.pth.tar"
+
+transform = A.Compose(
+    [
+        A.Resize(height=48, width=48),
+        A.Normalize(mean=[0.5], std=[0.5], max_pixel_value=255.0, ),
+        ToTensorV2(),
+    ]
+)

dataset.py

@@ -0,0 +1,39 @@
+import os
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+import config
+
+
+class EmotionDataset(Dataset):
+    def __init__(self, root_dir, transform=config.transform):
+        self.root_dir = root_dir
+        self.transform = transform
+        self.classes = sorted(os.listdir(root_dir))
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+        self.file_list = self.create_file_list()
+
+    def create_file_list(self):
+        file_list = []
+        for cls in self.classes:
+            class_path = os.path.join(self.root_dir, cls)
+            for file_name in os.listdir(class_path):
+                file_path = os.path.join(class_path, file_name)
+                file_list.append((file_path, self.class_to_idx[cls]))
+        return file_list
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, idx):
+        img_path, label = self.file_list[idx]
+        image = np.array(Image.open(img_path).convert('L'))
+        target = [0]*7
+        target[label] = 1
+        target = torch.FloatTensor(target)
+
+        if self.transform:
+            image = config.transform(image=image)["image"]
+
+        return image, target

model.py

@@ -0,0 +1,151 @@
+# from torch import nn as nn
+#
+#
+# class EmotionModel(nn.Module):
+#     def __init__(self, in_channels=1, num_classes=7):
+#         super(EmotionModel, self).__init__()
+#         self.conv1 = nn.Conv2d(
+#             in_channels=in_channels, out_channels=256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)
+#         )
+#         self.relu1 = nn.ReLU()
+#         self.pool1 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.drop1 = nn.Dropout2d(0.4)
+#
+#         self.conv2 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.relu2 = nn.ReLU()
+#         self.pool2 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.drop2 = nn.Dropout2d(0.4)
+#
+#         self.conv3 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.relu3 = nn.ReLU()
+#         self.pool3 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.drop3 = nn.Dropout2d(0.4)
+#
+#         self.conv4 = nn.Conv2d(
+#             in_channels=512, out_channels=512 * 4 * 4, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)
+#         )
+#         self.relu4 = nn.ReLU()
+#         self.pool4 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+#         self.drop4 = nn.Dropout2d(0.4)
+#
+#         self.fc1 = nn.Linear(in_features=512 * 4 * 4, out_features=512)
+#         self.relu5 = nn.ReLU()
+#         self.drop5 = nn.Dropout(0.3)
+#         self.fc2 = nn.Linear(in_features=512, out_features=256)
+#         self.relu6 = nn.ReLU()
+#         self.drop6 = nn.Dropout(0.3)
+#         self.fc3 = nn.Linear(in_features=256, out_features=num_classes)
+#         self.softmax = nn.Softmax(dim=1)
+#
+#     def forward(self, x):
+#         x = self.conv1(x)
+#         x = self.relu1(x)
+#         x = self.pool1(x)
+#         x = self.drop1(x)
+#         x = self.conv2(x)
+#         x = self.relu2(x)
+#         x = self.pool2(x)
+#         x = self.drop2(x)
+#         x = self.conv3(x)
+#         x = self.relu3(x)
+#         x = self.pool3(x)
+#         x = self.drop3(x)
+#         x = self.conv4(x)
+#         x = self.relu4(x)
+#         x = self.pool4(x)
+#         x = self.drop4(x)
+#         x = x.view(-1, 512 * 4 * 4)
+#         x = self.fc1(x)
+#         x = self.relu5(x)
+#         x = self.drop5(x)
+#         x = self.fc2(x)
+#         x = self.relu6(x)
+#         x = self.drop6(x)
+#         x = self.fc3(x)
+#         x = self.softmax(x)
+#         return x
+
+import torch.nn as nn
+
+
+class EmotionModel(nn.Module):
+    def __init__(self, in_channels=1, num_classes=7):
+        super(EmotionModel, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 256, kernel_size=3, stride=2, padding=1)
+        self.bn1 = nn.BatchNorm2d(256)
+        self.relu1 = nn.ReLU()
+        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.drop1 = nn.Dropout2d(0.4)
+
+        self.conv2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
+        self.bn2 = nn.BatchNorm2d(512)
+        self.relu2 = nn.ReLU()
+        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.drop2 = nn.Dropout2d(0.4)
+
+        self.conv3 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+        self.bn3 = nn.BatchNorm2d(512)
+        self.relu3 = nn.ReLU()
+        self.pool3 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.drop3 = nn.Dropout2d(0.4)
+
+        self.conv4 = nn.Conv2d(512, 512 * 4 * 4, kernel_size=3, stride=2, padding=1)
+        self.bn4 = nn.BatchNorm2d(512 * 4 * 4)
+        self.relu4 = nn.ReLU()
+        self.pool4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.drop4 = nn.Dropout2d(0.4)
+
+        self.fc1 = nn.Linear(512 * 4 * 4, 512)
+        self.bn_fc1 = nn.BatchNorm1d(512)
+        self.relu5 = nn.ReLU()
+        self.drop5 = nn.Dropout(0.3)
+
+        self.fc2 = nn.Linear(512, 256)
+        self.bn_fc2 = nn.BatchNorm1d(256)
+        self.relu6 = nn.ReLU()
+        self.drop6 = nn.Dropout(0.3)
+
+        self.fc3 = nn.Linear(256, num_classes)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu1(x)
+        x = self.pool1(x)
+        x = self.drop1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu2(x)
+        x = self.pool2(x)
+        x = self.drop2(x)
+
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu3(x)
+        x = self.pool3(x)
+        x = self.drop3(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu4(x)
+        x = self.pool4(x)
+        x = self.drop4(x)
+
+        x = x.view(-1, 512 * 4 * 4)
+
+        x = self.fc1(x)
+        x = self.bn_fc1(x)
+        x = self.relu5(x)
+        x = self.drop5(x)
+
+        x = self.fc2(x)
+        x = self.bn_fc2(x)
+        x = self.relu6(x)
+        x = self.drop6(x)
+
+        x = self.fc3(x)
+        x = self.softmax(x)
+
+        return x

requirements.txt

@@ -0,0 +1,7 @@
+torch
+streamlit
+tqdm
+albumentations
+numpy
+Pillow
+matplotlib

train.py

@@ -0,0 +1,144 @@
+import time
+import numpy as np
+import torch
+from PIL import Image
+from matplotlib import pyplot as plt
+from torch import nn as nn, optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import config
+from dataset import EmotionDataset
+from model import EmotionModel
+from utils import load_checkpoint, save_checkpoint
+
+
+def train_fn(model, loader, opt, criterion, epoch):
+    loop = tqdm(loader, leave=True)
+    model.train()
+    epoch_loss = 0.0
+
+    for idx, (image, label) in enumerate(loop):
+        total_acc, total_count = 0, 0
+        image = image.to(config.DEVICE)
+        label = label.to(config.DEVICE)
+        opt.zero_grad()
+        predicted_label = model(image)
+        loss = criterion(predicted_label, label)
+        epoch_loss += loss.item()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
+        opt.step()
+        total_acc += (predicted_label.argmax(1) == label.argmax(1)).sum().item()
+        total_count += label.size(0)
+        loop.set_postfix({"epoch": epoch, "loss": epoch_loss / len(loader), "accuracy": total_acc / total_count})
+
+
+def main():
+    model = EmotionModel().to(config.DEVICE)
+    opt = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, betas=(0.5, 0.999), )
+    criterion = nn.CrossEntropyLoss()
+
+    # if config.LOAD_MODEL:
+    #     load_checkpoint(
+    #         config.CHECKPOINT, model, opt, config.LEARNING_RATE,
+    #     )
+
+    train_dataset = EmotionDataset(root_dir=config.TRAIN_DIR)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=config.BATCH_SIZE,
+        shuffle=True,
+        num_workers=config.NUM_WORKERS,
+    )
+    val_dataset = EmotionDataset(root_dir=config.VAL_DIR)
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=config.BATCH_SIZE,
+        shuffle=True,
+        num_workers=config.NUM_WORKERS,
+    )
+
+    total_accu = None
+    # scheduler = torch.optim.lr_scheduler.StepLR(opt, 1, gamma=0.5)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, patience=2, verbose=True)
+    for epoch in range(config.NUM_EPOCHS):
+        epoch_start_time = time.time()
+        train_fn(
+            model, train_loader, opt, criterion, epoch
+        )
+        accu_val, loss_val = evaluate(model, criterion, val_loader)
+        # if total_accu is not None and total_accu > accu_val:
+        #     scheduler.step()
+        # else:
+        #     total_accu = accu_val
+        scheduler.step(loss_val)
+        print("+" + "-" * 19 + "+" + "-" * 15 + "+" + "-" * 20 + "+" + "-" * 24 + "+")
+        print(
+            "| end of epoch: {:3d} | time: {:6.2f}s | val_loss: {:8.3f} | "
+            "val_accuracy: {:8.3f} |".format(
+                epoch, time.time() - epoch_start_time, loss_val, accu_val
+            )
+        )
+        print("+" + "-" * 19 + "+" + "-" * 15 + "+" + "-" * 20 + "+" + "-" * 24 + "+")
+
+        if config.SAVE_MODEL:
+            save_checkpoint(model, opt, filename=config.CHECKPOINT)
+
+
+def test():
+    model = EmotionModel().to(config.DEVICE)
+    opt = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, betas=(0.5, 0.999), )
+    criterion = nn.CrossEntropyLoss()
+    if config.LOAD_MODEL:
+        load_checkpoint(
+            config.CHECKPOINT, model, opt, config.LEARNING_RATE,
+        )
+
+    val_dataset = EmotionDataset(root_dir=config.VAL_DIR)
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=config.BATCH_SIZE,
+        shuffle=True,
+        num_workers=config.NUM_WORKERS,
+    )
+    # print(evaluate(model, criterion, val_loader))
+    model.eval()
+    print(val_dataset.class_to_idx)
+    image = np.array(Image.open("images/validation/angry/245.jpg").convert('L'))
+    plt.imshow(image)
+    image = config.transform(image=image)["image"]
+    image = image.to(config.DEVICE)
+    image = torch.unsqueeze(image, dim=0)
+    score = model(image)
+    print(torch.argmax(score))
+    plt.show()
+
+
+def evaluate(model, criterion, dataloader):
+    model.eval()
+    total_correct = 0
+    total_samples = 0
+    total_loss = 0.0
+
+    with torch.no_grad():
+        for inputs, labels in dataloader:
+            inputs, labels = inputs.to(config.DEVICE), labels.to(config.DEVICE)
+
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+
+            total_loss += loss.item()
+
+            _, predicted = torch.max(outputs, 1)
+
+            total_correct += (predicted == labels.argmax(1)).sum().item()
+            total_samples += labels.size(0)
+
+    accuracy = total_correct / total_samples
+    average_loss = total_loss / len(dataloader)
+
+    return accuracy, average_loss
+
+
+if __name__ == "__main__":
+    test()

utils.py

@@ -0,0 +1,21 @@
+import torch
+import config
+
+
+def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
+    print("=> Saving checkpoint")
+    checkpoint = {
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+    }
+    torch.save(checkpoint, filename)
+
+
+def load_checkpoint(checkpoint_file, model, optimizer, lr):
+    print("=> Loading checkpoint")
+    checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
+    model.load_state_dict(checkpoint["state_dict"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr