(tree: ProtoTree, project_info: dict, project_loader: DataLoader, folder_name: str, args: argparse.Namespace, log: Log)
| 12 | |
| 13 | # adapted from protopnet |
| 14 | def upsample(tree: ProtoTree, project_info: dict, project_loader: DataLoader, folder_name: str, args: argparse.Namespace, log: Log): |
| 15 | dir = os.path.join(os.path.join(args.log_dir, args.dir_for_saving_images), folder_name) |
| 16 | if not os.path.exists(dir): |
| 17 | os.makedirs(dir) |
| 18 | with torch.no_grad(): |
| 19 | sim_maps, project_info = get_similarity_maps(tree, project_info, log) |
| 20 | log.log_message("\nUpsampling prototypes for visualization...") |
| 21 | imgs = project_loader.dataset.imgs |
| 22 | for node, j in tree._out_map.items(): |
| 23 | if node in tree.branches: #do not upsample when node is pruned |
| 24 | prototype_info = project_info[j] |
| 25 | decision_node_idx = prototype_info['node_ix'] |
| 26 | x = Image.open(imgs[prototype_info['input_image_ix']][0]) |
| 27 | x.save(os.path.join(dir,'%s_original_image.png'%str(decision_node_idx))) |
| 28 | |
| 29 | x_np = np.asarray(x) |
| 30 | x_np = np.float32(x_np)/ 255 |
| 31 | if x_np.ndim == 2: #convert grayscale to RGB |
| 32 | x_np = np.stack((x_np,)*3, axis=-1) |
| 33 | |
| 34 | img_size = x_np.shape[:2] |
| 35 | similarity_map = sim_maps[j] |
| 36 | |
| 37 | rescaled_sim_map = similarity_map - np.amin(similarity_map) |
| 38 | rescaled_sim_map= rescaled_sim_map / np.amax(rescaled_sim_map) |
| 39 | similarity_heatmap = cv2.applyColorMap(np.uint8(255*rescaled_sim_map), cv2.COLORMAP_JET) |
| 40 | similarity_heatmap = np.float32(similarity_heatmap) / 255 |
| 41 | similarity_heatmap = similarity_heatmap[...,::-1] |
| 42 | plt.imsave(fname=os.path.join(dir,'%s_heatmap_latent_similaritymap.png'%str(decision_node_idx)), arr=similarity_heatmap, vmin=0.0,vmax=1.0) |
| 43 | |
| 44 | upsampled_act_pattern = cv2.resize(similarity_map, |
| 45 | dsize=(img_size[1], img_size[0]), |
| 46 | interpolation=cv2.INTER_CUBIC) |
| 47 | rescaled_act_pattern = upsampled_act_pattern - np.amin(upsampled_act_pattern) |
| 48 | rescaled_act_pattern = rescaled_act_pattern / np.amax(rescaled_act_pattern) |
| 49 | heatmap = cv2.applyColorMap(np.uint8(255*rescaled_act_pattern), cv2.COLORMAP_JET) |
| 50 | heatmap = np.float32(heatmap) / 255 |
| 51 | heatmap = heatmap[...,::-1] |
| 52 | |
| 53 | overlayed_original_img = 0.5 * x_np + 0.2 * heatmap |
| 54 | plt.imsave(fname=os.path.join(dir,'%s_heatmap_original_image.png'%str(decision_node_idx)), arr=overlayed_original_img, vmin=0.0,vmax=1.0) |
| 55 | |
| 56 | # save the highly activated patch |
| 57 | masked_similarity_map = np.zeros(similarity_map.shape) |
| 58 | prototype_index = prototype_info['patch_ix'] |
| 59 | W, H = prototype_info['W'], prototype_info['H'] |
| 60 | assert W == H |
| 61 | masked_similarity_map[prototype_index // W, prototype_index % W] = 1 #mask similarity map such that only the nearest patch z* is visualized |
| 62 | |
| 63 | upsampled_prototype_pattern = cv2.resize(masked_similarity_map, |
| 64 | dsize=(img_size[1], img_size[0]), |
| 65 | interpolation=cv2.INTER_CUBIC) |
| 66 | plt.imsave(fname=os.path.join(dir,'%s_masked_upsampled_heatmap.png'%str(decision_node_idx)), arr=upsampled_prototype_pattern, vmin=0.0,vmax=1.0) |
| 67 | |
| 68 | high_act_patch_indices = find_high_activation_crop(upsampled_prototype_pattern, args.upsample_threshold) |
| 69 | high_act_patch = x_np[high_act_patch_indices[0]:high_act_patch_indices[1], |
| 70 | high_act_patch_indices[2]:high_act_patch_indices[3], :] |
| 71 | plt.imsave(fname=os.path.join(dir,'%s_nearest_patch_of_image.png'%str(decision_node_idx)), arr=high_act_patch, vmin=0.0,vmax=1.0) |
no test coverage detected