| 106 | |
| 107 | |
| 108 | def visualize(cfg, |
| 109 | vis_file_list=None, |
| 110 | use_gpu=False, |
| 111 | vis_dir="visual_predict", |
| 112 | ckpt_dir=None, |
| 113 | log_writer=None, |
| 114 | local_test=False, |
| 115 | **kwargs): |
| 116 | if vis_file_list is None: |
| 117 | vis_file_list = cfg.DATASET.TEST_FILE_LIST |
| 118 | dataset = SegDataset( |
| 119 | file_list=vis_file_list, |
| 120 | mode=ModelPhase.VISUAL, |
| 121 | data_dir=cfg.DATASET.DATA_DIR) |
| 122 | |
| 123 | startup_prog = fluid.Program() |
| 124 | test_prog = fluid.Program() |
| 125 | pred, logit = build_model(test_prog, startup_prog, phase=ModelPhase.VISUAL) |
| 126 | # Clone forward graph |
| 127 | test_prog = test_prog.clone(for_test=True) |
| 128 | |
| 129 | # Generator full colormap for maximum 256 classes |
| 130 | color_map = get_color_map_list(256) |
| 131 | |
| 132 | # Get device environment |
| 133 | place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() |
| 134 | exe = fluid.Executor(place) |
| 135 | exe.run(startup_prog) |
| 136 | |
| 137 | ckpt_dir = cfg.TEST.TEST_MODEL if not ckpt_dir else ckpt_dir |
| 138 | |
| 139 | fluid.io.load_params(exe, ckpt_dir, main_program=test_prog) |
| 140 | |
| 141 | save_dir = os.path.join('visual', vis_dir) |
| 142 | makedirs(save_dir) |
| 143 | |
| 144 | fetch_list = [pred.name] |
| 145 | test_reader = dataset.batch(dataset.generator, batch_size=1, is_test=True) |
| 146 | img_cnt = 0 |
| 147 | for imgs, grts, img_names, valid_shapes, org_shapes in test_reader: |
| 148 | pred_shape = (imgs.shape[2], imgs.shape[3]) |
| 149 | pred, = exe.run( |
| 150 | program=test_prog, |
| 151 | feed={'image': imgs}, |
| 152 | fetch_list=fetch_list, |
| 153 | return_numpy=True) |
| 154 | |
| 155 | num_imgs = pred.shape[0] |
| 156 | # TODO: use multi-thread to write images |
| 157 | for i in range(num_imgs): |
| 158 | # Add more comments |
| 159 | res_map = np.squeeze(pred[i, :, :, :]).astype(np.uint8) |
| 160 | img_name = img_names[i] |
| 161 | res_shape = (res_map.shape[0], res_map.shape[1]) |
| 162 | if res_shape[0] != pred_shape[0] or res_shape[1] != pred_shape[1]: |
| 163 | res_map = cv2.resize( |
| 164 | res_map, pred_shape, interpolation=cv2.INTER_NEAREST) |
| 165 | valid_shape = (valid_shapes[i, 0], valid_shapes[i, 1]) |