if args.saves_prob:
# Save probability tensors
prob_outdir = os.path.join(base_outdir, "prob")
mkdir_if_not_exist(prob_outdir)
prob_outfn = os.path.join(prob_outdir, path.split('/')[-1].replace('png', 'npy'))
np.save(prob_outfn, outputs[0].data.cpu().numpy())
# Save predicted pixel labels(pngs)
pred = outputs[0, :args.n_class - 1].data.max(0)[1].cpu().numpy()
img = Image.fromarray(np.uint8(pred))
img = img.resize(test_img_shape, Image.NEAREST)
label_outdir = os.path.join(base_outdir, "label")
mkdir_if_not_exist(label_outdir)
label_fn = os.path.join(label_outdir, path.split('/')[-1])
img.save(label_fn)
# Save visualized predicted pixel labels(pngs)
vis_outdir = os.path.join(base_outdir, "vis")
mkdir_if_not_exist(vis_outdir)
vis_fn = os.path.join(vis_outdir, path.split('/')[-1])
save_colorized_lbl(img, vis_fn, args.tgt_dataset)
exec_eval(args.tgt_dataset, label_outdir)
ave_ent = total_ent / len(target_loader)
print ("average entropy: %s" % ave_ent)
with open(os.path.join(base_outdir, "ave_ent_%s.txt" % ave_ent), "w") as f:
f.write(str(ave_ent))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Refine raw boundary imgs')
parser.add_argument('segdir',
type=str,
help="Directory that contains segmentation results")
parser.add_argument('bwbddir',
type=str,
help="Directory that contains bwboundary results")
parser.add_argument('--dataset',
type=str,
default="nyu",
help="dataset (used for colorizing indexed label)")
parser.add_argument('--min_thre',
type=int,
default=500,
help='the minimum number of pixel in a region')
parser.add_argument(
'--max_thre',
type=int,
default=79333, # 79333 = 425 * 560 / 3
help='the maximum number of pixel in a region')
args = parser.parse_args()
out_seg_dir = refine_by_bwboundary(segdir=args.segdir,
bwbddir=args.bwbddir,
dataset=args.dataset,
min_thre=args.min_thre,
max_thre=args.max_thre)
exec_eval(args.dataset, out_seg_dir)