| 17 | |
| 18 | |
| 19 | def test(model, test_loader, device, results_dir): |
| 20 | alpha_max = 60 |
| 21 | kappa_max = 30 |
| 22 | |
| 23 | with torch.no_grad(): |
| 24 | for data_dict in tqdm(test_loader): |
| 25 | |
| 26 | img = data_dict['img'].to(device) |
| 27 | norm_out_list, _, _ = model(img) |
| 28 | norm_out = norm_out_list[-1] |
| 29 | |
| 30 | pred_norm = norm_out[:, :3, :, :] |
| 31 | pred_kappa = norm_out[:, 3:, :, :] |
| 32 | |
| 33 | # to numpy arrays |
| 34 | img = img.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 3) |
| 35 | pred_norm = pred_norm.detach().cpu().permute(0, 2, 3, 1).numpy() # (B, H, W, 3) |
| 36 | pred_kappa = pred_kappa.cpu().permute(0, 2, 3, 1).numpy() |
| 37 | |
| 38 | # save results |
| 39 | img_name = data_dict['img_name'][0] |
| 40 | |
| 41 | # 1. save input image |
| 42 | img = utils.unnormalize(img[0, ...]) |
| 43 | |
| 44 | target_path = '%s/%s_img.png' % (results_dir, img_name) |
| 45 | plt.imsave(target_path, img) |
| 46 | |
| 47 | # 2. predicted normal |
| 48 | pred_norm_rgb = ((pred_norm + 1) * 0.5) * 255 |
| 49 | pred_norm_rgb = np.clip(pred_norm_rgb, a_min=0, a_max=255) |
| 50 | pred_norm_rgb = pred_norm_rgb.astype(np.uint8) # (B, H, W, 3) |
| 51 | |
| 52 | target_path = '%s/%s_pred_norm.png' % (results_dir, img_name) |
| 53 | plt.imsave(target_path, pred_norm_rgb[0, :, :, :]) |
| 54 | |
| 55 | # 3. predicted kappa (concentration parameter) |
| 56 | target_path = '%s/%s_pred_kappa.png' % (results_dir, img_name) |
| 57 | plt.imsave(target_path, pred_kappa[0, :, :, 0], vmin=0.0, vmax=kappa_max, cmap='gray') |
| 58 | |
| 59 | # 4. predicted uncertainty |
| 60 | pred_alpha = utils.kappa_to_alpha(pred_kappa) |
| 61 | target_path = '%s/%s_pred_alpha.png' % (results_dir, img_name) |
| 62 | plt.imsave(target_path, pred_alpha[0, :, :, 0], vmin=0.0, vmax=alpha_max, cmap='jet') |
| 63 | |
| 64 | # 5. concatenated results |
| 65 | image_path_list = ['img', 'pred_norm', 'pred_alpha'] |
| 66 | image_path_list = ['%s/%s_%s.png' % (results_dir, img_name, i) for i in image_path_list] |
| 67 | target_path = '%s/%s_concat.png' % (results_dir, img_name) |
| 68 | utils.concat_image(image_path_list, target_path) |
| 69 | |
| 70 | |
| 71 | if __name__ == '__main__': |