def visualizeSingleBatch(fp_loader_test,
opt,
exp_folder,
batches_done,
batch_size=8):
print('Loading saved model ... \n{}'.format(
'./checkpoints/{}_{}.pth'.format(exp_folder, batches_done)))
generatorTest = Generator()
generatorTest.load_state_dict(
torch.load('./checkpoints/{}_{}.pth'.format(exp_folder, batches_done)))
generatorTest = generatorTest.eval()
generatorTest.cuda()
with torch.no_grad():
# Unpack batch
mks, nds, eds, nd_to_sample, ed_to_sample = next(iter(fp_loader_test))
real_mks = Variable(mks.type(Tensor))
given_nds = Variable(nds.type(Tensor))
given_eds = eds
# Generate a batch of images
z_shape = [real_mks.shape[0], opt.latent_dim]
z = Variable(Tensor(np.random.normal(0, 1, tuple(z_shape))))
# Select random nodes
ind_fixed_nodes, bin_fixed_nodes = selectNodesTypes(
nd_to_sample, batch_size, nds)
# Generate a batch of images
given_masks = torch.tensor(real_mks)
ind_not_fixed_nodes = torch.tensor([
k for k in range(given_masks.shape[0]) if k not in ind_fixed_nodes
])
## Set non fixed masks to -1.0
given_masks[ind_not_fixed_nodes.long()] = -1.0
given_masks = given_masks.unsqueeze(1)
## Add channel to indicate given nodes
inds_masks = torch.zeros_like(given_masks)
inds_masks[ind_not_fixed_nodes.long()] = 0.0
inds_masks[ind_fixed_nodes.long()] = 1.0
given_masks = torch.cat([given_masks, inds_masks], 1)
gen_mks = generatorTest(z, None, given_masks, given_nds, given_eds)
# Generate image tensors
real_imgs_tensor = combine_images_maps(real_mks, given_nds, given_eds, \
nd_to_sample, ed_to_sample)
fake_imgs_tensor = combine_images_maps(gen_mks, given_nds, given_eds, \
nd_to_sample, ed_to_sample)
# Save images
save_image(real_imgs_tensor, "./exps/{}/{}_real.png".format(exp_folder, batches_done), \
nrow=12, normalize=False)
save_image(fake_imgs_tensor, "./exps/{}/{}_fake.png".format(exp_folder, batches_done), \
nrow=12, normalize=False)
return
x, y = np.array(junctions[j])
dwg.add(
dwg.circle(center=(float(x), float(y)),
r=3,
stroke='red',
fill='white',
stroke_width=2,
opacity=1.0))
return
# Create folder
os.makedirs(opt.exp_folder, exist_ok=True)
# Initialize generator and discriminator
generator = Generator()
generator.load_state_dict(torch.load(checkpoint), strict=False)
generator = generator.eval()
# Initialize variables
cuda = True if torch.cuda.is_available() else False
if cuda:
generator.cuda()
rooms_path = '../'
# Initialize dataset iterator
fp_dataset_test = FloorplanGraphDataset(rooms_path,
transforms.Normalize(mean=[0.5],
std=[0.5]),
target_set=target_set,
split=phase)