def vis_using_Colorize(indir_list, outdir):
indir = indir_list[0]
# outdir = os.path.join(os.path.split(indir)[0], "vis_labels")
mkdir_if_not_exist(outdir)
for one_file in tqdm(os.listdir(indir)):
fullpath = os.path.join(indir, one_file)
hard_to_see_img = m.imread(fullpath)
# outputs = outputs[0, :19].data.max(0)[1]
# outputs = outputs.view(1, outputs.size()[0], outputs.size()[1])
outputs = hard_to_see_img # TODO this should be fixed
output = Colorize()(outputs)
output = np.transpose(output.cpu().numpy(), (1, 2, 0))
img = Image.fromarray(output, "RGB")
img = img.resize(hard_to_see_img.shape, Image.NEAREST)
outfn = os.path.join(outdir, one_file)
plt.savefig(outfn, transparent=True, bbox_inches='tight', pad_inches=0)
img.save(outfn)
loss.backward()
optimizer_seg.step()
optimizer_feat.step()
print "epoch is:[{}|{}],index is:[{}|{}],loss:{}".\
format(epoch,epoch_num,i,len(dataloader),loss)
win = visutils.visualize_loss(epoch, loss.cpu().detach(), env, win)
if epoch % 40 == 0:
#save model
torch.save(vgg16.state_dict(),
'%s/vgg16_%03d.pkl' % (result_directory, epoch))
torch.save(Seg.state_dict(),
'%s/Seg_%03d.pkl' % (result_directory, epoch))
#save result
input = make_image_grid(img, mean, std)
label = make_label_grid(labels.data)
label = Colorize()(label).type(torch.FloatTensor)
output = make_label_grid(torch.max(logits, dim=1)[1].data)
output = Colorize()(output).type(torch.FloatTensor)
vutils.save_image(
label.cpu().detach(),
'%s/labels_epoch_%03d.png' % (result_directory, epoch))
vutils.save_image(
output.cpu().detach(),
'%s/output_epoch_%03d.png' % (result_directory, epoch))
vutils.save_image(input.cpu().detach(),
'%s/img_epoch_%03d.png' % (result_directory, epoch))
