├── .gitignore ├── FaceRec ├── dataset.py ├── main.py ├── model.py ├── test.py ├── train.py └── trainedmodel.pt └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | __* 2 | -------------------------------------------------------------------------------- /FaceRec/dataset.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, division 2 | import os 3 | import torch 4 | import pandas as pd 5 | from skimage import io, transform 6 | import numpy as np 7 | import matplotlib.pyplot as plt 8 | from torch.utils.data import Dataset, DataLoader 9 | from torchvision import transforms, utils 10 | import random 11 | 12 | class FaceDataset(Dataset): 13 | def __init__(self, csv_file,root_dir, forTrain, transform = None): 14 | self.image_name = pd.read_csv(csv_file) 15 | self.root_dir = root_dir 16 | self.transform = transform 17 | self.forTrain = forTrain 18 | if forTrain: 19 | self.labelSet = {} 20 | for idx in range(len(self.image_name)): 21 | i = self.image_name.iloc[idx,0] 22 | name, label = i.split(' ') 23 | if(not label in self.labelSet): 24 | self.labelSet[label] = [name] 25 | else: 26 | self.labelSet[label].append(name) 27 | 28 | def __len__(self): 29 | return len(self.image_name) 30 | def __getitem__(self, idx): # TODO: Change item selecting. 31 | if self.forTrain: 32 | anchor_name = os.path.join(self.root_dir, self.image_name.iloc[idx,0].split(' ')[0]) 33 | anchor = io.imread(anchor_name) 34 | label = self.image_name.iloc[idx,0].split(' ')[1] 35 | positive_name = anchor_name 36 | count = 0 37 | while anchor_name == positive_name: 38 | count+=1 39 | positive_name = np.random.choice(self.labelSet[label]) 40 | if count == 10: 41 | break 42 | positive = io.imread(os.path.join(self.root_dir, positive_name)) 43 | index_neg = random.randint(0,len(self)-1) 44 | while(self.image_name.iloc[index_neg,0].split(' ')[1]==label): 45 | index_neg = random.randint(0,len(self)-1) 46 | negative = io.imread(os.path.join(self.root_dir, self.image_name.iloc[index_neg,0].split(' ')[0])) 47 | if self.transform: 48 | anchor = self.transform(anchor) 49 | positive = self.transform(positive) 50 | negative = self.transform(negative) 51 | return anchor, positive, negative 52 | 53 | else: 54 | sample = io.imread(os.path.join(self.root_dir, self.image_name.iloc[idx,0].split(' ')[0])) 55 | label = self.image_name.iloc[idx,0].split(' ')[1] 56 | if self.transform: 57 | sample = self.transform(sample) 58 | return sample,label 59 | 60 | class ToTensor(object): 61 | def __call__(self,sample): 62 | image = sample 63 | image = image.transpose((2,0,1)) 64 | return image 65 | 66 | 67 | class Rescale(object): 68 | """Rescale the image in a sample to a given size. 69 | 70 | Args: 71 | output_size (tuple or int): Desired output size. If tuple, output is 72 | matched to output_size. If int, smaller of image edges is matched 73 | to output_size keeping aspect ratio the same. 74 | """ 75 | 76 | def __init__(self, output_size): 77 | assert isinstance(output_size, list) 78 | self.output_size = output_size 79 | 80 | def __call__(self, image): 81 | return transform.resize(image, (self.output_size[0], self.output_size[1])) 82 | -------------------------------------------------------------------------------- /FaceRec/main.py: -------------------------------------------------------------------------------- 1 | from test import * 2 | import warnings 3 | warnings.filterwarnings("ignore") 4 | 5 | if __name__ == '__main__': 6 | runTrain() 7 | runTest() 8 | -------------------------------------------------------------------------------- /FaceRec/model.py: -------------------------------------------------------------------------------- 1 | #This file is expected to be a prototype of convolutional network 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | class Maxout(nn.Module): 8 | 9 | def __init__(self, in_size, out_size, pool_size): 10 | super().__init__() 11 | self.in_size, self.out_size, self.pool_size = in_size, out_size, pool_size 12 | self.lin = nn.Linear(in_size, out_size * pool_size) 13 | 14 | 15 | def forward(self, inputs): 16 | shape = list(inputs.size()) 17 | shape[-1] = self.out_size 18 | shape.append(self.pool_size) 19 | max_dim = len(shape) - 1 20 | out = self.lin(inputs) 21 | m, i = out.view(*shape).max(max_dim) 22 | return m 23 | 24 | class SiameseNetwork(nn.Module): 25 | def __init__(self): 26 | super(SiameseNetwork, self).__init__() 27 | #Explanation of Conv2d(x,y,z): 28 | # Build a convolutional layer with 29 | # x input channels 30 | # y output channels 31 | # z*z filter 32 | self.model1 = nn.Sequential( 33 | nn.Conv2d(3,64,7,stride=2,padding=3), 34 | nn.MaxPool2d(3,stride=2,padding=1) 35 | ) 36 | 37 | self.model2 = nn.Sequential( 38 | nn.Conv2d(64,64,1,stride=1), 39 | nn.Conv2d(64,192,3,stride=1,padding=1) 40 | ) 41 | 42 | self.model3 = nn.Sequential( 43 | nn.MaxPool2d(3,stride=2,padding=1), 44 | nn.Conv2d(192,192,1,stride=1), 45 | nn.Conv2d(192,384,3,stride=1,padding=1), 46 | nn.MaxPool2d(3,stride=2,padding=1), 47 | nn.Conv2d(384,384,1,stride=1), 48 | nn.Conv2d(384,256,3,stride=1,padding=1), 49 | nn.Conv2d(256,256,1,stride=1), 50 | nn.Conv2d(256,256,3,stride=1,padding=1), 51 | nn.Conv2d(256,256,1,stride=1), 52 | nn.Conv2d(256,256,3,stride=1,padding=1), 53 | nn.MaxPool2d(3,stride=2,padding=1) 54 | ) 55 | 56 | self.model4 = nn.Sequential( 57 | Maxout(7*7*256, 1*32*128, 2), 58 | Maxout(1*32*128, 1*32*128, 2), 59 | nn.Linear(1*32*128, 1*1*128), 60 | nn.Linear(128,128) 61 | ) 62 | 63 | 64 | #Fucntion to feed forward 65 | 66 | def forward(self, x): #Go through network once 67 | # Max pooling over a (2, 2) window 68 | x = x.type('torch.cuda.FloatTensor') 69 | x = F.normalize(x) 70 | x = self.model1(x) 71 | x = F.normalize(x) 72 | x = self.model2(x) 73 | x = F.normalize(x) 74 | x = self.model3(x) 75 | x = x.view(-1, self.num_flat_features(x)) 76 | x = self.model4(x) 77 | 78 | return x 79 | 80 | def forward_triple(self, anchor, positive, negative): #Use for computing triplet loss 81 | anchor_output = self.forward(anchor) 82 | positive_output = self.forward(positive) 83 | negative_output = self.forward(negative) 84 | return anchor_output, positive_output, negative_output 85 | 86 | def num_flat_features(self, x): 87 | size = x.size()[1:] # all dimensions except the batch dimension 88 | num_features = 1 89 | for s in size: 90 | num_features *= s 91 | return num_features 92 | -------------------------------------------------------------------------------- /FaceRec/test.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | import torch.optim as optim 4 | from model import * 5 | from dataset import * 6 | from train import * 7 | 8 | def runTest(): 9 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 10 | 11 | composed = transforms.Compose([Rescale(IMAGE_SIZE),ToTensor()]) 12 | network = SiameseNetwork() 13 | network.load_state_dict(torch.load("./"+MODELNAME)) 14 | network.eval() 15 | network.to(device) 16 | labelSet = {} 17 | 18 | #Give each label a standard output 19 | trainedDataset = FaceDataset(csv_file = TEST_FILE, root_dir = IMAGE_DIRECTORY, forTrain = False, transform = composed) 20 | trainedDataloader = DataLoader(trainedDataset, batch_size=1, shuffle=False, num_workers=10) 21 | for i, data in enumerate(trainedDataloader, 0): 22 | sample, label = data 23 | if i%1000 == 0: 24 | print("Go through"+ str(i) + "images") 25 | sample.to(device) 26 | labelSet[label] = network(sample).data 27 | print('Standard output setting: Done.') 28 | 29 | #Go through all trainning data and test the accuracy 30 | testDataset = FaceDataset(csv_file = TEST_FILE, root_dir = IMAGE_DIRECTORY, forTrain = False, transform = composed) 31 | testDataloader = DataLoader(testDataset, batch_size=1, shuffle=False, num_workers=10) 32 | 33 | total_count = 0 34 | correct_count = 0 35 | for i, data in enumerate(trainedDataloader, 0): 36 | sample, label = data 37 | sample.to(device) 38 | output = network(sample).data 39 | distance = 100000 40 | curr_label = '0' 41 | for q in labelSet: 42 | currDis = (output - labelSet[q]).pow(2).sum(1) 43 | if currDis