improve visualization code

utils/visualize_att_maps.py

@@ -26,7 +26,6 @@ class VisualizeAttentionMaps:
 
         self.resize_unnorm = inverse_normalize_w_resize(resize_resolution=self.save_resolution)
         self.num_parts = num_parts
-        self.req_colors = colors[:num_parts]
         self.figs_size = (10, 10)
 
     @torch.no_grad()
@@ -45,7 +44,11 @@ class VisualizeAttentionMaps:
         map_argmax = torch.nn.functional.interpolate(maps.clone().detach(), size=self.save_resolution,
                                                      mode='bilinear',
                                                      align_corners=True).argmax(dim=1).cpu().numpy()
-
-        curr_map = skimage.color.label2rgb(label=map_argmax[0], image=ims[0], colors=self.req_colors,
+        # Select colors for parts which are present
+        parts_present = np.unique(map_argmax).tolist()
+        if self.bg_label in parts_present:
+            parts_present.remove(self.bg_label)
+        colors_present = [colors[i] for i in parts_present]
+        curr_map = skimage.color.label2rgb(label=map_argmax[0], image=ims[0], colors=colors_present,
                                            bg_label=self.bg_label, alpha=self.alpha)
         return curr_map