-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
69 lines (52 loc) · 1.97 KB
/
test.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
from model.classifier import ActivityClassifier, anno2padded, normalize_res, get_label
from model.discriminator import Discriminator
from model.generator import Generator
import torch
import cv2
import pickle
import numpy as np
import matplotlib.pyplot as plt
opt = {'b1': 0.5,
'b2': 0.999,
'batch_size': 64,
'channels': 1,
'img_size': 64,
'latent_dim': 32,
'lr': 0.0002,
'n_classes': 200,
'n_cpu': 8,
'n_epochs': 1,
'sample_interval': 400}
mydevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load Model:
print("Loading Model...")
classifier = torch.load('./intermediate/classifier_200.pth', map_location = mydevice)
netG_pose = torch.load('./intermediate/pose_netG_uni.pth', map_location = mydevice)
netG_parsing = torch.load('./intermediate/parsing_netG_uni.pth', map_location = mydevice)
with open('./intermediate/word2idx.pk', 'rb') as f:
word2idx = pickle.load(f)
# User entered annotation:
annotate = input('Please enter an annotation (<16 words): ')
while len(annotate.split(' ')) > 15:
annotate = input('[!] Max length exceeded. Please enter an annotation (<16 words): ')
# Retrieving the potential cluster label based on pose distribution
s = anno2padded(annotate, word2idx, mydevice)
with torch.no_grad():
res = classifier(s)
res = normalize_res(res)
label = get_label(res)
# Generating pose and parsing based on retrieved label
z = torch.randn(1, opt['latent_dim'], device = mydevice)
label = torch.tensor([label], device = mydevice)
netG_parsing.eval()
netG_pose.eval()
with torch.no_grad():
gen_parsing = netG_parsing(z, label)
gen_pose = netG_pose(z, label)
img_parsing = np.reshape(gen_parsing[0].cpu().detach().numpy(), (64, 64, 3))
img_pose = np.reshape(gen_pose[0].cpu().detach().numpy(), (64, 64, 3))
plt.imshow(img_parsing)
plt.savefig('./output/sample_parsing.png')
plt.imshow(img_pose)
plt.savefig('./output/sample_pose.png')
print('Pose, Parsing Generated.')