(args, img, gt_norm, gt_norm_mask, norm_out_list, total_iter)
| 138 | |
| 139 | # visualize during training |
| 140 | def visualize(args, img, gt_norm, gt_norm_mask, norm_out_list, total_iter): |
| 141 | B, _, H, W = gt_norm.shape |
| 142 | |
| 143 | pred_norm_list = [] |
| 144 | pred_kappa_list = [] |
| 145 | for norm_out in norm_out_list: |
| 146 | norm_out = F.interpolate(norm_out, size=[gt_norm.size(2), gt_norm.size(3)], mode='nearest') |
| 147 | pred_norm = norm_out[:, :3, :, :] # (B, 3, H, W) |
| 148 | pred_norm = pred_norm.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 3) |
| 149 | pred_norm_list.append(pred_norm) |
| 150 | |
| 151 | pred_kappa = norm_out[:, 3:, :, :] # (B, 1, H, W) |
| 152 | pred_kappa = pred_kappa.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 1) |
| 153 | pred_kappa_list.append(pred_kappa) |
| 154 | |
| 155 | # to numpy arrays |
| 156 | img = img.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 3) |
| 157 | gt_norm = gt_norm.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 3) |
| 158 | gt_norm_mask = gt_norm_mask.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 1) |
| 159 | |
| 160 | # input image |
| 161 | target_path = '%s/%08d_img.jpg' % (args.exp_vis_dir, total_iter) |
| 162 | img = unnormalize(img[0, ...]) |
| 163 | plt.imsave(target_path, img) |
| 164 | |
| 165 | # gt norm |
| 166 | gt_norm_rgb = ((gt_norm[0, ...] + 1) * 0.5) * 255 |
| 167 | gt_norm_rgb = np.clip(gt_norm_rgb, a_min=0, a_max=255) |
| 168 | gt_norm_rgb = gt_norm_rgb.astype(np.uint8) |
| 169 | |
| 170 | target_path = '%s/%08d_gt_norm.jpg' % (args.exp_vis_dir, total_iter) |
| 171 | plt.imsave(target_path, gt_norm_rgb * gt_norm_mask[0, ...]) |
| 172 | |
| 173 | # pred_norm |
| 174 | for i in range(len(pred_norm_list)): |
| 175 | pred_norm = pred_norm_list[i] |
| 176 | pred_norm_rgb = norm_to_rgb(pred_norm) |
| 177 | target_path = '%s/%08d_pred_norm_%d.jpg' % (args.exp_vis_dir, total_iter, i) |
| 178 | plt.imsave(target_path, pred_norm_rgb) |
| 179 | |
| 180 | pred_kappa = pred_kappa_list[i] |
| 181 | pred_alpha = kappa_to_alpha(pred_kappa) |
| 182 | target_path = '%s/%08d_pred_alpha_%d.jpg' % (args.exp_vis_dir, total_iter, i) |
| 183 | plt.imsave(target_path, pred_alpha[0, :, :, 0], vmin=0, vmax=60, cmap='jet') |
| 184 | |
| 185 | # error in angles |
| 186 | DP = np.sum(gt_norm * pred_norm, axis=3, keepdims=True) # (B, H, W, 1) |
| 187 | DP = np.clip(DP, -1, 1) |
| 188 | E = np.degrees(np.arccos(DP)) # (B, H, W, 1) |
| 189 | E = E * gt_norm_mask |
| 190 | target_path = '%s/%08d_pred_error_%d.jpg' % (args.exp_vis_dir, total_iter, i) |
| 191 | plt.imsave(target_path, E[0, :, :, 0], vmin=0, vmax=60, cmap='jet') |
nothing calls this directly
no test coverage detected