MaskedFaceRepresentation/predict.py at master · sachith500/MaskedFaceRepresentation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import os
import torchvision
from matplotlib import pyplot as plt
import torch
import glob
import os
from shutil import copyfile
from io import BytesIO
from PIL import Image
import time
import heapq
from heapq import heappush
from heapq import heappop
from torchvision import datasets
import sys


import gc
gc.collect()
torch.cuda.empty_cache()

class MyImageFolder(datasets.ImageFolder):
   def getitem(self, index):
      return super(MyImageFolder, self).getitem(index), self.imgs[index]#return image path

class ImageFolderWithPaths(datasets.ImageFolder):
    """Custom dataset that includes image file paths. Extends
    torchvision.datasets.ImageFolder
    """

    # override the __getitem__ method. this is the method that dataloader calls
    def __getitem__(self, index):
        # this is what ImageFolder normally returns
        original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
        # the image file path
        path = self.imgs[index][0]
        # make a new tuple that includes original and the path
        tuple_with_path = (original_tuple + (path,))
        return tuple_with_path

checkpoint_path = 'model_best.pth.tar'
data_folder = "/data/scratch/projects/punim1439/Mask/data/gender_separate/test/"

num_classes = 2
shape = 224
BATCHSIZE = 64


transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize(shape),
        torchvision.transforms.CenterCrop(shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225]),
])

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

arch = 'resnet50'
model = torchvision.models.__dict__[arch](pretrained=False)
model.fc = torch.nn.Linear(2048, num_classes)
model.fc.weight.data.normal_(mean=0.0, std=0.01)
model.fc.bias.data.zero_()

state_dict = torch.load(checkpoint_path)

parallel_model = torch.nn.DataParallel(model)
parallel_model.load_state_dict(
    torch.load(checkpoint_path, map_location=str(device))['state_dict']
)

model = parallel_model.module

#print(model)

model.eval()
model.to(device)

data_transforms = {
    'predict': transform
    }

#MyImageFolder
#dataset = datasets.ImageFolder(data_folder, transform=transform)#
dataset = ImageFolderWithPaths(data_folder, transform=transform)
dataloader = {'predict': torch.utils.data.DataLoader(dataset, batch_size = BATCHSIZE , shuffle=False, num_workers=8, pin_memory=True)}

count = 0
total_correct = 0
with open('submission.txt', 'w') as the_file:
  the_file.write('Name, Label, Predicted' + "\n")
  for inputs, labels, filenames in dataloader['predict']:
    #print(filenames)
    count += 1
    if (count % (1000) == 0):
        print(count)
        print(time.time() - since)
    inputs = inputs.to(device)
    output = model(inputs)
    output = output.to(device)
    out_labels = torch.argmax(output, dim=1)
    out_labels = out_labels.to('cpu')
    labels = labels.to('cpu')
    correct = (out_labels == labels).float().sum()
    for i in range(len(output)):
        #out_str = str(count) + "_" + str(i) + "," + str(labels[i].item()) + "," + str(out_labels[i].item())
        out_str = filenames[i].split("/")[-1] + "," + str(labels[i].item()) + "," + str(out_labels[i].item())
        the_file.write(out_str + "\n")

    #print(correct)
    total_correct += correct
    #print(out_labels)

print("Accuracy = " + str(total_correct) + "/" + str(len(dataloader['predict'].dataset)))
print(" ===  " + str(total_correct/len(dataloader['predict'].dataset)) + "  ===")