hackathon-4 / app /Hackathon_setup /face_recognition_model.py
care2achieve's picture
updated arch
5e64ee3
Raw
History Blame Contribute Delete
4.09 kB
import math
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
# Add more imports if required
# Sample Transformation function
# YOUR CODE HERE for changing the Transformation values.
trnscm = transforms.Compose([
transforms.Grayscale(
num_output_channels=1
),
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.5],
std=[0.5]
)
])
##Example Network
class Siamese(torch.nn.Module):
def __init__(self):
super(Siamese, self).__init__()
################################################################
# CNN Feature Extractor
################################################################
self.cnn1 = nn.Sequential(
# Block 1
nn.Conv2d(
in_channels=1,
out_channels=32,
kernel_size=3,
padding=1
),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
# 224 -> 112
# Block 2
nn.Conv2d(
in_channels=32,
out_channels=64,
kernel_size=3,
padding=1
),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
# 112 -> 56
# Block 3
nn.Conv2d(
in_channels=64,
out_channels=128,
kernel_size=3,
padding=1
),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
# 56 -> 28
# Block 4
nn.Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
padding=1
),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
################################################################
# Makes output independent of image size
################################################################
nn.AdaptiveAvgPool2d((1, 1))
)
################################################################
# Embedding Head
################################################################
self.fc1 = nn.Sequential(
nn.Flatten(),
nn.Linear(256, 128),
nn.ReLU(inplace=True),
nn.Dropout(0.3),
nn.Linear(128, 128)
)
####################################################################
# Generate embedding for one image
####################################################################
def forward_once(self, x):
output = self.cnn1(x)
output = self.fc1(output)
################################################################
# L2 normalize embeddings
################################################################
output = F.normalize(output, p=2, dim=1)
return output
####################################################################
# Siamese forward
####################################################################
def forward(self, input1, input2):
output1 = self.forward_once(input1)
output2 = self.forward_once(input2)
return output1, output2
##########################################################################################################
## Sample classification network (Specify if you are using a pytorch classifier during the training) ##
## classifier = nn.Sequential(nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Linear...) ##
##########################################################################################################
# YOUR CODE HERE for pytorch classifier
# Definition of classes as dictionary
classes = ['person1','person2','person3','person4','person5','person6','person7']