Added model.py

Author: Mayurji

model.py

@@ -0,0 +1,196 @@
+import random
+import torch
+import torchvision
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, Dataset
+import torch.nn as nn
+from torch import optim
+import torch.nn.functional as F
+from PIL import Image
+import PIL.ImageOps
+import pytorch_lightning as pl
+import numpy as np
+from pytorch_lightning.callbacks import ModelCheckpoint
+import argparse
+
+
+class ContrastiveLoss(torch.nn.Module):
+    """
+    Contrastive loss function.
+    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
+    """
+
+    def __init__(self, margin):
+        super().__init__()
+        self.margin = margin
+
+    def forward(self, output1, output2, label):
+    	euclidean_distance = F.pairwise_distance(output1, output2, keepdim = True)
+    	loss_contrastive = torch.mean((1-label) * torch.pow(euclidean_distance, 2) + (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
+
+    	return loss_contrastive
+
+class SiameseNetworkDataset(Dataset):
+    
+    def __init__(self,imageFolderDataset,transform=None,should_invert=None):
+        self.imageFolderDataset = imageFolderDataset    
+        self.transform = transform
+        self.should_invert = should_invert
+        
+    def __getitem__(self,index):
+        img0_tuple = random.choice(self.imageFolderDataset.imgs)
+        #we need to make sure approx 50% of images are in the same class
+        should_get_same_class = random.randint(0,1) 
+        if should_get_same_class:
+            while True:
+                #keep looping till the same class image is found
+                img1_tuple = random.choice(self.imageFolderDataset.imgs) 
+                if img0_tuple[1]==img1_tuple[1]:
+                    break
+        else:
+            while True:
+                #keep looping till a different class image is found
+                
+                img1_tuple = random.choice(self.imageFolderDataset.imgs) 
+                if img0_tuple[1] !=img1_tuple[1]:
+                    break
+
+        img0 = Image.open(img0_tuple[0])
+        img1 = Image.open(img1_tuple[0])
+        img0 = img0.convert("L")
+        img1 = img1.convert("L")
+        
+        if self.should_invert:
+            img0 = PIL.ImageOps.invert(img0)
+            img1 = PIL.ImageOps.invert(img1)
+
+        if self.transform is not None:
+            img0 = self.transform(img0)
+            img1 = self.transform(img1)
+        
+        return img0, img1 , torch.from_numpy(np.array([int(img1_tuple[1]!=img0_tuple[1])],dtype=np.float32))
+    
+    def __len__(self):
+        return len(self.imageFolderDataset.imgs)
+
+
+class SiameseNetwork(pl.LightningModule):
+    def __init__(self, margin, learning_rate, resize, imageFolderTrain, imageFolderTest, batch_size, should_invert):
+        super().__init__()
+        self.imageFolderTrain = imageFolderTrain
+        self.imageFolderTest= imageFolderTest
+        self.learning_rate = learning_rate
+        self.criterion = ContrastiveLoss(margin=margin)
+        self.batch_size = batch_size
+        self.should_invert = should_invert
+        self.transform = transforms.Compose([transforms.Resize((resize,resize)),
+                                             transforms.RandomHorizontalFlip(),
+                                             transforms.ToTensor()])
+        
+        self.cnn1 = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(1, 4, kernel_size=3),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(4),
+            
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(4, 8, kernel_size=3),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(8),
+
+
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(8, 8, kernel_size=3),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(8),
+
+        )
+
+        self.fc1 = nn.Sequential(
+            nn.Linear(8*100*100, 500),
+            nn.ReLU(inplace=True),
+
+            nn.Linear(500, 500),
+            nn.ReLU(inplace=True),
+
+            nn.Linear(500, 5))
+    
+    def forward_once(self, x):
+        output = self.cnn1(x)
+        output = output.view(output.size()[0], -1)
+        output = self.fc1(output)
+        return output
+
+    
+    def forward(self, input1, input2):
+        output1 = self.forward_once(input1)
+        output2 = self.forward_once(input2)
+        return output1, output2
+
+    
+    def training_step(self, batch, batch_idx):
+        x0, x1 , y = batch
+        output1,output2 = self(x0, x1)
+        loss = self.criterion(output1,output2, y)
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        x0, x1 , y = batch
+        output1,output2 = self(x0, x1)
+        loss = self.criterion(output1,output2,y)
+
+        self.log('val_loss', loss, prog_bar=True)
+        return loss
+    
+    def test_step(self, batch, batch_idx):
+        return self.validation_step(batch, batch_idx)
+    
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        return optimizer
+            
+    def prepare_data(self):
+        self.DatasetFolder = dset.ImageFolder(self.imageFolderTrain)
+        self.DatasetFolder_testing = dset.ImageFolder(self.imageFolderTest)
+    
+    def setup(self, stage=None):
+        
+        self.siamese_dataset_train = SiameseNetworkDataset(imageFolderDataset=self.DatasetFolder,
+                                        transform=self.transform
+                                       ,should_invert=self.should_invert)
+        self.siamese_dataset_test = SiameseNetworkDataset(imageFolderDataset=self.DatasetFolder_testing,
+                                        transform=self.transform
+                                       ,should_invert=self.should_invert)
+
+    def train_dataloader(self):
+        return DataLoader(self.siamese_dataset_train, batch_size=self.batch_size)
+    
+    def test_dataloader(self):
+        return DataLoader(self.siamese_dataset_test, batch_size=self.batch_size)
+
+if __name__=='__main__':
+
+	parser = argparse.ArgumentParser(
+        description='Siamese Network - Face Recognition',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+	parser.add_argument('--gpus', default=1, type=int)
+	parser.add_argument('--batch_size', default=64, type=int)
+	parser.add_argument('--pretrain_epochs', default=5000, type=int)
+	parser.add_argument('--margin', default=1.0, type=float)
+	parser.add_argument('--should_invert', default=False)
+	parser.add_argument('--imageFolderTrain', default=None)
+	parser.add_argument('--imageFolderTest', default=None)
+	parser.add_argument('--learning_rate', default=2e-2, type=float)
+	parser.add_argument('--resize', default=100, type=int)
+
+
+	args = parser.parse_args()
+	print(args)
+
+	model = SiameseNetwork(margin= args.margin, learning_rate=args.learning_rate, resize=args.resize, imageFolderTrain=args.imageFolderTrain,
+							imageFolderTest=args.imageFolderTest, batch_size=args.batch_size, should_invert=args.should_invert)
+	trainer = pl.Trainer(gpus=args.gpus, max_epochs=args.pretrain_epochs, progress_bar_refresh_rate=20)
+	trainer.fit(model)
+	trainer.save_checkpoint("siamese_face_recognition.ckpt")
+	trainer.test()

utils.py

@@ -0,0 +1,46 @@
+import torch
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, Dataset
+import torchvision
+import matplotlib.pyplot as plt
+import torch.nn.functional as F
+import numpy as np
+from model import SiameseNetwork
+%matplotlib inline
+
+model = SiameseNetwork.load_from_checkpoint("siamese_face_recognition.ckpt")
+
+def imshow(img,text=None,should_save=False):
+    npimg = img.numpy()
+    plt.axis("off")
+    if text:
+        plt.text(75, 8, text, style='italic',fontweight='bold',
+            bbox={'facecolor':'white', 'alpha':0.8, 'pad':10})
+    plt.imshow(np.transpose(npimg, (1, 2, 0)))
+    plt.show()    
+
+def show_plot(iteration,loss):
+    plt.plot(iteration,loss)
+    plt.show()
+
+
+# def view_dissimilarity(testDirectory):
+#     folder_dataset_test = dset.ImageFolder(root=testDirectory)
+#     siamese_dataset = SiameseNetworkDataset(imageFolderDataset=folder_dataset_test,
+#                                             transform=transforms.Compose([transforms.Resize((100,100)),
+#                                                                           transforms.ToTensor()
+#                                                                           ])
+#                                            ,should_invert=False)
+
+#     test_dataloader = DataLoader(siamese_dataset,num_workers=6,batch_size=1,shuffle=True)
+#     dataiter = iter(test_dataloader)
+#     x0,_,_ = next(dataiter)
+
+#     for i in range(4):
+#         _,x1,label2 = next(dataiter)
+#         concatenated = torch.cat((x0,x1),0)
+#         model.eval()
+#         output1,output2 = model(Variable(x0).cuda(),Variable(x1).cuda())
+#         euclidean_distance = F.pairwise_distance(output1, output2)
+#         imshow(torchvision.utils.make_grid(concatenated),'Dissimilarity: {:.2f}'.format(euclidean_distance.item()))