except:
pass
gaborext = GaborWavelet()
gaborext.eval()
extractor = GaborVAE()
extractor.load_state_dict(torch.load('checkpoints/gabor_step_i3_vae_mu/model_90.pth'))
# extractor.load_state_dict(torch.load('checkpoints/gabor_kl_step_1_i3_3c/model_500.pth'))
extractor.eval()
model = GaborAE()
# print(model)
if args.start>1:
model.load_state_dict(torch.load('checkpoints2/%s/model_%d.pth' % (args.outf, args.start-1)))
# solver.load_state_dict(torch.load('checkpoints2/%s/optimizer_%d.pth' % (args.outf, args.start-1)))
else:
init_weights(model, init_type='kaiming')
solver = optim.Adam(model.parameters(), lr=args.lr, betas=(0.5,0.9))#, weight_decay=1e-5)
scheduler = get_scheduler(solver, args)
train_writer = tensorboardX.SummaryWriter("./logs2/%s/"%args.outf)
gaborext.cuda()
extractor.cuda()
model.cuda()
for epoch in range(args.start, args.nepoch+1):
train(epoch,args)
with torch.no_grad():
test(epoch,args)
scheduler.step()
out = (mu[0]+1)*127.5
cv2.imwrite(data[1][0].replace('data/','results/').replace('.jpg','.png'), cv2.cvtColor(out.cpu().numpy().transpose(1,2,0), cv2.COLOR_RGB2BGR))
# out = torch.clamp(inter[0]*200+200,0,400)*255//400
# cv2.imwrite('results/%s/%d-in.png'%(args.outf,i), (data[0,0].cpu().numpy()+1)*127.5)
# cv2.imwrite('results/%s/%d-out.png'%(args.outf,i), cv2.cvtColor(out[:3,:,:].cpu().numpy().transpose(1,2,0), cv2.COLOR_RGB2BGR))
batch_size=args.batchSize
# gaborext = GaborWavelet()
# gaborext.eval()
model = GaborAE()
model.eval()
try:
os.mkdir('results/%s'%args.outf)
except:
pass
# print(model)
if args.start>1:
model.load_state_dict(torch.load('checkpoints2/%s/model_%d.pth' % (args.outf, args.start)))
# solver.load_state_dict(torch.load('checkpoints/%s/optimizer_%d.pth' % (args.outf, args.start-1)))
else:
init_weights(model, init_type='xavier')
# gaborext.cuda()
model.cuda()
with torch.no_grad():
test(args)