Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torchvision import transforms | |
| #################################################################################################################### | |
| # Expression Recognition Model | |
| #################################################################################################################### | |
| # Definition of classes as dictionary | |
| classes = { | |
| 0: 'ANGER', | |
| 1: 'DISGUST', | |
| 2: 'FEAR', | |
| 3: 'HAPPINESS', | |
| 4: 'NEUTRAL', | |
| 5: 'SADNESS', | |
| 6: 'SURPRISE' | |
| } | |
| # ============================================================ | |
| # RESIDUAL BLOCK | |
| # ============================================================ | |
| class ResidualBlock(nn.Module): | |
| def __init__(self, in_channels, out_channels, stride=1): | |
| super().__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_channels, | |
| out_channels, | |
| kernel_size=3, | |
| stride=stride, | |
| padding=1, | |
| bias=False | |
| ) | |
| self.bn1 = nn.BatchNorm2d(out_channels) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv2 = nn.Conv2d( | |
| out_channels, | |
| out_channels, | |
| kernel_size=3, | |
| padding=1, | |
| bias=False | |
| ) | |
| self.bn2 = nn.BatchNorm2d(out_channels) | |
| self.shortcut = nn.Sequential() | |
| if stride != 1 or in_channels != out_channels: | |
| self.shortcut = nn.Sequential( | |
| nn.Conv2d( | |
| in_channels, | |
| out_channels, | |
| kernel_size=1, | |
| stride=stride, | |
| bias=False | |
| ), | |
| nn.BatchNorm2d(out_channels) | |
| ) | |
| def forward(self, x): | |
| identity = self.shortcut(x) | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = self.conv2(out) | |
| out = self.bn2(out) | |
| out += identity | |
| out = self.relu(out) | |
| return out | |
| # ============================================================ | |
| # MODEL | |
| # ============================================================ | |
| class facExpRec(nn.Module): | |
| def __init__(self, num_classes=7): | |
| super().__init__() | |
| self.features = nn.Sequential( | |
| nn.Conv2d( | |
| 1, | |
| 32, | |
| kernel_size=3, | |
| padding=1 | |
| ), | |
| nn.BatchNorm2d(32), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(2), | |
| ResidualBlock(32, 64), | |
| nn.MaxPool2d(2), | |
| ResidualBlock(64, 128), | |
| nn.MaxPool2d(2), | |
| ResidualBlock(128, 256), | |
| nn.MaxPool2d(2), | |
| ResidualBlock(256, 512), | |
| nn.AdaptiveAvgPool2d((1, 1)) | |
| ) | |
| self.classifier = nn.Sequential( | |
| nn.Flatten(), | |
| nn.Linear(512, 256), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(0.4), | |
| nn.Linear(256, 128), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(0.3), | |
| nn.Linear(128, num_classes) | |
| ) | |
| def forward(self, x): | |
| x = self.features(x) | |
| x = self.classifier(x) | |
| return x | |
| # ============================================================ | |
| # HELPER FUNCTIONS | |
| # ============================================================ | |
| def rgb2gray(image): | |
| return image.convert('L') | |
| # ============================================================ | |
| # TRANSFORMS | |
| # ============================================================ | |
| trnscm = transforms.Compose([ | |
| rgb2gray, | |
| transforms.Resize((72, 72)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.5], | |
| std=[0.5] | |
| ) | |
| ]) | |