-
Notifications
You must be signed in to change notification settings - Fork 89
/
vis.py
136 lines (112 loc) · 4.76 KB
/
vis.py
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
from __future__ import print_function
import argparse
import torch
from torch.autograd import Variable
from torchvision import datasets, transforms
from model import FaceModel
from torchvision.datasets import ImageFolder
from TripletFaceDataset import FaceDataset
import math
import os
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as PathEffects
parser = argparse.ArgumentParser(description='PyTorch face recognition Example')
parser.add_argument('--test_batch_size', type=int, default=64, metavar='N',
help='input batch size for testing (default: 64)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--root', type=str,
help='path to the data directory containing aligned face patches. Multiple directories are separated with colon.',
default='/media/lior/LinuxHDD/datasets/vgg_face_dataset/aligned')
parser.add_argument('--resume', type=str,
help='model path to the resume training',
default='/home/lior/dev/workspace/face_recognition_seminar/facenet_pytorch/logs/run-optim_adagrad-n1000000-lr0.125-wd0.0-m0.5/checkpoint_1.pth')
def visual_feature_space(features, labels, num_classes, name_dict):
num = len(labels)
title_font = {'fontname':'Arial', 'size':'20', 'color':'black', 'weight':'normal',
'verticalalignment':'bottom'} # Bottom vertical alignment for more space
axis_font = {'fontname':'Arial', 'size':'20'}
# draw
palette = np.array(sns.color_palette("hls", num_classes))
# We create a scatter plot.
f = plt.figure(figsize=(8, 8))
ax = plt.subplot(aspect='equal')
sc = ax.scatter(features[:,0], features[:,1], lw=0, s=40,
c=palette[labels.astype(np.int)])
# ax.axis('off')
# ax.axis('tight')
# We add the labels for each digit.
txts = []
for i in range(num_classes):
# Position of each label.
xtext, ytext = np.median(features[labels == i, :], axis=0)
txt = ax.text(xtext, ytext, name_dict[i])
txt.set_path_effects([
PathEffects.Stroke(linewidth=5, foreground="w"),
PathEffects.Normal()])
txts.append(txt)
ax.set_xlabel('Activation of the 1st neuron', **axis_font)
ax.set_ylabel('Activation of the 2nd neuron', **axis_font)
ax.set_title('softmax_loss + center_loss', **title_font)
ax.set_axis_bgcolor('grey')
f.savefig('center_loss.png')
plt.show()
return f, ax, sc, txts
def validation_iterator(dataLoader):
for data, target in dataLoader:
yield data, target
def main():
args = parser.parse_args()
cuda = torch.cuda.is_available()
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed(args.seed)
# 1. dataset
root = args.root
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
test_transforms = transforms.Compose([transforms.Scale(96),
transforms.ToTensor(),
transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])
test_dataset = ImageFolder(root, transform=test_transforms)
test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=args.test_batch_size, shuffle=False, **kwargs)
val_iterator = validation_iterator(test_loader)
# 2. model
#train_dir = FaceDataset(dir='/media/lior/LinuxHDD/datasets/MSCeleb-cleaned',n_triplets=10)
print('construct model')
model = FaceModel(embedding_size=128,
num_classes=3367,
pretrained=False)
model = model.cuda()
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}'".format(args.resume))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# extract feature
print('extracting feature')
embeds = []
labels = []
for data, target in val_iterator:
if cuda:
data, target = data.cuda(), target.cuda(async=True)
data_var = Variable(data, volatile=True)
# compute output
output = model(data_var)
embeds.append( output.data.cpu().numpy() )
labels.append( target.cpu().numpy() )
embeds = np.vstack(embeds)
labels = np.hstack(labels)
print('embeds shape is ', embeds.shape)
print('labels shape is ', labels.shape)
# prepare dict for display
namedict = dict()
for i in range(10):
namedict[i]=str(i)
visual_feature_space(embeds, labels, len(test_dataset.classes), namedict)
if __name__ == '__main__':
main()