def test(cfg_file, embedding_t7_path):
cfg_from_file(cfg_file)
cfg.GPU_ID = 0
print('Using config:')
pprint.pprint(cfg)
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=[0])
print(netG)
#state_dict = torch.load('../gan/models/birds_3stages/netG_26000.pth')
print(cfg.TRAIN.NET_G)
state_path = '/home/ubuntu/GANTextToImage/gan/' + cfg.TRAIN.NET_G
state_dict = torch.load(state_path, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
# print('Load ', '../gan/models/flowers_1stage/netG_6000.pth')
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(1, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
t_embedding = load_lua(embedding_t7_path)
t_embedding = t_embedding.unsqueeze(0)
print(t_embedding.size())
if cfg.CUDA:
t_embedding = Variable(t_embedding).cuda()
else:
t_embedding = Variable(t_embedding)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embedding.size(1)
noise.data.resize_(1, nz)
noise.data.normal_(0, 1)
fake_img_list = []
# for i in range(embedding_dim):
fake_imgs, _, _ = netG(noise, t_embedding[:, 0, :])
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
save_singleimages(fake_imgs[-1], '/home/ubuntu/GANTextToImage/static', 0, 256)
def load_network_stageII(self):
from model import G_NET, D_NET
netG = G_NET()
netG.apply(weights_init)
print(netG)
if cfg.NET_G != '':
state_dict = \
torch.load(cfg.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load from: ', cfg.NET_G)
netD = D_NET()
netD.apply(weights_init)
if cfg.NET_D != '':
state_dict = \
torch.load(cfg.NET_D,
map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load from: ', cfg.NET_D)
print(netD)
if cfg.CUDA:
netG.cuda()
netD.cuda()
return netG, netD
def load_network(gpus):
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
if cfg.TREE.BRANCH_NUM > 3:
netsD.append(D_NET512())
if cfg.TREE.BRANCH_NUM > 4:
netsD.append(D_NET1024())
# TODO: if cfg.TREE.BRANCH_NUM > 5:
for i in range(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
# print(netsD[i])
print('# of netsD', len(netsD))
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
try:
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count)
except:
last_run_dir = cfg.DATA_DIR + '/' + cfg.LAST_RUN_DIR + '/Model'
with open(last_run_dir + '/count.txt', 'r') as f:
count = int(f.read())
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
netsD[i].load_state_dict(state_dict)
inception_model = INCEPTION_V3()
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
inception_model.eval()
return netG, netsD, len(netsD), inception_model, count
def evaluate_finegan(self):
if cfg.TRAIN.NET_G == '':
print('Error: the path for model not found!')
else:
# Build and load the generator
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
model_dict = netG.state_dict()
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)['birds128']
state_dict = {k: v for k, v in state_dict.items() if k in model_dict}
model_dict.update(state_dict)
netG.load_state_dict(model_dict)
print('Load ', cfg.TRAIN.NET_G)
# Uncomment this to print Generator layers
print(netG)
sys
nz = cfg.GAN.Z_DIM
noise = torch.FloatTensor(self.batch_size, nz)
noise.data.normal_(0, 1)
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
netG.eval()
background_class = cfg.TEST_BACKGROUND_CLASS
parent_class = cfg.TEST_PARENT_CLASS
child_class = cfg.TEST_CHILD_CLASS
bg_code = torch.zeros([self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
p_code = torch.zeros([self.batch_size, cfg.SUPER_CATEGORIES])
c_code = torch.zeros([self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
for j in range(self.batch_size):
bg_code[j][background_class] = 1
p_code[j][parent_class] = 1
c_code[j][child_class] = 1
fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(noise, c_code, p_code, bg_code) # Forward pass through the generator
self.save_image(fake_imgs[0][0], self.save_dir, 'background')
self.save_image(fake_imgs[1][0], self.save_dir, 'parent_final')
self.save_image(fake_imgs[2][0], self.save_dir, 'child_final')
self.save_image(fg_imgs[0][0], self.save_dir, 'parent_foreground')
self.save_image(fg_imgs[1][0], self.save_dir, 'child_foreground')
self.save_image(mk_imgs[0][0], self.save_dir, 'parent_mask')
self.save_image(mk_imgs[1][0], self.save_dir, 'child_mask')
self.save_image(fgmk_imgs[0][0], self.save_dir, 'parent_foreground_masked')
self.save_image(fgmk_imgs[1][0], self.save_dir, 'child_foreground_masked')
def load_network(gpus, dictionary=None):
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
if cfg.TREE.BRANCH_NUM > 3:
netsD.append(D_NET512())
if cfg.TREE.BRANCH_NUM > 4:
netsD.append(D_NET1024())
# TODO: if cfg.TREE.BRANCH_NUM > 5:
for i in range(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
# print(netsD[i])
print('# of netsD', len(netsD))
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
netsD[i].load_state_dict(state_dict)
inception_model = INCEPTION_V3()
embedding_model = load_embedding_model(dictionary)
for param in embedding_model.parameters():
param.requires_grad = False
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
embedding_model.cuda()
inception_model.eval()
return netG, netsD, len(netsD), inception_model, embedding_model, count
def load_network(gpus):
netEn_img = MLP_ENCODER_IMG()
netEn_img.apply(weights_init)
netEn_img = torch.nn.DataParallel(netEn_img, device_ids=gpus)
print(netEn_img)
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
for i in xrange(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
print('# of netsD', len(netsD))
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in xrange(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
netsD[i].load_state_dict(state_dict)
if cfg.TRAIN.NET_MLP_IMG != '':
state_dict = torch.load(cfg.TRAIN.NET_MLP_IMG)
netEn_img.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_MLP_IMG)
inception_model = INCEPTION_V3()
if cfg.CUDA:
netG.cuda()
netEn_img = netEn_img.cuda()
for i in xrange(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
inception_model.eval()
return netG, netsD, netEn_img, inception_model, len(netsD), count
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d/%s' % (s_tmp, iteration, split_dir)
if cfg.TEST.B_EXAMPLE:
folder = '%s/super' % (save_dir)
else:
folder = '%s/single' % (save_dir)
print('Make a new folder: ', folder)
mkdir_p(folder)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
num_batches = int(cfg.TEST.SAMPLE_NUM / self.batch_size)
cnt = 0
for step in range(num_batches):
noise.data.normal_(0, 1)
#hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxstart
fake_imgs, layers_output, _, _ = netG(noise)
if len(layers_output) != len(lamdas):
print("please check lamdas length")
#hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxend
if cfg.TEST.B_EXAMPLE:
self.save_superimages(fake_imgs[-1], folder, cnt, 256)
else:
self.save_singleimages(fake_imgs[-1], folder, cnt, 256)
# self.save_singleimages(fake_imgs[-2], folder, 128)
# self.save_singleimages(fake_imgs[-3], folder, 64)
cnt += self.batch_size
def load_network(gpus):
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
if cfg.TREE.BRANCH_NUM > 3:
netsD.append(D_NET512())
if cfg.TREE.BRANCH_NUM > 4:
netsD.append(D_NET1024())
# TODO: if cfg.TREE.BRANCH_NUM > 5:
for i in range(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
# print(netsD[i])
print('# of netsD', len(netsD))
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
netsD[i].load_state_dict(state_dict)
inception_model = INCEPTION_V3()
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
inception_model.eval()
return netG, netsD, len(netsD), inception_model, count
def build_models(self):
netG = G_NET(len(self.cats_index_dict))
netINSD = INS_D_NET(len(self.cats_index_dict))
netGLBD = GLB_D_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netINSD.apply(weights_init)
netGLBD.apply(weights_init)
if cfg.CUDA:
netG.cuda()
netINSD.cuda()
netGLBD.cuda()
if len(cfg.GPU_IDS) > 1:
netG = nn.DataParallel(netG)
netG.to(self.device)
netINSD = nn.DataParallel(netINSD)
netINSD.to(self.device)
netGLBD = nn.DataParallel(netGLBD)
netGLBD.to(self.device)
# ########################################################### #
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
filename = path_leaf(cfg.TRAIN.NET_G)
istart = filename.rfind('_') + 1
iend = filename.rfind('.')
epoch = filename[istart:iend]
epoch = int(epoch) + 1
Gname = cfg.TRAIN.NET_G
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netINSD.pth' % (s_tmp)
print('Load INSD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netINSD.load_state_dict(state_dict)
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netGLBD.pth' % (s_tmp)
print('Load GLBD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netGLBD.load_state_dict(state_dict)
return [netG, netINSD, netGLBD, epoch]
def load_network(gpus):
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
# 128 * 128
if cfg.TREE.BRANCH_NUM > 1:
for i in range(
3): # 3 discriminators for background, parent and child stage
netsD.append(D_NET128(i))
# 256 * 256
if cfg.TREE.BRANCH_NUM > 2:
for i in range(
3): # 3 discriminators for background, parent and child stage
netsD.append(D_NET256(i))
# for i in range(3): # 3 discriminators for background, parent and child stage
# netsD.append(D_NET128(i))
for i in range(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
print(netsD[i])
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s_%d.pth' % (cfg.TRAIN.NET_D, i))
netsD[i].load_state_dict(state_dict)
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return netG, netsD, len(netsD), count
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d/%s' % (s_tmp, iteration, split_dir)
if cfg.TEST.B_EXAMPLE:
folder = '%s/super' % (save_dir)
else:
folder = '%s/single' % (save_dir)
print('Make a new folder: ', folder)
mkdir_p(folder)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
num_batches = int(cfg.TEST.SAMPLE_NUM / self.batch_size)
cnt = 0
for step in xrange(num_batches):
noise.data.normal_(0, 1)
fake_imgs, _, _ = netG(noise)
if cfg.TEST.B_EXAMPLE:
self.save_superimages(fake_imgs[-1], folder, cnt, 256)
else:
self.save_singleimages(fake_imgs[-1], folder, cnt, 256)
# self.save_singleimages(fake_imgs[-2], folder, 128)
# self.save_singleimages(fake_imgs[-3], folder, 64)
cnt += self.batch_size
def load_checkpoint(modelpath):
s_gpus = cfg.GPU_ID.split(',')
gpus = [int(ix) for ix in s_gpus]
torch.cuda.set_device(gpus[0])
state_dict = torch.load(modelpath, map_location=lambda storage, loc: storage)
#print(checkpoint.keys())
#model = checkpoint['model']
#model.load_state_dict(checkpoint['state_dict'])
#for parameter in model.parameters():
# parameter.requires_grad = False
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
netG.load_state_dict(state_dict)
netG.eval()
return netG
def load_network(gpus):
netG = G_NET(start_depth)
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
netsD.append(D_NET_BG(start_depth))
netsD.append(D_NET_PC(1, start_depth))
netsD.append(D_NET_PC(2, start_depth))
netsD.append(D_NET_BG_PG(start_depth))
for i in range(len(netsD)):
netsD[i].apply(weights_init)
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
print(netsD[i])
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('netG_') + 5
iend = cfg.TRAIN.NET_G.rfind('_depth')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count)
istart = cfg.TRAIN.NET_G.rfind('depth')
iend = cfg.TRAIN.NET_G.rfind('.')
_depth = cfg.TRAIN.NET_G[istart:iend]
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s%d_%s.pth' % (cfg.TRAIN.NET_D, i, _depth))
state_dict = torch.load('%s%d_%s.pth' %
(cfg.TRAIN.NET_D, i, _depth))
netsD[i].load_state_dict(state_dict)
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return netG, netsD, len(netsD), count
def sampling(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
netG = G_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netG.eval()
if cfg.CUDA:
netG.cuda()
if len(cfg.GPU_IDS) > 1:
netG = nn.DataParallel(netG)
netG.to(self.device)
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(
torch.FloatTensor(batch_size, cfg.ROI.BOXES_NUM,
len(self.cats_index_dict) * 4))
noise = noise.cuda()
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
# state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save generated images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s/%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
cnt = 0
for _ in range(1): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if step % 100 == 0:
print('step: ', step)
# if step > 50:
# break
imgs, pooled_hmaps, hmaps, bbox_maps_fwd, bbox_maps_bwd, bbox_fmaps, \
rois, fm_rois, num_rois, class_ids, keys = prepare_data(data)
num_rois = num_rois.data.cpu().numpy()
cats_list = []
for batch_index in range(self.batch_size):
cats = []
for roi_index in range(num_rois[batch_index]):
rela_cat_id = int(rois[batch_index, roi_index, 4])
abs_cat_id = self.cats_dict[rela_cat_id][0]
cat = self.ixtoword[abs_cat_id].encode(
'ascii', 'ignore').decode('ascii')
cats.append(cat)
cats_list.append(cats)
#######################################################
# (2) Generate fake images
######################################################
max_num_roi = max(num_rois)
noise.data.normal_(0, 1)
fake_hmaps = netG(noise[:, :max_num_roi], bbox_maps_fwd,
bbox_maps_bwd, bbox_fmaps)
fake_hmaps = fake_hmaps.repeat(1, 1, 3, 1, 1)
for j in range(batch_size):
s_tmp = '%s/single/%s' % (save_dir, keys[j])
folder = s_tmp[:s_tmp.rfind('/')]
if not os.path.isdir(folder):
print('Make a new folder: ', folder)
mkdir_p(folder)
k = 0
# for k in range(len(fake_imgs)):
im = fake_hmaps[j][k].data.cpu().numpy()
minV = im.min()
maxV = im.max()
im = (im - minV) / (maxV - minV)
im *= 255
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
cat = cats_list[j][k]
fullpath = '{0}_{1}.png'.format(s_tmp, cat)
im.save(fullpath)
def build_models(self):
# ############################## encoders ############################# #
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
# ########### image generator and (potential) shape generator ########## #
netG = G_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netG.eval()
netShpG = None
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG = SHP_G_NET(len(self.cats_index_dict))
netShpG.apply(weights_init)
netShpG.eval()
# ################### parallization and initialization ################## #
if cfg.CUDA:
text_encoder.cuda()
image_encoder.cuda()
netG.cuda()
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG.cuda()
if len(cfg.GPU_IDS) > 1:
text_encoder = nn.DataParallel(text_encoder)
text_encoder.to(self.device)
image_encoder = nn.DataParallel(image_encoder)
image_encoder.to(self.device)
netG = nn.DataParallel(netG)
netG.to(self.device)
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG = nn.DataParallel(netShpG)
netShpG.to(self.device)
state_dict = torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
if cfg.TEST.USE_GT_BOX_SEG > 0:
state_dict = torch.load(cfg.TEST.NET_SHP_G,
map_location=lambda storage, loc: storage)
netShpG.load_state_dict(state_dict)
print('Load Shape G from: ', cfg.TEST.NET_SHP_G)
return [text_encoder, image_encoder, netG, netShpG]
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(self.data_loader, 0):
imgs, t_embeddings, filenames = data
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
batch_size = imgs[0].size(0)
noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
fake_img_list = []
for i in range(embedding_dim):
fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], filenames,
save_dir, split_dir, i, 256)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
self.save_superimages(fake_img_list, filenames,
save_dir, split_dir, 256)
def evaluate(self):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
self.num_Ds = cfg.TREE.BRANCH_NUM
self.base_num = 135
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(self.data_loader, 0):
imgs, t_embeddings, filenames, _ = data
embedding_dim = t_embeddings.size(1)
batch_size = imgs[0].size(0)
noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
crop_vbase = []
crop_base_imgs = torch.zeros(batch_size, 3, self.img_size,
self.img_size)
for step, (base_img_list) in enumerate(data[3]):
if cfg.DATASET_NAME.find('flower') != -1:
base_ix = random.randint(1, self.base_num)
base_img_name = '%s/%s.jpg' % (base_img_list,
str(base_ix))
else:
temp_base_list = os.listdir(base_img_list)
base_ix = random.randint(0, len(temp_base_list) - 1)
base_img_name = '%s/%s.jpg' % (base_img_list,
str(base_ix))
base_img = Image.open(base_img_name).convert('RGB')
crop_base = base_img.resize([self.img_size, self.img_size])
crop_base = Torchtransform(crop_base)
crop_base_imgs[step, :] = crop_base
if cfg.CUDA:
crop_vbase.append(Variable(crop_base_imgs).cuda())
else:
crop_vbase.append(Variable(crop_base_imgs))
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
for i in range(embedding_dim):
fake_imgs, fake_segs, _, _ = netG(noise,
t_embeddings[:, i, :],
crop_vbase)
self.save_singleimages(fake_imgs, fake_segs[-1],
crop_vbase[0], filenames, save_dir,
i, self.img_size)
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
#state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
# s_tmp = cfg.TRAIN.NET_G
# istart = s_tmp.rfind('_') + 1
# iend = s_tmp.rfind('.')
# iteration = int(s_tmp[istart:iend])
# s_tmp = s_tmp[:s_tmp.rfind('/')]
# save_dir = '%s/iteration%d' % (s_tmp, iteration)
save_dir = 'J:\\qimao\\Text-to-image\\results\\Plugin-v1-210K-random50'
nz = cfg.GAN.Z_DIM
n_samples = 50
# noise = Variable(torch.FloatTensor(self.batch_size, nz))
noise = Variable(torch.FloatTensor(n_samples, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(self.data_loader, 0):
imgs, t_embeddings, filenames = data
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
# batch_size = imgs[0].size(0)
# noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
fake_img_list = []
for i in range(embedding_dim):
for j in range(n_samples):
noise_j = noise[j].unsqueeze(0)
t_embeddings_i = t_embeddings[:, i, :]
#hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxstart
fake_imgs, layers_output, _, _ = netG(
noise_j, t_embeddings_i)
if len(layers_output) != len(lamdas):
print("please check lamdas length")
#hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxend
# filenames_number ='_sample_%2.2d'%(j)
# filenames_new = []
# filenames_new.append(filenames[-1]+filenames_number)
# filenames_new = tuple(filenames_new)
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], filenames, j,
save_dir, split_dir, i, 256)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
self.save_superimages(fake_img_list, filenames, save_dir,
split_dir, 256)
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load('/content/drive/My Drive/Colab Notebooks/StackGAN-v2-master/models/netG_210000.pth',
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(self.data_loader, 0):
imgs, t_embeddings, filenames = data
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
batch_size = imgs[0].size(0)
noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
fake_img_list = []
for i in range(embedding_dim):
fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], filenames,
save_dir, split_dir, i, 256)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
self.save_superimages(fake_img_list, filenames,
save_dir, split_dir, 256)
def gen_example(n_words, wordtoix, ixtoword, model_dir):
'''generate images from example sentences'''
# filepath = 'example_captions.txt'
filepath = 'caption.txt'
data_dic = {}
with open(filepath, "r") as f:
filenames = f.read().split('\n')
captions = []
cap_lens = []
for sent in filenames:
if len(sent) == 0:
continue
sent = sent.replace("\ufffd\ufffd", " ")
tokenizer = RegexpTokenizer(r'\w+')
tokens = tokenizer.tokenize(sent.lower())
if len(tokens) == 0:
print('sentence token == 0 !')
continue
rev = []
for t in tokens:
t = t.encode('ascii', 'ignore').decode('ascii')
if len(t) > 0 and t in wordtoix:
rev.append(wordtoix[t])
captions.append(rev)
cap_lens.append(len(rev))
max_len = np.max(cap_lens)
sorted_indices = np.argsort(cap_lens)[::-1]
cap_lens = np.asarray(cap_lens)
cap_lens = cap_lens[sorted_indices]
cap_array = np.zeros((len(captions), max_len), dtype='int64')
for i in range(len(captions)):
idx = sorted_indices[i]
cap = captions[idx]
c_len = len(cap)
cap_array[i, :c_len] = cap
# key = name[(name.rfind('/') + 1):]
key = 0
data_dic[key] = [cap_array, cap_lens, sorted_indices]
# algo.gen_example(data_dic)
text_encoder = RNN_ENCODER(n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
netG = G_NET()
netG.apply(weights_init)
# netG.cuda()
netG.eval()
state_dict = torch.load(model_dir,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
save_dir = 'results'
mkdir_p(save_dir)
for key in data_dic:
captions, cap_lens, sorted_indices = data_dic[key]
batch_size = captions.shape[0]
nz = cfg.GAN.Z_DIM
with torch.no_grad():
captions = Variable(torch.from_numpy(captions))
cap_lens = Variable(torch.from_numpy(cap_lens))
# captions = captions.cuda()
# cap_lens = cap_lens.cuda()
for i in range(image_per_caption): # 16
with torch.no_grad():
noise = Variable(torch.FloatTensor(batch_size, nz))
# noise = noise.cuda()
# (1) Extract text embeddings
hidden = text_encoder.init_hidden(batch_size)
words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
mask = (captions == 0)
# (2) Generate fake images
noise.data.normal_(0, 1)
fake_imgs, attention_maps, _, _ = netG(noise, sent_emb, words_embs,
mask)
cap_lens_np = cap_lens.data.numpy()
for j in range(batch_size):
save_name = '%s/%d_%d' % (save_dir, i, sorted_indices[j])
for k in range(len(fake_imgs)):
im = fake_imgs[k][j].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
# print('im', im.shape)
im = np.transpose(im, (1, 2, 0))
# print('im', im.shape)
im = Image.fromarray(im)
fullpath = '%s_g%d.png' % (save_name, k)
im.save(fullpath)
for k in range(len(attention_maps)):
if len(fake_imgs) > 1:
im = fake_imgs[k + 1]
else:
im = fake_imgs[0]
attn_maps = attention_maps[k]
att_sze = attn_maps.size(2)
img_set, sentences = \
build_super_images2(im[j].unsqueeze(0),
captions[j].unsqueeze(0),
[cap_lens_np[j]], ixtoword,
[attn_maps[j]], att_sze)
if img_set is not None:
im = Image.fromarray(img_set)
fullpath = '%s_a%d_attention.png' % (save_name, k)
im.save(fullpath)
def sample_images(self):
sample_size = 24
save_dir = '../sample_images/'
save_final = '../sample_finals/'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
if not os.path.exists(save_final):
os.makedirs(save_final)
random.seed(datetime.now())
depth = cfg.TEST_DEPTH
res = 32 * 2**depth
if cfg.TRAIN.NET_G == '':
print('Error: the path for model not found!')
else:
# Build and load the generator
netG = G_NET(depth)
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
model_dict = netG.state_dict()
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
state_dict = {
k: v
for k, v in state_dict.items() if k in model_dict
}
model_dict.update(state_dict)
netG.load_state_dict(model_dict)
print('Load ', cfg.TRAIN.NET_G)
# Uncomment this to print Generator layers
# print(netG)
nz = cfg.GAN.Z_DIM
noise = torch.FloatTensor(1, nz)
# noise.data.normal_(0, 1)
# noise = noise.repeat(1, 1)
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
netG.eval()
for i in tqdm(range(sample_size)):
noise.data.normal_(0, 1)
bg_code = torch.zeros([1, cfg.FINE_GRAINED_CATEGORIES]).cuda()
p_code = torch.zeros([1, cfg.SUPER_CATEGORIES]).cuda()
c_code = torch.zeros([1, cfg.FINE_GRAINED_CATEGORIES]).cuda()
b = random.randint(0, cfg.FINE_GRAINED_CATEGORIES - 1)
p = random.randint(0, cfg.SUPER_CATEGORIES - 1)
c = random.randint(0, cfg.FINE_GRAINED_CATEGORIES - 1)
bg_code[0][b] = 1
p_code[0][p] = 1
c_code[0][c] = 1
fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(
noise, c_code, 1, p_code,
bg_code) # Forward pass through the generator
self.save_image(fake_imgs[3 * depth + 0][0], save_dir,
'%d_bg' % i)
self.save_image(fake_imgs[3 * depth + 1][0], save_dir,
'%d_pf' % i)
self.save_image(fake_imgs[3 * depth + 2][0], save_dir,
'%d_cf' % i)
self.save_image(fake_imgs[3 * depth + 2][0], save_final,
'%d' % i)
# self.save_image(fg_imgs[2 * depth + 0][0], save_dir, 'parent_foreground')
# self.save_image(fg_imgs[2 * depth + 1][0], save_dir, 'child_foreground')
self.save_image(mk_imgs[2 * depth + 0][0], save_dir,
'%d_pmk' % i)
self.save_image(mk_imgs[2 * depth + 1][0], save_dir,
'%d_cmk' % i)
def evaluate_finegan(self):
self.save_dir = os.path.join(cfg.SAVE_DIR, 'images')
mkdir_p(self.save_dir)
random.seed(datetime.now())
depth = cfg.TEST_DEPTH
res = 32 * 2**depth
if cfg.TRAIN.NET_G == '':
print('Error: the path for model not found!')
else:
# Build and load the generator
netG = G_NET(depth)
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
model_dict = netG.state_dict()
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
state_dict = {
k: v
for k, v in state_dict.items() if k in model_dict
}
model_dict.update(state_dict)
netG.load_state_dict(model_dict)
print('Load ', cfg.TRAIN.NET_G)
# Uncomment this to print Generator layers
# print(netG)
nrow = 6
ncol = 4
z_std = 0.1
p_vs_c = False
reprod = False
if not reprod:
torch.manual_seed(random.randint(-9999, 9999))
bg_li = []
pf_li = []
cf_li = []
pk_li = []
ck_li = []
pfg_li = []
cfg_li = []
pfgmk_li = []
cfgmk_li = []
b = random.randint(0, cfg.FINE_GRAINED_CATEGORIES - 1)
nz = cfg.GAN.Z_DIM
noise = torch.FloatTensor(1, nz)
noise.data.normal_(0, z_std)
# noise = noise.repeat(self.batch_size, 1)
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
netG.eval()
c_li = np.random.randint(0,
cfg.FINE_GRAINED_CATEGORIES - 1,
size=nrow)
p_li = np.random.randint(0, cfg.SUPER_CATEGORIES - 1, size=nrow)
for k in range(ncol):
p = p_li[k]
# p = random.randint(0, cfg.SUPER_CATEGORIES-1)
for i in range(nrow):
bg_code = torch.zeros(
[self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
p_code = torch.zeros(
[self.batch_size, cfg.SUPER_CATEGORIES])
c_code = torch.zeros(
[self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
c = c_li[i]
for j in range(self.batch_size):
bg_code[j][b] = 1
p_code[j][p] = 1
c_code[j][c] = 1
fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(
noise, c_code, None, p_code,
bg_code) # Forward pass through the generator
bg_li.append(fake_imgs[3 * depth][0])
pf_li.append(fake_imgs[3 * depth + 1][0])
cf_li.append(fake_imgs[3 * depth + 2][0])
pk_li.append(mk_imgs[2 * depth][0])
ck_li.append(mk_imgs[2 * depth + 1][0])
pfg_li.append(fg_imgs[2 * depth][0])
cfg_li.append(fg_imgs[2 * depth + 1][0])
pfgmk_li.append(fgmk_imgs[2 * depth][0])
cfgmk_li.append(fgmk_imgs[2 * depth + 1][0])
save_image(bg_li, self.save_dir, 'background_pvc', nrow, res)
save_image(pf_li, self.save_dir, 'parent_final_pvc', nrow, res)
save_image(cf_li, self.save_dir, 'child_final_pvc', nrow, res)
save_image(pfg_li, self.save_dir, 'parent_foreground_pvc', nrow,
res)
save_image(cfg_li, self.save_dir, 'child_foreground_pvc', nrow,
res)
save_image(pk_li, self.save_dir, 'parent_mask_pvc', nrow, res)
save_image(ck_li, self.save_dir, 'child_mask_pvc', nrow, res)
save_image(pfgmk_li, self.save_dir, 'parent_foreground_masked_pvc',
nrow, res)
save_image(cfgmk_li, self.save_dir, 'child_foreground_masked_pvc',
nrow, res)
bg_li = []
pf_li = []
cf_li = []
pk_li = []
ck_li = []
pfg_li = []
cfg_li = []
pfgmk_li = []
cfgmk_li = []
for _ in range(ncol):
noise.data.normal_(0, z_std)
for i in range(nrow):
bg_code = torch.zeros(
[self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
p_code = torch.zeros(
[self.batch_size, cfg.SUPER_CATEGORIES])
c_code = torch.zeros(
[self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
c = c_li[i]
p = p_li[i]
for j in range(self.batch_size):
bg_code[j][b] = 1
p_code[j][p] = 1
c_code[j][c] = 1
fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(
noise, c_code, None, p_code,
bg_code) # Forward pass through the generator
bg_li.append(fake_imgs[3 * depth][0])
pf_li.append(fake_imgs[3 * depth + 1][0])
cf_li.append(fake_imgs[3 * depth + 2][0])
pk_li.append(mk_imgs[2 * depth][0])
ck_li.append(mk_imgs[2 * depth + 1][0])
pfg_li.append(fg_imgs[2 * depth][0])
cfg_li.append(fg_imgs[2 * depth + 1][0])
pfgmk_li.append(fgmk_imgs[2 * depth][0])
cfgmk_li.append(fgmk_imgs[2 * depth + 1][0])
save_image(bg_li, self.save_dir, 'background_zvpc', nrow, res)
save_image(pf_li, self.save_dir, 'parent_final_zvpc', nrow, res)
save_image(cf_li, self.save_dir, 'child_final_zvpc', nrow, res)
save_image(pfg_li, self.save_dir, 'parent_foreground_zvpc', nrow,
res)
save_image(cfg_li, self.save_dir, 'child_foreground_zvpc', nrow,
res)
save_image(pk_li, self.save_dir, 'parent_mask_zvpc', nrow, res)
save_image(ck_li, self.save_dir, 'child_mask_zvpc', nrow, res)
save_image(pfgmk_li, self.save_dir,
'parent_foreground_masked_zvpc', nrow, res)
save_image(cfgmk_li, self.save_dir, 'child_foreground_masked_zvpc',
nrow, res)
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netG = G_NET(len(self.cats_index_dict))
netsPatD, netsShpD = [], []
if cfg.TREE.BRANCH_NUM > 0:
netsPatD.append(PAT_D_NET64())
netsShpD.append(SHP_D_NET64(len(self.cats_index_dict)))
if cfg.TREE.BRANCH_NUM > 1:
netsPatD.append(PAT_D_NET128())
netsShpD.append(SHP_D_NET128(len(self.cats_index_dict)))
if cfg.TREE.BRANCH_NUM > 2:
netsPatD.append(PAT_D_NET256())
netsShpD.append(SHP_D_NET256(len(self.cats_index_dict)))
netObjSSD = OBJ_SS_D_NET(len(self.cats_index_dict))
netObjLSD = OBJ_LS_D_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netObjSSD.apply(weights_init)
netObjLSD.apply(weights_init)
for i in range(len(netsPatD)):
netsPatD[i].apply(weights_init)
netsShpD[i].apply(weights_init)
print('# of netsPatD', len(netsPatD))
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
netObjSSD.cuda()
netObjLSD.cuda()
for i in range(len(netsPatD)):
netsPatD[i].cuda()
netsShpD[i].cuda()
if len(cfg.GPU_IDS) > 1:
text_encoder = nn.DataParallel(text_encoder)
text_encoder.to(self.device)
image_encoder = nn.DataParallel(image_encoder)
image_encoder.to(self.device)
netG = nn.DataParallel(netG)
netG.to(self.device)
netObjSSD = nn.DataParallel(netObjSSD)
netObjSSD.to(self.device)
netObjLSD = nn.DataParallel(netObjLSD)
netObjLSD.to(self.device)
for i in range(len(netsPatD)):
netsPatD[i] = nn.DataParallel(netsPatD[i])
netsPatD[i].to(self.device)
netsShpD[i] = nn.DataParallel(netsShpD[i])
netsShpD[i].to(self.device)
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
Gname = cfg.TRAIN.NET_G
for i in range(len(netsPatD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netPatD%d.pth' % (s_tmp, i)
print('Load PatD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsPatD[i].load_state_dict(state_dict)
Dname = '%s/netShpD%d.pth' % (s_tmp, i)
print('Load ShpD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsShpD[i].load_state_dict(state_dict)
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netObjSSD.pth' % (s_tmp)
print('Load ObjSSD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netObjSSD.load_state_dict(state_dict)
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netObjLSD.pth' % (s_tmp)
print('Load ObjLSD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netObjLSD.load_state_dict(state_dict)
return [
text_encoder, image_encoder, netG, netsPatD, netsShpD, netObjSSD,
netObjLSD, epoch
]
def sample(self, split_dir, num_samples=25, draw_bbox=False):
from PIL import Image, ImageDraw, ImageFont
import cPickle as pickle
import torchvision
import torchvision.utils as vutils
if cfg.TRAIN.NET_G == '':
print('Error: the path for model NET_G is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
text_encoder = RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
batch_size = cfg.TRAIN.BATCH_SIZE
nz = cfg.GAN.Z_DIM
model_dir = cfg.TRAIN.NET_G
state_dict = torch.load(model_dir, map_location=lambda storage, loc: storage)
# state_dict = torch.load(cfg.TRAIN.NET_G)
netG = G_NET()
print('Load G from: ', model_dir)
netG.apply(weights_init)
netG.load_state_dict(state_dict["netG"])
netG.cuda()
netG.eval()
# the path to save generated images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s_%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
#######################################
noise = Variable(torch.FloatTensor(9, nz))
imsize = 256
for step, data in enumerate(self.data_loader, 0):
if step >= num_samples:
break
imgs, captions, cap_lens, class_ids, keys, transformation_matrices, label_one_hot, bbox = \
prepare_data(data, eval=True)
transf_matrices_inv = transformation_matrices[1][0].unsqueeze(0)
label_one_hot = label_one_hot[0].unsqueeze(0)
img = imgs[-1][0]
val_image = img.view(1, 3, imsize, imsize)
hidden = text_encoder.init_hidden(batch_size)
# words_embs: batch_size x nef x seq_len
# sent_emb: batch_size x nef
words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
words_embs, sent_emb = words_embs[0].unsqueeze(0).detach(), sent_emb[0].unsqueeze(0).detach()
words_embs = words_embs.repeat(9, 1, 1)
sent_emb = sent_emb.repeat(9, 1)
mask = (captions == 0)
mask = mask[0].unsqueeze(0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
mask = mask.repeat(9, 1)
transf_matrices_inv = transf_matrices_inv.repeat(9, 1, 1, 1)
label_one_hot = label_one_hot.repeat(9, 1, 1)
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
inputs = (noise, sent_emb, words_embs, mask, transf_matrices_inv, label_one_hot)
with torch.no_grad():
fake_imgs, _, mu, logvar = nn.parallel.data_parallel(netG, inputs, self.gpus)
data_img = torch.FloatTensor(10, 3, imsize, imsize).fill_(0)
data_img[0] = val_image
data_img[1:10] = fake_imgs[-1]
if draw_bbox:
for idx in range(3):
x, y, w, h = tuple([int(imsize*x) for x in bbox[0, idx]])
w = imsize-1 if w > imsize-1 else w
h = imsize-1 if h > imsize-1 else h
if x <= -1:
break
data_img[:10, :, y, x:x + w] = 1
data_img[:10, :, y:y + h, x] = 1
data_img[:10, :, y+h, x:x + w] = 1
data_img[:10, :, y:y + h, x + w] = 1
# get caption
cap = captions[0].data.cpu().numpy()
sentence = ""
for j in range(len(cap)):
if cap[j] == 0:
break
word = self.ixtoword[cap[j]].encode('ascii', 'ignore').decode('ascii')
sentence += word + " "
sentence = sentence[:-1]
vutils.save_image(data_img, '{}/{}_{}.png'.format(save_dir, sentence, step), normalize=True, nrow=10)
print("Saved {} files to {}".format(step, save_dir))
def evaluate(self, split_dir):
sample_num_per_image = 1
self.inception_model = INCEPTION_V3().eval()
if cfg.CUDA:
self.inception_model = self.inception_model.cuda()
# self.inception_model = torch.cuda.parallel.DistributedDataParallel(self.inception_model, device_ids=self.gpus, output_device=self.gpus[0])
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise_raw = (torch.FloatTensor(self.batch_size, nz).requires_grad_())
if cfg.CUDA:
netG.cuda()
noise_raw = noise_raw.cuda()
else:
pass
predictions_g = []
predictions_r = []
activate_g = []
activate_r = []
IS_mean_add = 0
FID_add = 0
NLPP_mean_add = 0
metric_cnt = 0
count = 0
# switch to evaluate mode
netG.eval()
loader_bar = tqdm.tqdm(self.data_loader)
for step, data in enumerate(loader_bar, 0):
count += 1
imgs, t_embeddings, filenames = data
# print(t_embeddings.shape)
if cfg.CUDA:
if isinstance(t_embeddings, list):
t_embeddings = [emb.requires_grad_(False).float().cuda() for emb in t_embeddings]
else:
t_embeddings = t_embeddings.requires_grad_(False).float().cuda()
if isinstance(imgs, list):
imgs = [img.requires_grad_(False).float().cuda() for img in imgs]
else:
imgs = imgs.requires_grad_(False).float().cuda()
else:
if isinstance(t_embeddings, list):
t_embeddings = [emb.requires_grad_(False).float() for emb in t_embeddings]
else:
t_embeddings = t_embeddings.requires_grad_(False).float()
if isinstance(imgs, list):
imgs = [img.requires_grad_(False).float() for img in imgs]
else:
imgs = imgs.requires_grad_(False).float()
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
batch_size = imgs[0].size(0)
noise = noise_raw[:batch_size]
noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
# trunc
# noise[noise<-cfg.TEST.TRUNC] = -cfg.TEST.TRUNC
# noise[noise>cfg.TEST.TRUNC] = cfg.TEST.TRUNC
fake_img_list = []
for i in range(embedding_dim):
for sample_idx in range(sample_num_per_image):
noise.data.normal_(0, 1)
with torch.autograd.no_grad():
fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
real_imgs = imgs
# pred_g, pool3_g = self.inception_model(fake_imgs[-1].detach())
# pred_r, pool3_r = self.inception_model(real_imgs[-1].detach())
# predictions_g.append(pred_g.data.cpu().numpy())
# predictions_r.append(pred_r.data.cpu().numpy())
# activate_g.append(pool3_g.data.cpu().numpy())
# activate_r.append(pool3_r.data.cpu().numpy())
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], filenames,
save_dir, split_dir, i, 256, sample_idx)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
self.save_superimages(fake_img_list, filenames,
save_dir, split_dir, 256)
print("len pred: {}".format(len(predictions_g)))
return [{'mu':0, 'sigma':0},{'mu':0, 'sigma':0}] # to save time
predictions_g = np.concatenate(predictions_g, 0)
predictions_r = np.concatenate(predictions_r, 0)
activate_g = np.concatenate(activate_g, 0)
activate_r = np.concatenate(activate_r, 0)
mean, std = compute_inception_score(predictions_g, 10)
# print('mean:', mean, 'std', std)
self.summary_writer.add_scalar('Inception_mean', mean, count)
#
mean_nlpp, std_nlpp = \
negative_log_posterior_probability(predictions_g, 10)
self.summary_writer.add_scalar('NLPP_mean', mean_nlpp, count)
# FID
fid, fid_data = compute_frethet_distance(activate_g, activate_r)
self.summary_writer.add_scalar("FID", fid, count)
IS_mean_add += mean
FID_add += fid
NLPP_mean_add += mean_nlpp
metric_cnt += 1
IS_mean, FID_mean, NLPP_mean = IS_mean_add/metric_cnt, FID_add/metric_cnt, NLPP_mean_add/metric_cnt
print("total, IS mean:{}, FID mean:{}, NLPP mean:{}".format(IS_mean, FID_mean, NLPP_mean))
return fid_data
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = 1 #int(s_tmp[istart:iend])
s_tmp = "firstEvaluation" #s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(self.batch_size, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(self.data_loader, 0):
# print(80*'@')
# print(data)
imgs, t_embeddings, rec_id, im_id, _ = data
rec_id = [t.tostring().decode('UTF-8') for t in rec_id.numpy()]
im_id = [t.tostring().decode('UTF-8') for t in im_id.numpy()]
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
batch_size = imgs[0].size(0)
noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
fake_img_list = []
#for i in range(embedding_dim):
fake_imgs, _, _ = netG(
noise,
t_embeddings) #licht -this wat t_rmbeddings[:, i, :])
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], rec_id + '_' + im_id,
save_dir, split_dir, 0, 256)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
#self.save_superimages(fake_img_list, [rec_id[i]+'_'+im_id[i] for i in range(len(rec_id))],
# save_dir, split_dir, 256)
save_img_results(imgs, fake_imgs, 3, 1, save_dir,
self.summary_writer, rec_id, im_id)
def evaluate(self, split_dir):
inception_model = INCEPTION_V3()
# fid_model = FID_INCEPTION()
if cfg.CUDA:
inception_model.cuda()
# fid_model.cuda()
inception_model.eval()
# fid_model.eval()
if cfg.TRAIN.NET_G == '':
print('Error: the path for models is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
# print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
# the path to save generated images
# s_tmp = cfg.TRAIN.NET_G
# istart = s_tmp.rfind('_') + 1
# iend = s_tmp.rfind('.')
# iteration = int(s_tmp[istart:iend])
# s_tmp = s_tmp[:s_tmp.rfind('/')]
# save_dir = '%s/iteration%d' % (s_tmp, iteration)
# save_dir = 'C:\\Users\\alper\\PycharmProjects\\MSGAN\\StackGAN++-Mode-Seeking\\results'
save_dir = "D:\\results"
nz = cfg.GAN.Z_DIM
n_samples = 50
# noise = Variable(torch.FloatTensor(self.batch_size, nz))
noise = Variable(torch.FloatTensor(n_samples, nz))
if cfg.CUDA:
netG.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
for step, data in enumerate(tqdm(self.data_loader)):
# if step == 8:
# break
imgs, t_embeddings, filenames = data
if cfg.CUDA:
t_embeddings = Variable(t_embeddings).cuda()
else:
t_embeddings = Variable(t_embeddings)
# print(t_embeddings[:, 0, :], t_embeddings.size(1))
embedding_dim = t_embeddings.size(1)
# batch_size = imgs[0].size(0)
# noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
fake_img_list = []
inception_score_list = []
fid_list = []
score_list = []
predictions = []
fids = []
for i in range(embedding_dim):
inception_score_list.append([])
fid_list.append([])
score_list.append([])
emb_imgs = []
for j in range(n_samples):
noise_j = noise[j].unsqueeze(0)
t_embeddings_i = t_embeddings[:, i, :]
fake_imgs, _, _ = netG(noise_j, t_embeddings_i)
# filenames_number ='_sample_%2.2d'%(j)
# filenames_new = []
# filenames_new.append(filenames[-1]+filenames_number)
# filenames_new = tuple(filenames_new)
# for selecting reasonable images
pred = inception_model(fake_imgs[-1].detach())
pred = pred.data.cpu().numpy()
predictions.append(pred)
bird_indices = [
7, 8, 9, 10, 11, 13, 15, 16, 17, 18, 19, 21, 23,
81, 84, 85, 86, 88, 90, 91, 93, 94, 95, 96, 97, 99,
129, 130, 133, 134, 135, 138, 141, 142, 143, 144,
146, 517
]
score = np.max(pred[0, bird_indices])
score_list[i].append((j, score))
emb_imgs.append(fake_imgs[2].data.cpu())
if cfg.TEST.B_EXAMPLE:
# fake_img_list.append(fake_imgs[0].data.cpu())
# fake_img_list.append(fake_imgs[1].data.cpu())
fake_img_list.append(fake_imgs[2].data.cpu())
else:
self.save_singleimages(fake_imgs[-1], filenames, j,
save_dir, split_dir, i, 256)
# self.save_singleimages(fake_imgs[-2], filenames,
# save_dir, split_dir, i, 128)
# self.save_singleimages(fake_imgs[-3], filenames,
# save_dir, split_dir, i, 64)
# break
score_list[i] = sorted(score_list[i],
key=lambda x: x[1],
reverse=True)[:5]
# for FID score
# ffi = [i[0].numpy() for i in emb_imgs]
fake_filtered_images = [
fake_img_list[i][0].numpy()
for i in range(len(fake_img_list))
]
img_dir = os.path.join(cfg.DATA_DIR, "CUB_200_2011",
"images",
filenames[0].split("/")[0])
img_files = [
os.path.join(img_dir, i) for i in os.listdir(img_dir)
]
# act_real = get_activations(img_files, fid_model)
# mu_real, sigma_real = get_fid_stats(act_real)
# print("mu_real: {}, sigma_real: {}".format(mu_real, sigma_real))
np_imgs = np.array(fake_filtered_images)
# print(np_imgs.shape)
# # print(type(np_imgs[0]))
# act_fake = get_activations(np_imgs, fid_model, img=True)
# mu_fake, sigma_fake = get_fid_stats(act_fake)
# fid_score = frechet_distance(mu_real, sigma_real, mu_fake, sigma_fake)
# fids.append(fid_score)
# print("mu_fake: {}, sigma_fake: {}".format(mu_fake, sigma_fake))
# print(inception_score_list)
# # calculate inception score
# predictions = np.concatenate(predictions, 0)
# mean, std = compute_inception_score(predictions, 10)
# mean_nlpp, std_nlpp = \
# negative_log_posterior_probability(predictions, 10)
# inception_score_list.append((mean, std, mean_nlpp, std_nlpp))
# # for FID score
# fake_filtered_images = [fake_img_list[i*n_samples + k[0]][0].numpy() for i, j in enumerate(score_list) for k in j]
# # fake_filtered_images = [fake_img_list[i][0].numpy() for i in range(len(fake_img_list))]
# img_dir = os.path.join(cfg.DATA_DIR, "CUB_200_2011", "images", filenames[0].split("/")[0])
# img_files = [os.path.join(img_dir, i) for i in os.listdir(img_dir)]
#
# act_real = get_activations(img_files, fid_model)
# mu_real, sigma_real = get_fid_stats(act_real)
# # print("mu_real: {}, sigma_real: {}".format(mu_real, sigma_real))
#
# np_imgs = np.array(fake_filtered_images)
# # print(np_imgs.shape)
#
# # print(type(np_imgs[0]))
# act_fake = get_activations(np_imgs, fid_model, img=True)
# mu_fake, sigma_fake = get_fid_stats(act_fake)
# # print("mu_fake: {}, sigma_fake: {}".format(mu_fake, sigma_fake))
#
# # fid_score = frechet_distance(mu_real, sigma_real, mu_fake, sigma_fake)
# fid_score = np.mean(fids)
# fid_list.append(fid_score)
# stats = 'step: {}, FID: {}, inception_score: {}, nlpp: {}\n'.format(step, fid_score, (mean, std), (mean_nlpp, std_nlpp))
# with open("results\\stats.txt", "a+") as f:
# f.write(stats)
# print(stats)
if cfg.TEST.B_EXAMPLE:
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 64)
# self.save_superimages(fake_img_list, filenames,
# save_dir, split_dir, 128)
if cfg.TEST.FILTER:
images_to_save = [
fake_img_list[i * n_samples + k[0]]
for i, j in enumerate(score_list) for k in j
]
else:
images_to_save = fake_img_list
self.save_superimages(images_to_save, filenames, save_dir,
split_dir, 256)
def evaluate(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
netE = load_embedding_model(self.data_loader.dataset.dictionary)
print(netE)
nz = cfg.GAN.Z_DIM
sample_size = cfg.TEST.NUM_IMAGES
noise = Variable(torch.FloatTensor(sample_size, nz))
if cfg.CUDA:
netG.cuda()
netE.cuda()
noise = noise.cuda()
# switch to evaluate mode
netG.eval()
count = 0
output_dir = os.path.join(cfg.OUTPUT_DIR, cfg.EXPERIMENT_NAME)
for step, data in enumerate(
tqdm(self.data_loader, desc='evaluate'), 0):
imgs, txt_ids, txts = data
if cfg.CUDA:
txt_ids = Variable(txt_ids).cuda()
else:
txt_ids = Variable(txt_ids)
txts_embeddings = netE(txt_ids)
batch_size = imgs[0].size(0)
imgs64, imgs128, imgs256 = [], [], []
for i in range(0, batch_size):
noise.data.normal_(0, 1)
txt_embedding = txts_embeddings[i].repeat(sample_size, 1)
fake_imgs, _, _ = netG(noise, txt_embedding)
imgs64.append(normalize_(fake_imgs[0]))
imgs128.append(normalize_(fake_imgs[1]))
imgs256.append(normalize_(fake_imgs[2]))
save_images_with_text(imgs64, imgs128, imgs256, imgs, txts,
batch_size, cfg.TEXT.MAX_LEN, count,
output_dir)
count = count + batch_size + 1
cfg_from_file('./stackGAN_code/cfg/eval_birds.yml')
save_dir = './display'
txt_dir = './embeddings/txt_embedding.t7'
manualSeed = random.randint(1, 120)
random.seed(manualSeed)
torch.manual_seed(manualSeed)
torch.cuda.manual_seed_all(manualSeed)
gpus = [0]
num_gpus = 1
torch.cuda.set_device(gpus[0])
cudnn.benchmark = True
batch_size = 2
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz))
netG.cuda()
netG.eval()
noise = noise.cuda()
t_embeddings = load_lua(txt_dir)
t_embeddings = t_embeddings.view(-1,1024)
def evaluate(self,
split_dir,
n_samples=4,
extractor='googlenet',
save_dir=None):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
# Build and load the generator
if split_dir == 'test':
split_dir = 'valid'
netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
mapper = EXTRACTOR_MAPPING[extractor]()
mapper = torch.nn.DataParallel(mapper, device_ids=self.gpus)
set_parameter_requires_grad(netG, False)
set_parameter_requires_grad(mapper, False)
print(netG)
# state_dict = torch.load(cfg.TRAIN.NET_G)
state_dict = \
torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
if save_dir is None:
# the path to save generated images
s_tmp = cfg.TRAIN.NET_G
istart = s_tmp.rfind('_') + 1
iend = s_tmp.rfind('.')
iteration = int(s_tmp[istart:iend])
s_tmp = s_tmp[:s_tmp.rfind('/')]
save_dir = '%s/iteration%d' % (s_tmp, iteration)
nz = cfg.GAN.Z_DIM
if cfg.CUDA:
netG.cuda()
mapper.cuda()
# switch to evaluate mode
netG.eval()
mapper.eval()
synthetic_ds = SyntheticDataset(save_dir)
for class_embeddings, synthetic_id in self.data_loader.dataset.embeddings_by_class(
):
if cfg.CUDA:
class_embeddings = class_embeddings.cuda()
class_embeddings = class_embeddings.mean(
dim=1) # mean of 10 captions per image
for i in range(class_embeddings.size(0)):
image_embeddings = class_embeddings[i].repeat(n_samples, 1)
noise = torch.randn(n_samples, nz)
if cfg.CUDA:
noise = noise.cuda()
imgs, _, _ = netG(noise, image_embeddings)
imgs = imgs[-1]
samples = mapper(imgs)
synthetic_ds.save_pairs(samples, synthetic_id)
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
epoch = 0
if self.resume:
checkpoint_list = sorted([ckpt for ckpt in glob.glob(self.model_dir + "/" + '*.pth')])
latest_checkpoint = checkpoint_list[-1]
state_dict = torch.load(latest_checkpoint, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict["netG"])
for i in range(len(netsD)):
netsD[i].load_state_dict(state_dict["netD"][i])
epoch = int(latest_checkpoint[-8:-4]) + 1
print("Resuming training from checkpoint {} at epoch {}.".format(latest_checkpoint, epoch))
#
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def load_network(gpus:list, distributed:bool):
netG = G_NET()
netG.apply(weights_init)
if distributed:
netG = netG.cuda()
netG = torch.nn.parallel.DistributedDataParallel(netG, device_ids=gpus, output_device=gpus[0], broadcast_buffers=True)
else:
if cfg.CUDA:
netG = netG.cuda()
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
if cfg.TREE.BRANCH_NUM > 3:
netsD.append(D_NET512())
if cfg.TREE.BRANCH_NUM > 4:
netsD.append(D_NET1024())
# TODO: if cfg.TREE.BRANCH_NUM > 5:
# netsD_module = nn.ModuleList(netsD)
# netsD_module.apply(weights_init)
# netsD_module = torch.nn.parallel.DistributedDataParallel(netsD_module.cuda(), device_ids=gpus, output_device=gpus[0])
for i in range(len(netsD)):
netsD[i].apply(weights_init)
if distributed:
netsD[i] = torch.nn.parallel.DistributedDataParallel(netsD[i].cuda(), device_ids=gpus, output_device=gpus[0], broadcast_buffers=True
# , process_group=pg_Ds[i]
)
else:
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
print(netsD[i])
print('# of netsD', len(netsD))
count = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i), map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
inception_model = INCEPTION_V3()
if not distributed:
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
inception_model = torch.nn.DataParallel(inception_model, device_ids=gpus)
else:
inception_model = torch.nn.parallel.DistributedDataParallel(inception_model.cuda(), device_ids=gpus, output_device=gpus[0])
pass
# inception_model = inception_model.cpu() #to(torch.device("cuda:{}".format(gpus[0])))
inception_model.eval()
print("model device, G:{}, D:{}, incep:{}".format(netG.device_ids, netsD[0].device_ids, inception_model.device_ids))
return netG, netsD, len(netsD), inception_model, count
