-
Notifications
You must be signed in to change notification settings - Fork 15
/
generate_images.py
92 lines (74 loc) · 3.63 KB
/
generate_images.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
from importlib import import_module
import pickle as pkl
from utils.data_utils import *
from utils.training_utils import ModelCheckpoint
from config import Config
import argparse
import matplotlib.pyplot as plt
class ImgGenerator:
def __init__(self, checkpt_path, config, char_map=None):
"""
:param checkpt_path: Path of the model checkpoint file to be used
:param config: Config with all the parameters to be used
"""
self.config = config
if char_map is None:
with open(config.data_file, 'rb') as f:
data = pkl.load(f)
self.char_map = data['char_map']
else:
self.char_map = char_map
# Model
print(f'Model: {config.architecture}')
model_type = import_module('models.' + self.config.architecture)
create_model = getattr(model_type, 'create_model')
self.model = create_model(self.config, self.char_map)
# print(self.model, end='\n\n')
self.model.to(self.config.device)
self.model.eval()
self.word_map = WordMap(self.char_map)
# Load model weights
self.model_checkpoint = ModelCheckpoint(config=self.config)
self.model, _, _, _ = self.model_checkpoint.load(self.model, epoch=None, checkpoint_path=checkpt_path)
def generate(self, random_num_imgs=5, word_list=None, z=None):
"""
Returns images generated by the trained generator model
:param random_num_imgs: Number of images to be randomly generated using lexicon (only valid if word_list=None)
:param word_list: List of words for which images need to be generated
:param z: Noise vector determining the style of the images to be generated (32 dimension vector)
"""
if word_list is None:
# Generate random images
with torch.no_grad():
self.model.forward_fake(z=None, b_size=random_num_imgs)
else:
encoded_words, _ = self.word_map.encode(word_list)
with torch.no_grad():
self.model.forward_fake(z=None, fake_y=encoded_words, b_size=len(word_list))
word_labels_decoded = self.word_map.decode(self.model.fake_y.cpu().numpy())
return self.model.fake_img.squeeze(1).cpu().numpy(), self.model.fake_y.cpu().numpy(), word_labels_decoded
if __name__ == "__main__":
# Construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-c", "--checkpt_path", required=True, type=str,
help="Path of the model checkpoint file to be used")
ap.add_argument("-m", "--char_map_path", required=True, type=str,
help="Path of the file with character mapping to be used")
ap.add_argument("-n", "--num_imgs", required=False, type=int,
help="number of sample points")
ap.add_argument("-w", "--word_list", required=False, nargs='+', default=[],
help="words for which images need to be generated")
args = vars(ap.parse_args())
checkpoint_path = args['checkpt_path']
char_map_path = args['char_map_path']
num_imgs = args['num_imgs'] if args['num_imgs'] is not None else 5
word_list = args['word_list'] if len(args['word_list']) > 0 else None
with open(f'{char_map_path}', 'rb') as f:
char_map = pkl.load(f)
config = Config
generator = ImgGenerator(checkpt_path=checkpoint_path, config=config, char_map=char_map)
generated_imgs, _, word_labels = generator.generate(num_imgs, word_list)
for label, img in zip(word_labels, generated_imgs):
plt.imshow(img, cmap='gray')
plt.title(label)
plt.show()