import torch
from parser_2 import arg_parse
import torchvision.utils as vutils
from GAN.Gen_GAN import Generator
import os
import numpy as np
from PIL import Image
device = torch.device('cuda')
#t = torch.randn(32, 100, 1, 1, device=device)
#torch.save(t, 'seed_GAN.pt')
t = torch.load('seed_GAN.pt')
args = arg_parse()
netG = Generator(args.ngpu).cuda()
#plt.figure()
#plt.subplot(1, 2, 2)
#plt.axis("off")
for i in range(53):
model_std = torch.load(
os.path.join(args.save_GAN, 'net_G_{}.pth.tar'.format(i)))
netG.load_state_dict(model_std)
img_list = []
with torch.no_grad():
fake = netG(t).detach().cpu()
img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
img_only = np.transpose(img_list[-1], (1, 2, 0)).numpy()
#print("batch predicted")
_, pred = torch.max(pred, dim=1)
pred = pred.cpu().numpy()
n = 0
dirs = []
dirs.append(name)
dirs = np.concatenate(dirs)
for i in pred:
result = Image.fromarray((i).astype(np.uint8))
result.save(args.save_img + '/' + dirs[n])
#torchvision.utils.save_image(i, 'preds/' + dirs[n])
n += 1
if __name__ == '__main__':
args = parser_2.arg_parse()
''' setup GPU '''
torch.cuda.set_device(args.gpu)
''' prepare data_loader '''
print('===> prepare data loader ...')
test_loader2 = torch.utils.data.DataLoader(data2.DATA2(args),
batch_size=args.test_batch,
num_workers=args.workers,
shuffle=False)
''' prepare mode '''
model = models.Net(args)
model_std = torch.load(os.path.join('model_best.pth.tar'))
model.load_state_dict(model_std)
model.cuda()
print("model loaded")