# -*- coding: utf-8 -*-
"""pytorch_IBA_inserted_NIH1_30.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1YbVMIdKCn6EPPR9wvI7OIUS3Bxn6jRpr

# PyTorch: IBA (Per-Sample Bottleneck)

This notebook shows how to apply the Per-Sample Bottleneck to pretrained ImageNet models. 

Ensure that `./imagenet` points to your copy of the ImageNet dataset. 

You might want to create a symlink:
"""

from google.colab import drive
import sys
drive.mount('/content/drive', force_remount=True)

sys.path.append('/content/drive/MyDrive/Prak_MLMI')
sys.path.append('/content/drive/MyDrive')
sys.path.append('/content/drive/MyDrive/Prak_MLMI/model')

import warnings
warnings.filterwarnings('ignore')

import matplotlib.patches as patches

# Commented out IPython magic to ensure Python compatibility.
# to set you cuda device
# %env CUDA_VISIBLE_DEVICES=0

# %reload_ext autoreload
# %autoreload 2

import torch
import torchvision.models 
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize, Normalize, Grayscale
import matplotlib.pyplot as plt 
import os
from tqdm import tqdm_notebook
import json
from PIL import Image
import numpy as np
# import sys

try:
    import IBA
except ModuleNotFoundError:
    sys.path.insert(0, '..')
    import IBA
  
from IBA.pytorch_new import IBA, tensor_to_np_img
from IBA.utils import plot_saliency_map

import cxr_dataset as CXR

# Commented out IPython magic to ensure Python compatibility.
# #!unzip "/content/drive/MyDrive/NIH_CXR14_Resized.zip" -d "/content/drive/MyDrive/NIH_CXR14_Resized"
# %%capture
# !unzip "/content/drive/MyDrive/NIH_CXR14_Resized.zip" -d /content/

# %load /content/drive/MyDrive/Prak_MLMI/IBA/pytorch_new.py

"""## Loading Data and Model"""

prak_dir = '/content/drive/MyDrive/Prak_MLMI'
imagenet_dir = '/content/drive/MyDrive/Prak_MLMI/imagenet'
PATH_TO_IMAGES = "/content/NIH small"
MODEL_PATH = prak_dir + '/model/results/checkpoint_best'
label_path = '/content/drive/MyDrive/Prak_MLMI/model/labels'

dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

checkpoint = torch.load(MODEL_PATH, map_location=dev)
model = checkpoint['model'].module
model.to(dev).eval()
print(model)

# load the data

# use imagenet mean,std for normalization
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

data_transforms = {
    'train': Compose([
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean, std)
    ]),
    'val': Compose([
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean, std)
    ]),
}


bounding_box_transform = CXR.RescaleBB(224, 1024)

def get_label(LA):
    labels = {   
        'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Mass': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Nodule': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Fibrosis': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pleural_Thickening': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Hernia': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    }

    return labels.get(LA, None)

def model_loss_closure(input):
  loss = nn.BCEWithLogitsLoss()
  mse_loss = loss(model(input), torch.tensor(_label).view(1,-1).expand(10, -1).to(dev).float())
  return mse_loss

# Commented out IPython magic to ensure Python compatibility.
# %matplotlib inline

def samples_display(LABEL, beta=6):



  dataset = CXR.CXRDataset(
      path_to_images=PATH_TO_IMAGES,
      fold='BBox',#fold='train'
      transform=data_transforms['train'],
      transform_bb=bounding_box_transform,
      fine_tune=False,
      label_path=label_path,
      finding=LABEL)  #finding=LABEL
  dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
  global _label
  _label = get_label(LABEL)


  seed = np.random.randint(0, 1000000000)

  np.random.seed(seed)

  length = len(dataset)
  print("20 Ransom samples of "+LABEL+" from " + str(length) + " images")

  for sample_idx in np.random.choice(length, 20):  #change with dataset length
      iba = IBA(model.features.denseblock2)
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length, progbar=False) #change with dataset length
      img, target, idx, bbox = dataset[sample_idx]
      img = img[None].to(dev)
  
      # reverse the data pre-processing for plotting the original image
      np_img = tensor_to_np_img(img[0])

      size = 4
      rows = 1
      cols = 5
      fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols*(size+1), rows*size))


      cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
      mean = np.array([0.485, 0.456, 0.406])
      std = np.array([0.229, 0.224, 0.225])
      cxr = std * cxr + mean
      cxr = np.clip(cxr, 0, 1)

  #    rect_original = patches.Rectangle((bbox[0, 0], bbox[0, 1]), bbox[0, 2], bbox[0, 3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
      rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)

      ax.imshow(cxr)
      ax.axis('off')
      ax.set_title(idx)
      ax.add_patch(rect_original) 


      iba.reverse_lambda = False
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
      _ = ax2.set_title("Old method, Block2")
      ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba.reverse_lambda = True
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
      _ = ax3.set_title("New method, Block2")
      ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba = IBA(model.features.denseblock3)
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length, progbar=False)
      
      iba.reverse_lambda = False
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
      _ = ax4.set_title("Old method, Block3")
      ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba.reverse_lambda = True
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
      _ = ax5.set_title("New method, Block3")
      ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))



      plt.show()
  return length

"""# Cardiomegaly"""

LABEL='Cardiomegaly'

length_data=samples_display(LABEL)

dataset = CXR.CXRDataset(
    path_to_images=PATH_TO_IMAGES,
    fold='BBox',#fold='train'
    transform=data_transforms['train'],
    transform_bb=bounding_box_transform,
    fine_tune=False,
    label_path=label_path,
    finding=LABEL)  #finding=LABEL
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
img, target, idx, bbox = dataset[0]

"""## Compare different betas (So I chose beta=6 instead of 10, the default)

###IB in denseblock3 with different betas
"""

iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""###IB in denseblock2 with different betas"""

iba = IBA(model.features.denseblock2)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""###IB in denseblock4 with different betas"""

iba = IBA(model.features.denseblock4)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""## Compare New and Old method with IB in output of denseblock3 with another image as well as compare masked and unmasked prediction probability"""

img = img[None].to(dev)
np_img = tensor_to_np_img(img[0])

iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)


print("Old method")
iba.reverse_lambda = False
iba.beta = 6
#model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
heatmap = iba.analyze(img, model_loss_closure)  
# show the heatmap
ax = plot_saliency_map(heatmap, np_img)
_ = ax.set_title(idx)
ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

# show the prediction probability
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1,-1).to(dev).float())
with iba.restrict_flow():
  print(torch.sigmoid(model(img)))

plt.show()



print("Ground truth with bounding box")
cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
cxr = std * cxr + mean
cxr = np.clip(cxr, 0, 1)

rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)

fig, ax2 = plt.subplots(1, 1, figsize=(5.5, 4.0))

ax2.imshow(cxr)
ax2.axis('off')
ax2.set_title(idx)
ax2.add_patch(rect_original) 
plt.show()

print("New method")
iba.reverse_lambda = True
iba.beta = 6
#model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
heatmap = iba.analyze(img, model_loss_closure) 
# show the heatmap
ax3 = plot_saliency_map(heatmap, np_img)
_ = ax3.set_title(idx)
ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
# show the prediction probability
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1,-1).to(dev).float())
with iba.restrict_flow():
  print(torch.sigmoid(model(img)))
plt.show()

"""## New: Search for a proper threshold to generate mask"""

length_data = 146
_label = get_label(LABEL)

def compare_threshold(index=0, new_method = False):

  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])


  iba = IBA(model.features.denseblock3)
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

  if new_method: 
      print("New method")
  else:
      print("Old method")
      
  iba.reverse_lambda = new_method
  iba.beta = 6
  #model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
  heatmap = iba.analyze(img, model_loss_closure)  
  # show the heatmap
  ax = plot_saliency_map(heatmap, np_img)
  _ = ax.set_title(idx)
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

  heatmap = heatmap / np.linalg.norm(heatmap)


  size = 4
  rows = 5
  cols = 5
  fig, axes = plt.subplots(5, 5, figsize=(cols*(size+1), rows*size))


  threshold = 0.0015
  for ax in axes.flatten():
    mask = heatmap.copy()
    mask[mask < threshold] = 0
    mask[mask >= threshold] = 1
    
    ax.matshow(mask)
    ax.set_title("thresh={}".format(threshold),pad = 20)
    threshold += 0.00025
    

  plt.subplots_adjust(top=1.2)    
  plt.show()

"""### old method"""

compare_threshold(0)

compare_threshold(1)

compare_threshold(2)

"""### New method"""

compare_threshold(index=0, new_method = True)

compare_threshold(index=1, new_method = True)

"""## New: Initialzation of IB after block1 with mask generated from IB after block3"""

def new_initit(position = "1", threshold = 0.006, index=0, more_iter = False, new_method = False):

  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  iba = IBA(model.features.denseblock3)
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

  if new_method: 
      print("New method")
  else:
      print("Old method")
      
  iba.reverse_lambda = new_method
  iba.beta = 6
  heatmap = iba.analyze(img, model_loss_closure)  

  size = 4
  rows = 1
  cols = 2
  fig, (ax, axe1) = plt.subplots(1, 2, figsize=(cols*(size+3), rows*size)) 

  # show the heatmap
  ax = plot_saliency_map(heatmap, np_img, ax = ax )
  _ = ax.set_title(idx)
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  heatmap = heatmap / np.linalg.norm(heatmap)

  mask = heatmap.copy()
  mask[mask < threshold] = -5
  mask[mask >= threshold] = 5    
  axe1.matshow(mask)
  axe1.set_title("thresh={}".format(threshold),pad = 20)
 

  plt.subplots_adjust(top=1.2)    
  plt.show()
  
  #alter and save the old mask
  alpha = torch.tensor(mask).to(dev)
  alpha = alpha.unsqueeze(0)
  alpha = alpha.unsqueeze(0) 

  iba = IBA(eval("model.features.denseblock{0}".format(*position)))
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)
  iba.reverse_lambda = new_method
  iba.beta = 6

  new_alpha = nn.functional.interpolate(alpha, iba.estimator.shape[1:])
  new_alpha = new_alpha.squeeze(0)
  new_alpha = new_alpha.expand(iba.estimator.shape[0],-1,-1).clone()

  iba.old_alpha= torch.nn.Parameter(new_alpha,  requires_grad=True) 
  print("With old mask generated from IB after denseblock3 for new IB after denseblock " + position)

  if more_iter:
    fig, axes = plt.subplots(3, 4, figsize=(20,15))

    iter_set = [ 1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
    for iter_value, ax in zip(iter_set, axes.flatten()) : 
        heatmap = iba.analyze(img, model_loss_closure, optimization_steps = iter_value) 
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title("optimization steps={}".format(iter_value))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
  else: 
    fig, axes = plt.subplots(2, 4, figsize=(20,10))

    iter_set = [ 1, 2, 4, 6, 8, 10, 12, 14]
    for iter_value, ax in zip(iter_set, axes.flatten()) : 
        heatmap = iba.analyze(img, model_loss_closure, optimization_steps = iter_value) 
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title("optimization steps={}".format(iter_value))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))



  plt.show()

"""### old method"""

new_initit()

new_initit(position = "1", threshold = 0.003, index=0)

new_initit(position = "1", threshold = 0.003, index=1)

"""### New method"""

new_initit(more_iter = True, new_method = True)

new_initit(index=1, more_iter = True, new_method = True)

"""## New: Initialzation of IB after block2 with mask generated from IB after block3"""

new_initit(position = "2", threshold = 0.006, index=0, more_iter = True)

new_initit(position = "2", threshold = 0.003, index=0)

"""## New: Initialzation of IB after block3 with mask generated from IB after block3"""

new_initit(position = "3", threshold = 0.003, index=0)

#wrong initialization, highlighted shoulder area
# new_initit(position = "1", shift = 2, index=0)

"""## New: Compare IB in different layers of Block4"""

def compare_layers(length, ind=0, beta = 6, more_layer = False, new_method = False):

  index = min (ind, length) 
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])
  if new_method: 
    print("New method")
  else:
    print("Old method")

  if more_layer:
    size = 4
    rows = 4
    cols = 4
    fig, axes = plt.subplots(4, 4, figsize=(cols*(size+1), rows*size))

    positions = [["4","1"], ["4","2"], ["4","3"],["4","4"],["4","5"],["4","6"],["4","7"],["4","8"],\
                 ["4","9"],["4","10"],["4","11"],["4","12"],["4","13"],["4","14"],["4","15"],["4","16"]]

    for position, ax in zip(positions, axes.flatten()):
        current_pos = "model.features.denseblock{0}.denselayer{1}".format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)    
        iba.reverse_lambda = new_method
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)  
        ax.set_title("denseblock{0}.denselayer{1}".format(*position))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

    plt.show()

  else: 
    size = 4
    rows = 2
    cols = 4
    fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

    positions = [["4","1"], ["4","2"], ["4","3"],["4","4"],["4","5"],["4","6"],["4","7"],["4","8"]]

    for position, ax in zip(positions, axes.flatten()):
        current_pos = "model.features.denseblock{0}.denselayer{1}".format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)          
        iba.reverse_lambda = new_method

        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)  
        ax.set_title("denseblock{0}.denselayer{1}".format(*position))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

    plt.show()

compare_layers(length_data,0)
compare_layers(length_data,1)
compare_layers(length_data,2)

compare_layers(length_data,0, 3)
compare_layers(length_data,0, 0.5)

# beta=0.1, beta has to be extremely small (default beta from the papaer: 10)
compare_layers(length_data,0, 0.1, more_layer = True)

"""## New: concatenated mask"""

def norm_concat_heatmaps(index=0, weight = [1, 1, 1], new_method = False):  

# normalized_concatenated_heatmaps_generator
  if new_method: 
      print("New method")
  else:
      print("Old method")
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  heatmap = []

  size = 4
  rows = 1
  cols = 3
  fig, axes = plt.subplots(1, 3, figsize=(cols*(size+1), rows*size)) 

  # show the original heatmaps
  for i in range(3):

   
    current_pos = "model.features.denseblock{0}".format(i+1)
    iba = IBA(eval(current_pos))

    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

    iba.reverse_lambda = new_method
    iba.beta = 6
    hm = iba.analyze(img, model_loss_closure)
    hm = hm / np.linalg.norm(hm)
    heatmap.append (hm)
    # heatmap.append (iba.analyze(img, model_loss_closure))

    ax = axes[i]
    ax = plot_saliency_map(heatmap[i], np_img, ax=ax)  
    ax.set_title("normalized heatmap from mask afte block "+ str(i+1))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

  heatmap = np.array(heatmap) 
  weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1] + heatmap[2] * weight[2]
  ax = plot_saliency_map(weighted_heatmap, np_img)
  _ = ax.set_title("weight=" + str(weight))
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
  plt.show()

"""### Old method"""

norm_concat_heatmaps()

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

norm_concat_heatmaps(1)

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

norm_concat_heatmaps(3)

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

"""### New method"""

norm_concat_heatmaps(new_method = True)

"""## Old method (try with different blocks of DenseNet) with other images

"""

def compare_blocks(length, ind=0):

  index = min (ind, length) 
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  size = 4
  rows = 1
  cols = 4
  fig, axes = plt.subplots(1, 4, figsize=(cols*(size+1), rows*size))

  positions = [["1"], ["2"], ["3"],["4"]]

  for position, ax in zip(positions, axes.flatten()):
      current_pos = "model.features.denseblock{0}".format(*position)
      iba = IBA(eval(current_pos))
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

      iba.beta = 6
      heatmap = iba.analyze(img, model_loss_closure)
      ax = plot_saliency_map(heatmap, np_img, ax=ax)  
      ax.set_title(current_pos)
      ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

compare_blocks(length_data,0)
compare_blocks(length_data,1)
compare_blocks(length_data,2)

"""# Atelectasis"""

LABEL='Atelectasis'

length_data=samples_display(LABEL)

"""# Effusion"""

LABEL='Effusion'
length_data=samples_display(LABEL)

"""# Mass"""

LABEL='Mass'
length_data=samples_display(LABEL)

"""# Nodule"""

LABEL='Nodule'
length_data=samples_display(LABEL)

"""# Pneumonia"""

LABEL='Pneumonia'
length_data=samples_display(LABEL)

"""# Pneumothorax"""

LABEL='Pneumothorax'
length_data=samples_display(LABEL)
3b4c0c250414a443af09fd5cf51e10b84b42b9e1
<mask token>


def get_label(LA):
    labels = {'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Mass':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Nodule': [0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0], 'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0], 'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0], 'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Fibrosis':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pleural_Thickening': [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Hernia': [0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0]}
    return labels.get(LA, None)


def model_loss_closure(input):
    loss = nn.BCEWithLogitsLoss()
    mse_loss = loss(model(input), torch.tensor(_label).view(1, -1).expand(
        10, -1).to(dev).float())
    return mse_loss


def samples_display(LABEL, beta=6):
    dataset = CXR.CXRDataset(path_to_images=PATH_TO_IMAGES, fold='BBox',
        transform=data_transforms['train'], transform_bb=
        bounding_box_transform, fine_tune=False, label_path=label_path,
        finding=LABEL)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1
        )
    global _label
    _label = get_label(LABEL)
    seed = np.random.randint(0, 1000000000)
    np.random.seed(seed)
    length = len(dataset)
    print('20 Ransom samples of ' + LABEL + ' from ' + str(length) + ' images')
    for sample_idx in np.random.choice(length, 20):
        iba = IBA(model.features.denseblock2)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        img, target, idx, bbox = dataset[sample_idx]
        img = img[None].to(dev)
        np_img = tensor_to_np_img(img[0])
        size = 4
        rows = 1
        cols = 5
        fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols *
            (size + 1), rows * size))
        cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        cxr = std * cxr + mean
        cxr = np.clip(cxr, 0, 1)
        rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox
            [3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
        ax.imshow(cxr)
        ax.axis('off')
        ax.set_title(idx)
        ax.add_patch(rect_original)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
        _ = ax2.set_title('Old method, Block2')
        ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
        _ = ax3.set_title('New method, Block2')
        ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba = IBA(model.features.denseblock3)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
        _ = ax4.set_title('Old method, Block3')
        ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
        _ = ax5.set_title('New method, Block3')
        ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        plt.show()
    return length


<mask token>


def compare_threshold(index=0, new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    ax = plot_saliency_map(heatmap, np_img)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = heatmap / np.linalg.norm(heatmap)
    size = 4
    rows = 5
    cols = 5
    fig, axes = plt.subplots(5, 5, figsize=(cols * (size + 1), rows * size))
    threshold = 0.0015
    for ax in axes.flatten():
        mask = heatmap.copy()
        mask[mask < threshold] = 0
        mask[mask >= threshold] = 1
        ax.matshow(mask)
        ax.set_title('thresh={}'.format(threshold), pad=20)
        threshold += 0.00025
    plt.subplots_adjust(top=1.2)
    plt.show()


<mask token>


def compare_layers(length, ind=0, beta=6, more_layer=False, new_method=False):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    if new_method:
        print('New method')
    else:
        print('Old method')
    if more_layer:
        size = 4
        rows = 4
        cols = 4
        fig, axes = plt.subplots(4, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8'], ['4', '9'], ['4',
            '10'], ['4', '11'], ['4', '12'], ['4', '13'], ['4', '14'], ['4',
            '15'], ['4', '16']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()
    else:
        size = 4
        rows = 2
        cols = 4
        fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()


<mask token>


def norm_concat_heatmaps(index=0, weight=[1, 1, 1], new_method=False):
    if new_method:
        print('New method')
    else:
        print('Old method')
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    heatmap = []
    size = 4
    rows = 1
    cols = 3
    fig, axes = plt.subplots(1, 3, figsize=(cols * (size + 1), rows * size))
    for i in range(3):
        current_pos = 'model.features.denseblock{0}'.format(i + 1)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.reverse_lambda = new_method
        iba.beta = 6
        hm = iba.analyze(img, model_loss_closure)
        hm = hm / np.linalg.norm(hm)
        heatmap.append(hm)
        ax = axes[i]
        ax = plot_saliency_map(heatmap[i], np_img, ax=ax)
        ax.set_title('normalized heatmap from mask afte block ' + str(i + 1))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = np.array(heatmap)
    weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1
        ] + heatmap[2] * weight[2]
    ax = plot_saliency_map(weighted_heatmap, np_img)
    _ = ax.set_title('weight=' + str(weight))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>


def compare_blocks(length, ind=0):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    size = 4
    rows = 1
    cols = 4
    fig, axes = plt.subplots(1, 4, figsize=(cols * (size + 1), rows * size))
    positions = [['1'], ['2'], ['3'], ['4']]
    for position, ax in zip(positions, axes.flatten()):
        current_pos = 'model.features.denseblock{0}'.format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.beta = 6
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title(current_pos)
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>

<mask token>


def get_label(LA):
    labels = {'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Mass':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Nodule': [0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0], 'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0], 'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0], 'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Fibrosis':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pleural_Thickening': [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Hernia': [0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0]}
    return labels.get(LA, None)


def model_loss_closure(input):
    loss = nn.BCEWithLogitsLoss()
    mse_loss = loss(model(input), torch.tensor(_label).view(1, -1).expand(
        10, -1).to(dev).float())
    return mse_loss


def samples_display(LABEL, beta=6):
    dataset = CXR.CXRDataset(path_to_images=PATH_TO_IMAGES, fold='BBox',
        transform=data_transforms['train'], transform_bb=
        bounding_box_transform, fine_tune=False, label_path=label_path,
        finding=LABEL)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1
        )
    global _label
    _label = get_label(LABEL)
    seed = np.random.randint(0, 1000000000)
    np.random.seed(seed)
    length = len(dataset)
    print('20 Ransom samples of ' + LABEL + ' from ' + str(length) + ' images')
    for sample_idx in np.random.choice(length, 20):
        iba = IBA(model.features.denseblock2)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        img, target, idx, bbox = dataset[sample_idx]
        img = img[None].to(dev)
        np_img = tensor_to_np_img(img[0])
        size = 4
        rows = 1
        cols = 5
        fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols *
            (size + 1), rows * size))
        cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        cxr = std * cxr + mean
        cxr = np.clip(cxr, 0, 1)
        rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox
            [3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
        ax.imshow(cxr)
        ax.axis('off')
        ax.set_title(idx)
        ax.add_patch(rect_original)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
        _ = ax2.set_title('Old method, Block2')
        ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
        _ = ax3.set_title('New method, Block2')
        ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba = IBA(model.features.denseblock3)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
        _ = ax4.set_title('Old method, Block3')
        ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
        _ = ax5.set_title('New method, Block3')
        ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        plt.show()
    return length


<mask token>


def compare_threshold(index=0, new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    ax = plot_saliency_map(heatmap, np_img)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = heatmap / np.linalg.norm(heatmap)
    size = 4
    rows = 5
    cols = 5
    fig, axes = plt.subplots(5, 5, figsize=(cols * (size + 1), rows * size))
    threshold = 0.0015
    for ax in axes.flatten():
        mask = heatmap.copy()
        mask[mask < threshold] = 0
        mask[mask >= threshold] = 1
        ax.matshow(mask)
        ax.set_title('thresh={}'.format(threshold), pad=20)
        threshold += 0.00025
    plt.subplots_adjust(top=1.2)
    plt.show()


<mask token>


def new_initit(position='1', threshold=0.006, index=0, more_iter=False,
    new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    size = 4
    rows = 1
    cols = 2
    fig, (ax, axe1) = plt.subplots(1, 2, figsize=(cols * (size + 3), rows *
        size))
    ax = plot_saliency_map(heatmap, np_img, ax=ax)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    heatmap = heatmap / np.linalg.norm(heatmap)
    mask = heatmap.copy()
    mask[mask < threshold] = -5
    mask[mask >= threshold] = 5
    axe1.matshow(mask)
    axe1.set_title('thresh={}'.format(threshold), pad=20)
    plt.subplots_adjust(top=1.2)
    plt.show()
    alpha = torch.tensor(mask).to(dev)
    alpha = alpha.unsqueeze(0)
    alpha = alpha.unsqueeze(0)
    iba = IBA(eval('model.features.denseblock{0}'.format(*position)))
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    iba.reverse_lambda = new_method
    iba.beta = 6
    new_alpha = nn.functional.interpolate(alpha, iba.estimator.shape[1:])
    new_alpha = new_alpha.squeeze(0)
    new_alpha = new_alpha.expand(iba.estimator.shape[0], -1, -1).clone()
    iba.old_alpha = torch.nn.Parameter(new_alpha, requires_grad=True)
    print(
        'With old mask generated from IB after denseblock3 for new IB after denseblock '
         + position)
    if more_iter:
        fig, axes = plt.subplots(3, 4, figsize=(20, 15))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    else:
        fig, axes = plt.subplots(2, 4, figsize=(20, 10))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    plt.show()


<mask token>


def compare_layers(length, ind=0, beta=6, more_layer=False, new_method=False):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    if new_method:
        print('New method')
    else:
        print('Old method')
    if more_layer:
        size = 4
        rows = 4
        cols = 4
        fig, axes = plt.subplots(4, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8'], ['4', '9'], ['4',
            '10'], ['4', '11'], ['4', '12'], ['4', '13'], ['4', '14'], ['4',
            '15'], ['4', '16']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()
    else:
        size = 4
        rows = 2
        cols = 4
        fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()


<mask token>


def norm_concat_heatmaps(index=0, weight=[1, 1, 1], new_method=False):
    if new_method:
        print('New method')
    else:
        print('Old method')
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    heatmap = []
    size = 4
    rows = 1
    cols = 3
    fig, axes = plt.subplots(1, 3, figsize=(cols * (size + 1), rows * size))
    for i in range(3):
        current_pos = 'model.features.denseblock{0}'.format(i + 1)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.reverse_lambda = new_method
        iba.beta = 6
        hm = iba.analyze(img, model_loss_closure)
        hm = hm / np.linalg.norm(hm)
        heatmap.append(hm)
        ax = axes[i]
        ax = plot_saliency_map(heatmap[i], np_img, ax=ax)
        ax.set_title('normalized heatmap from mask afte block ' + str(i + 1))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = np.array(heatmap)
    weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1
        ] + heatmap[2] * weight[2]
    ax = plot_saliency_map(weighted_heatmap, np_img)
    _ = ax.set_title('weight=' + str(weight))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>


def compare_blocks(length, ind=0):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    size = 4
    rows = 1
    cols = 4
    fig, axes = plt.subplots(1, 4, figsize=(cols * (size + 1), rows * size))
    positions = [['1'], ['2'], ['3'], ['4']]
    for position, ax in zip(positions, axes.flatten()):
        current_pos = 'model.features.denseblock{0}'.format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.beta = 6
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title(current_pos)
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>

<mask token>
drive.mount('/content/drive', force_remount=True)
sys.path.append('/content/drive/MyDrive/Prak_MLMI')
sys.path.append('/content/drive/MyDrive')
sys.path.append('/content/drive/MyDrive/Prak_MLMI/model')
<mask token>
warnings.filterwarnings('ignore')
<mask token>
try:
    import IBA
except ModuleNotFoundError:
    sys.path.insert(0, '..')
    import IBA
<mask token>
model.to(dev).eval()
print(model)
<mask token>


def get_label(LA):
    labels = {'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Mass':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Nodule': [0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0], 'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0], 'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0], 'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Fibrosis':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pleural_Thickening': [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Hernia': [0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0]}
    return labels.get(LA, None)


def model_loss_closure(input):
    loss = nn.BCEWithLogitsLoss()
    mse_loss = loss(model(input), torch.tensor(_label).view(1, -1).expand(
        10, -1).to(dev).float())
    return mse_loss


def samples_display(LABEL, beta=6):
    dataset = CXR.CXRDataset(path_to_images=PATH_TO_IMAGES, fold='BBox',
        transform=data_transforms['train'], transform_bb=
        bounding_box_transform, fine_tune=False, label_path=label_path,
        finding=LABEL)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1
        )
    global _label
    _label = get_label(LABEL)
    seed = np.random.randint(0, 1000000000)
    np.random.seed(seed)
    length = len(dataset)
    print('20 Ransom samples of ' + LABEL + ' from ' + str(length) + ' images')
    for sample_idx in np.random.choice(length, 20):
        iba = IBA(model.features.denseblock2)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        img, target, idx, bbox = dataset[sample_idx]
        img = img[None].to(dev)
        np_img = tensor_to_np_img(img[0])
        size = 4
        rows = 1
        cols = 5
        fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols *
            (size + 1), rows * size))
        cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        cxr = std * cxr + mean
        cxr = np.clip(cxr, 0, 1)
        rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox
            [3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
        ax.imshow(cxr)
        ax.axis('off')
        ax.set_title(idx)
        ax.add_patch(rect_original)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
        _ = ax2.set_title('Old method, Block2')
        ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
        _ = ax3.set_title('New method, Block2')
        ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba = IBA(model.features.denseblock3)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
        _ = ax4.set_title('Old method, Block3')
        ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
        _ = ax5.set_title('New method, Block3')
        ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        plt.show()
    return length


<mask token>
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
<mask token>
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
<mask token>
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
<mask token>
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=
    False)
print('Old method')
<mask token>
ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
    linewidth=2, edgecolor='r', facecolor='none', zorder=2))
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1, -1).to(dev).float())
with iba.restrict_flow():
    print(torch.sigmoid(model(img)))
plt.show()
print('Ground truth with bounding box')
<mask token>
ax2.imshow(cxr)
ax2.axis('off')
ax2.set_title(idx)
ax2.add_patch(rect_original)
plt.show()
print('New method')
<mask token>
ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
    linewidth=2, edgecolor='r', facecolor='none', zorder=2))
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1, -1).to(dev).float())
with iba.restrict_flow():
    print(torch.sigmoid(model(img)))
plt.show()
<mask token>


def compare_threshold(index=0, new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    ax = plot_saliency_map(heatmap, np_img)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = heatmap / np.linalg.norm(heatmap)
    size = 4
    rows = 5
    cols = 5
    fig, axes = plt.subplots(5, 5, figsize=(cols * (size + 1), rows * size))
    threshold = 0.0015
    for ax in axes.flatten():
        mask = heatmap.copy()
        mask[mask < threshold] = 0
        mask[mask >= threshold] = 1
        ax.matshow(mask)
        ax.set_title('thresh={}'.format(threshold), pad=20)
        threshold += 0.00025
    plt.subplots_adjust(top=1.2)
    plt.show()


<mask token>
compare_threshold(0)
compare_threshold(1)
compare_threshold(2)
<mask token>
compare_threshold(index=0, new_method=True)
compare_threshold(index=1, new_method=True)
<mask token>


def new_initit(position='1', threshold=0.006, index=0, more_iter=False,
    new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    size = 4
    rows = 1
    cols = 2
    fig, (ax, axe1) = plt.subplots(1, 2, figsize=(cols * (size + 3), rows *
        size))
    ax = plot_saliency_map(heatmap, np_img, ax=ax)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    heatmap = heatmap / np.linalg.norm(heatmap)
    mask = heatmap.copy()
    mask[mask < threshold] = -5
    mask[mask >= threshold] = 5
    axe1.matshow(mask)
    axe1.set_title('thresh={}'.format(threshold), pad=20)
    plt.subplots_adjust(top=1.2)
    plt.show()
    alpha = torch.tensor(mask).to(dev)
    alpha = alpha.unsqueeze(0)
    alpha = alpha.unsqueeze(0)
    iba = IBA(eval('model.features.denseblock{0}'.format(*position)))
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    iba.reverse_lambda = new_method
    iba.beta = 6
    new_alpha = nn.functional.interpolate(alpha, iba.estimator.shape[1:])
    new_alpha = new_alpha.squeeze(0)
    new_alpha = new_alpha.expand(iba.estimator.shape[0], -1, -1).clone()
    iba.old_alpha = torch.nn.Parameter(new_alpha, requires_grad=True)
    print(
        'With old mask generated from IB after denseblock3 for new IB after denseblock '
         + position)
    if more_iter:
        fig, axes = plt.subplots(3, 4, figsize=(20, 15))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    else:
        fig, axes = plt.subplots(2, 4, figsize=(20, 10))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    plt.show()


<mask token>
new_initit()
new_initit(position='1', threshold=0.003, index=0)
new_initit(position='1', threshold=0.003, index=1)
<mask token>
new_initit(more_iter=True, new_method=True)
new_initit(index=1, more_iter=True, new_method=True)
<mask token>
new_initit(position='2', threshold=0.006, index=0, more_iter=True)
new_initit(position='2', threshold=0.003, index=0)
<mask token>
new_initit(position='3', threshold=0.003, index=0)
<mask token>


def compare_layers(length, ind=0, beta=6, more_layer=False, new_method=False):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    if new_method:
        print('New method')
    else:
        print('Old method')
    if more_layer:
        size = 4
        rows = 4
        cols = 4
        fig, axes = plt.subplots(4, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8'], ['4', '9'], ['4',
            '10'], ['4', '11'], ['4', '12'], ['4', '13'], ['4', '14'], ['4',
            '15'], ['4', '16']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()
    else:
        size = 4
        rows = 2
        cols = 4
        fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()


compare_layers(length_data, 0)
compare_layers(length_data, 1)
compare_layers(length_data, 2)
compare_layers(length_data, 0, 3)
compare_layers(length_data, 0, 0.5)
compare_layers(length_data, 0, 0.1, more_layer=True)
<mask token>


def norm_concat_heatmaps(index=0, weight=[1, 1, 1], new_method=False):
    if new_method:
        print('New method')
    else:
        print('Old method')
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    heatmap = []
    size = 4
    rows = 1
    cols = 3
    fig, axes = plt.subplots(1, 3, figsize=(cols * (size + 1), rows * size))
    for i in range(3):
        current_pos = 'model.features.denseblock{0}'.format(i + 1)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.reverse_lambda = new_method
        iba.beta = 6
        hm = iba.analyze(img, model_loss_closure)
        hm = hm / np.linalg.norm(hm)
        heatmap.append(hm)
        ax = axes[i]
        ax = plot_saliency_map(heatmap[i], np_img, ax=ax)
        ax.set_title('normalized heatmap from mask afte block ' + str(i + 1))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = np.array(heatmap)
    weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1
        ] + heatmap[2] * weight[2]
    ax = plot_saliency_map(weighted_heatmap, np_img)
    _ = ax.set_title('weight=' + str(weight))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>
norm_concat_heatmaps()
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
norm_concat_heatmaps(1)
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
norm_concat_heatmaps(3)
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
<mask token>
norm_concat_heatmaps(new_method=True)
<mask token>


def compare_blocks(length, ind=0):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    size = 4
    rows = 1
    cols = 4
    fig, axes = plt.subplots(1, 4, figsize=(cols * (size + 1), rows * size))
    positions = [['1'], ['2'], ['3'], ['4']]
    for position, ax in zip(positions, axes.flatten()):
        current_pos = 'model.features.denseblock{0}'.format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.beta = 6
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title(current_pos)
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


compare_blocks(length_data, 0)
compare_blocks(length_data, 1)
compare_blocks(length_data, 2)
<mask token>

<mask token>
from google.colab import drive
import sys
drive.mount('/content/drive', force_remount=True)
sys.path.append('/content/drive/MyDrive/Prak_MLMI')
sys.path.append('/content/drive/MyDrive')
sys.path.append('/content/drive/MyDrive/Prak_MLMI/model')
import warnings
warnings.filterwarnings('ignore')
import matplotlib.patches as patches
import torch
import torchvision.models
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize, Normalize, Grayscale
import matplotlib.pyplot as plt
import os
from tqdm import tqdm_notebook
import json
from PIL import Image
import numpy as np
try:
    import IBA
except ModuleNotFoundError:
    sys.path.insert(0, '..')
    import IBA
from IBA.pytorch_new import IBA, tensor_to_np_img
from IBA.utils import plot_saliency_map
import cxr_dataset as CXR
<mask token>
prak_dir = '/content/drive/MyDrive/Prak_MLMI'
imagenet_dir = '/content/drive/MyDrive/Prak_MLMI/imagenet'
PATH_TO_IMAGES = '/content/NIH small'
MODEL_PATH = prak_dir + '/model/results/checkpoint_best'
label_path = '/content/drive/MyDrive/Prak_MLMI/model/labels'
dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
checkpoint = torch.load(MODEL_PATH, map_location=dev)
model = checkpoint['model'].module
model.to(dev).eval()
print(model)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
data_transforms = {'train': Compose([Resize(224), CenterCrop(224), ToTensor
    (), Normalize(mean, std)]), 'val': Compose([Resize(224), CenterCrop(224
    ), ToTensor(), Normalize(mean, std)])}
bounding_box_transform = CXR.RescaleBB(224, 1024)


def get_label(LA):
    labels = {'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Mass':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Nodule': [0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0], 'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0], 'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0], 'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Fibrosis':
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Pleural_Thickening': [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'Hernia': [0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0]}
    return labels.get(LA, None)


def model_loss_closure(input):
    loss = nn.BCEWithLogitsLoss()
    mse_loss = loss(model(input), torch.tensor(_label).view(1, -1).expand(
        10, -1).to(dev).float())
    return mse_loss


def samples_display(LABEL, beta=6):
    dataset = CXR.CXRDataset(path_to_images=PATH_TO_IMAGES, fold='BBox',
        transform=data_transforms['train'], transform_bb=
        bounding_box_transform, fine_tune=False, label_path=label_path,
        finding=LABEL)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1
        )
    global _label
    _label = get_label(LABEL)
    seed = np.random.randint(0, 1000000000)
    np.random.seed(seed)
    length = len(dataset)
    print('20 Ransom samples of ' + LABEL + ' from ' + str(length) + ' images')
    for sample_idx in np.random.choice(length, 20):
        iba = IBA(model.features.denseblock2)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        img, target, idx, bbox = dataset[sample_idx]
        img = img[None].to(dev)
        np_img = tensor_to_np_img(img[0])
        size = 4
        rows = 1
        cols = 5
        fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols *
            (size + 1), rows * size))
        cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        cxr = std * cxr + mean
        cxr = np.clip(cxr, 0, 1)
        rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox
            [3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
        ax.imshow(cxr)
        ax.axis('off')
        ax.set_title(idx)
        ax.add_patch(rect_original)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
        _ = ax2.set_title('Old method, Block2')
        ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
        _ = ax3.set_title('New method, Block2')
        ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba = IBA(model.features.denseblock3)
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length,
            progbar=False)
        iba.reverse_lambda = False
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
        _ = ax4.set_title('Old method, Block3')
        ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        iba.reverse_lambda = True
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
        _ = ax5.set_title('New method, Block3')
        ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3
            ], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
        plt.show()
    return length


<mask token>
LABEL = 'Cardiomegaly'
length_data = samples_display(LABEL)
dataset = CXR.CXRDataset(path_to_images=PATH_TO_IMAGES, fold='BBox',
    transform=data_transforms['train'], transform_bb=bounding_box_transform,
    fine_tune=False, label_path=label_path, finding=LABEL)
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
img, target, idx, bbox = dataset[0]
<mask token>
iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size))
beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
iba = IBA(model.features.denseblock2)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size))
beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
iba = IBA(model.features.denseblock4)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar
    =False)
size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size))
beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()):
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta=
        beta_value)
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title('beta={}'.format(beta_value))
plt.show()
<mask token>
img = img[None].to(dev)
np_img = tensor_to_np_img(img[0])
iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=
    False)
print('Old method')
iba.reverse_lambda = False
iba.beta = 6
heatmap = iba.analyze(img, model_loss_closure)
ax = plot_saliency_map(heatmap, np_img)
_ = ax.set_title(idx)
ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
    linewidth=2, edgecolor='r', facecolor='none', zorder=2))
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1, -1).to(dev).float())
with iba.restrict_flow():
    print(torch.sigmoid(model(img)))
plt.show()
print('Ground truth with bounding box')
cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
cxr = std * cxr + mean
cxr = np.clip(cxr, 0, 1)
rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
    linewidth=2, edgecolor='r', facecolor='none', zorder=2)
fig, ax2 = plt.subplots(1, 1, figsize=(5.5, 4.0))
ax2.imshow(cxr)
ax2.axis('off')
ax2.set_title(idx)
ax2.add_patch(rect_original)
plt.show()
print('New method')
iba.reverse_lambda = True
iba.beta = 6
heatmap = iba.analyze(img, model_loss_closure)
ax3 = plot_saliency_map(heatmap, np_img)
_ = ax3.set_title(idx)
ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
    linewidth=2, edgecolor='r', facecolor='none', zorder=2))
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1, -1).to(dev).float())
with iba.restrict_flow():
    print(torch.sigmoid(model(img)))
plt.show()
<mask token>
length_data = 146
_label = get_label(LABEL)


def compare_threshold(index=0, new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    ax = plot_saliency_map(heatmap, np_img)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = heatmap / np.linalg.norm(heatmap)
    size = 4
    rows = 5
    cols = 5
    fig, axes = plt.subplots(5, 5, figsize=(cols * (size + 1), rows * size))
    threshold = 0.0015
    for ax in axes.flatten():
        mask = heatmap.copy()
        mask[mask < threshold] = 0
        mask[mask >= threshold] = 1
        ax.matshow(mask)
        ax.set_title('thresh={}'.format(threshold), pad=20)
        threshold += 0.00025
    plt.subplots_adjust(top=1.2)
    plt.show()


<mask token>
compare_threshold(0)
compare_threshold(1)
compare_threshold(2)
<mask token>
compare_threshold(index=0, new_method=True)
compare_threshold(index=1, new_method=True)
<mask token>


def new_initit(position='1', threshold=0.006, index=0, more_iter=False,
    new_method=False):
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    iba = IBA(model.features.denseblock3)
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    if new_method:
        print('New method')
    else:
        print('Old method')
    iba.reverse_lambda = new_method
    iba.beta = 6
    heatmap = iba.analyze(img, model_loss_closure)
    size = 4
    rows = 1
    cols = 2
    fig, (ax, axe1) = plt.subplots(1, 2, figsize=(cols * (size + 3), rows *
        size))
    ax = plot_saliency_map(heatmap, np_img, ax=ax)
    _ = ax.set_title(idx)
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    heatmap = heatmap / np.linalg.norm(heatmap)
    mask = heatmap.copy()
    mask[mask < threshold] = -5
    mask[mask >= threshold] = 5
    axe1.matshow(mask)
    axe1.set_title('thresh={}'.format(threshold), pad=20)
    plt.subplots_adjust(top=1.2)
    plt.show()
    alpha = torch.tensor(mask).to(dev)
    alpha = alpha.unsqueeze(0)
    alpha = alpha.unsqueeze(0)
    iba = IBA(eval('model.features.denseblock{0}'.format(*position)))
    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data,
        progbar=False)
    iba.reverse_lambda = new_method
    iba.beta = 6
    new_alpha = nn.functional.interpolate(alpha, iba.estimator.shape[1:])
    new_alpha = new_alpha.squeeze(0)
    new_alpha = new_alpha.expand(iba.estimator.shape[0], -1, -1).clone()
    iba.old_alpha = torch.nn.Parameter(new_alpha, requires_grad=True)
    print(
        'With old mask generated from IB after denseblock3 for new IB after denseblock '
         + position)
    if more_iter:
        fig, axes = plt.subplots(3, 4, figsize=(20, 15))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    else:
        fig, axes = plt.subplots(2, 4, figsize=(20, 10))
        iter_set = [1, 2, 4, 6, 8, 10, 12, 14]
        for iter_value, ax in zip(iter_set, axes.flatten()):
            heatmap = iba.analyze(img, model_loss_closure,
                optimization_steps=iter_value)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('optimization steps={}'.format(iter_value))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
    plt.show()


<mask token>
new_initit()
new_initit(position='1', threshold=0.003, index=0)
new_initit(position='1', threshold=0.003, index=1)
<mask token>
new_initit(more_iter=True, new_method=True)
new_initit(index=1, more_iter=True, new_method=True)
<mask token>
new_initit(position='2', threshold=0.006, index=0, more_iter=True)
new_initit(position='2', threshold=0.003, index=0)
<mask token>
new_initit(position='3', threshold=0.003, index=0)
<mask token>


def compare_layers(length, ind=0, beta=6, more_layer=False, new_method=False):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    if new_method:
        print('New method')
    else:
        print('Old method')
    if more_layer:
        size = 4
        rows = 4
        cols = 4
        fig, axes = plt.subplots(4, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8'], ['4', '9'], ['4',
            '10'], ['4', '11'], ['4', '12'], ['4', '13'], ['4', '14'], ['4',
            '15'], ['4', '16']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()
    else:
        size = 4
        rows = 2
        cols = 4
        fig, axes = plt.subplots(2, 4, figsize=(cols * (size + 1), rows * size)
            )
        positions = [['4', '1'], ['4', '2'], ['4', '3'], ['4', '4'], ['4',
            '5'], ['4', '6'], ['4', '7'], ['4', '8']]
        for position, ax in zip(positions, axes.flatten()):
            current_pos = 'model.features.denseblock{0}.denselayer{1}'.format(*
                position)
            iba = IBA(eval(current_pos))
            iba.reset_estimate()
            iba.estimate(model, dataloader, device=dev, n_samples=
                length_data, progbar=False)
            iba.reverse_lambda = new_method
            iba.beta = beta
            heatmap = iba.analyze(img, model_loss_closure)
            ax = plot_saliency_map(heatmap, np_img, ax=ax)
            ax.set_title('denseblock{0}.denselayer{1}'.format(*position))
            ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2],
                bbox[3], linewidth=2, edgecolor='r', facecolor='none',
                zorder=2))
        plt.show()


compare_layers(length_data, 0)
compare_layers(length_data, 1)
compare_layers(length_data, 2)
compare_layers(length_data, 0, 3)
compare_layers(length_data, 0, 0.5)
compare_layers(length_data, 0, 0.1, more_layer=True)
<mask token>


def norm_concat_heatmaps(index=0, weight=[1, 1, 1], new_method=False):
    if new_method:
        print('New method')
    else:
        print('Old method')
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    heatmap = []
    size = 4
    rows = 1
    cols = 3
    fig, axes = plt.subplots(1, 3, figsize=(cols * (size + 1), rows * size))
    for i in range(3):
        current_pos = 'model.features.denseblock{0}'.format(i + 1)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.reverse_lambda = new_method
        iba.beta = 6
        hm = iba.analyze(img, model_loss_closure)
        hm = hm / np.linalg.norm(hm)
        heatmap.append(hm)
        ax = axes[i]
        ax = plot_saliency_map(heatmap[i], np_img, ax=ax)
        ax.set_title('normalized heatmap from mask afte block ' + str(i + 1))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()
    heatmap = np.array(heatmap)
    weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1
        ] + heatmap[2] * weight[2]
    ax = plot_saliency_map(weighted_heatmap, np_img)
    _ = ax.set_title('weight=' + str(weight))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
        linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


<mask token>
norm_concat_heatmaps()
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
norm_concat_heatmaps(1)
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
norm_concat_heatmaps(3)
norm_concat_heatmaps(index=0, weight=[3, 2, 1])
<mask token>
norm_concat_heatmaps(new_method=True)
<mask token>


def compare_blocks(length, ind=0):
    index = min(ind, length)
    print('results for image ' + str(index))
    img, target, idx, bbox = dataset[index]
    img = img[None].to(dev)
    np_img = tensor_to_np_img(img[0])
    size = 4
    rows = 1
    cols = 4
    fig, axes = plt.subplots(1, 4, figsize=(cols * (size + 1), rows * size))
    positions = [['1'], ['2'], ['3'], ['4']]
    for position, ax in zip(positions, axes.flatten()):
        current_pos = 'model.features.denseblock{0}'.format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data,
            progbar=False)
        iba.beta = 6
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title(current_pos)
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
            linewidth=2, edgecolor='r', facecolor='none', zorder=2))
    plt.show()


compare_blocks(length_data, 0)
compare_blocks(length_data, 1)
compare_blocks(length_data, 2)
<mask token>
LABEL = 'Atelectasis'
length_data = samples_display(LABEL)
<mask token>
LABEL = 'Effusion'
length_data = samples_display(LABEL)
<mask token>
LABEL = 'Mass'
length_data = samples_display(LABEL)
<mask token>
LABEL = 'Nodule'
length_data = samples_display(LABEL)
<mask token>
LABEL = 'Pneumonia'
length_data = samples_display(LABEL)
<mask token>
LABEL = 'Pneumothorax'
length_data = samples_display(LABEL)

# -*- coding: utf-8 -*-
"""pytorch_IBA_inserted_NIH1_30.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1YbVMIdKCn6EPPR9wvI7OIUS3Bxn6jRpr

# PyTorch: IBA (Per-Sample Bottleneck)

This notebook shows how to apply the Per-Sample Bottleneck to pretrained ImageNet models. 

Ensure that `./imagenet` points to your copy of the ImageNet dataset. 

You might want to create a symlink:
"""

from google.colab import drive
import sys
drive.mount('/content/drive', force_remount=True)

sys.path.append('/content/drive/MyDrive/Prak_MLMI')
sys.path.append('/content/drive/MyDrive')
sys.path.append('/content/drive/MyDrive/Prak_MLMI/model')

import warnings
warnings.filterwarnings('ignore')

import matplotlib.patches as patches

# Commented out IPython magic to ensure Python compatibility.
# to set you cuda device
# %env CUDA_VISIBLE_DEVICES=0

# %reload_ext autoreload
# %autoreload 2

import torch
import torchvision.models 
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize, Normalize, Grayscale
import matplotlib.pyplot as plt 
import os
from tqdm import tqdm_notebook
import json
from PIL import Image
import numpy as np
# import sys

try:
    import IBA
except ModuleNotFoundError:
    sys.path.insert(0, '..')
    import IBA
  
from IBA.pytorch_new import IBA, tensor_to_np_img
from IBA.utils import plot_saliency_map

import cxr_dataset as CXR

# Commented out IPython magic to ensure Python compatibility.
# #!unzip "/content/drive/MyDrive/NIH_CXR14_Resized.zip" -d "/content/drive/MyDrive/NIH_CXR14_Resized"
# %%capture
# !unzip "/content/drive/MyDrive/NIH_CXR14_Resized.zip" -d /content/

# %load /content/drive/MyDrive/Prak_MLMI/IBA/pytorch_new.py

"""## Loading Data and Model"""

prak_dir = '/content/drive/MyDrive/Prak_MLMI'
imagenet_dir = '/content/drive/MyDrive/Prak_MLMI/imagenet'
PATH_TO_IMAGES = "/content/NIH small"
MODEL_PATH = prak_dir + '/model/results/checkpoint_best'
label_path = '/content/drive/MyDrive/Prak_MLMI/model/labels'

dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

checkpoint = torch.load(MODEL_PATH, map_location=dev)
model = checkpoint['model'].module
model.to(dev).eval()
print(model)

# load the data

# use imagenet mean,std for normalization
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

data_transforms = {
    'train': Compose([
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean, std)
    ]),
    'val': Compose([
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean, std)
    ]),
}


bounding_box_transform = CXR.RescaleBB(224, 1024)

def get_label(LA):
    labels = {   
        'Atelectasis': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Cardiomegaly': [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Effusion': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Infiltration': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Mass': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Nodule': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pneumonia': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pneumothorax': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Consolidation': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Edema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Emphysema': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Fibrosis': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Pleural_Thickening': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Hernia': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    }

    return labels.get(LA, None)

def model_loss_closure(input):
  loss = nn.BCEWithLogitsLoss()
  mse_loss = loss(model(input), torch.tensor(_label).view(1,-1).expand(10, -1).to(dev).float())
  return mse_loss

# Commented out IPython magic to ensure Python compatibility.
# %matplotlib inline

def samples_display(LABEL, beta=6):



  dataset = CXR.CXRDataset(
      path_to_images=PATH_TO_IMAGES,
      fold='BBox',#fold='train'
      transform=data_transforms['train'],
      transform_bb=bounding_box_transform,
      fine_tune=False,
      label_path=label_path,
      finding=LABEL)  #finding=LABEL
  dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
  global _label
  _label = get_label(LABEL)


  seed = np.random.randint(0, 1000000000)

  np.random.seed(seed)

  length = len(dataset)
  print("20 Ransom samples of "+LABEL+" from " + str(length) + " images")

  for sample_idx in np.random.choice(length, 20):  #change with dataset length
      iba = IBA(model.features.denseblock2)
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length, progbar=False) #change with dataset length
      img, target, idx, bbox = dataset[sample_idx]
      img = img[None].to(dev)
  
      # reverse the data pre-processing for plotting the original image
      np_img = tensor_to_np_img(img[0])

      size = 4
      rows = 1
      cols = 5
      fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize=(cols*(size+1), rows*size))


      cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
      mean = np.array([0.485, 0.456, 0.406])
      std = np.array([0.229, 0.224, 0.225])
      cxr = std * cxr + mean
      cxr = np.clip(cxr, 0, 1)

  #    rect_original = patches.Rectangle((bbox[0, 0], bbox[0, 1]), bbox[0, 2], bbox[0, 3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)
      rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)

      ax.imshow(cxr)
      ax.axis('off')
      ax.set_title(idx)
      ax.add_patch(rect_original) 


      iba.reverse_lambda = False
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax2 = plot_saliency_map(heatmap, np_img, ax=ax2)
      _ = ax2.set_title("Old method, Block2")
      ax2.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba.reverse_lambda = True
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax3 = plot_saliency_map(heatmap, np_img, ax=ax3)
      _ = ax3.set_title("New method, Block2")
      ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba = IBA(model.features.denseblock3)
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length, progbar=False)
      
      iba.reverse_lambda = False
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax4 = plot_saliency_map(heatmap, np_img, ax=ax4)
      _ = ax4.set_title("Old method, Block3")
      ax4.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))


      iba.reverse_lambda = True
      iba.beta = beta
      heatmap = iba.analyze(img, model_loss_closure) 
      # show the heatmap
      ax5 = plot_saliency_map(heatmap, np_img, ax=ax5)
      _ = ax5.set_title("New method, Block3")
      ax5.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))



      plt.show()
  return length

"""# Cardiomegaly"""

LABEL='Cardiomegaly'

length_data=samples_display(LABEL)

dataset = CXR.CXRDataset(
    path_to_images=PATH_TO_IMAGES,
    fold='BBox',#fold='train'
    transform=data_transforms['train'],
    transform_bb=bounding_box_transform,
    fine_tune=False,
    label_path=label_path,
    finding=LABEL)  #finding=LABEL
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
img, target, idx, bbox = dataset[0]

"""## Compare different betas (So I chose beta=6 instead of 10, the default)

###IB in denseblock3 with different betas
"""

iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""###IB in denseblock2 with different betas"""

iba = IBA(model.features.denseblock2)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""###IB in denseblock4 with different betas"""

iba = IBA(model.features.denseblock4)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=len(dataset), progbar=False)

size = 4
rows = 2
cols = 4
fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

beta_set = [0.5, 1, 2, 4, 6, 8, 10, 12]
for beta_value, ax in zip(beta_set, axes.flatten()) : 
    heatmap = iba.analyze(img[None].to(dev), model_loss_closure, beta = beta_value) 
    ax = plot_saliency_map(heatmap, tensor_to_np_img(img), ax=ax)
    ax.set_title("beta={}".format(beta_value))
plt.show()

"""## Compare New and Old method with IB in output of denseblock3 with another image as well as compare masked and unmasked prediction probability"""

img = img[None].to(dev)
np_img = tensor_to_np_img(img[0])

iba = IBA(model.features.denseblock3)
iba.reset_estimate()
iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)


print("Old method")
iba.reverse_lambda = False
iba.beta = 6
#model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
heatmap = iba.analyze(img, model_loss_closure)  
# show the heatmap
ax = plot_saliency_map(heatmap, np_img)
_ = ax.set_title(idx)
ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

# show the prediction probability
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1,-1).to(dev).float())
with iba.restrict_flow():
  print(torch.sigmoid(model(img)))

plt.show()



print("Ground truth with bounding box")
cxr = img.data.cpu().numpy().squeeze().transpose(1, 2, 0)
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
cxr = std * cxr + mean
cxr = np.clip(cxr, 0, 1)

rect_original = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2)

fig, ax2 = plt.subplots(1, 1, figsize=(5.5, 4.0))

ax2.imshow(cxr)
ax2.axis('off')
ax2.set_title(idx)
ax2.add_patch(rect_original) 
plt.show()

print("New method")
iba.reverse_lambda = True
iba.beta = 6
#model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
heatmap = iba.analyze(img, model_loss_closure) 
# show the heatmap
ax3 = plot_saliency_map(heatmap, np_img)
_ = ax3.set_title(idx)
ax3.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
# show the prediction probability
print(torch.sigmoid(model(img)))
print(torch.tensor(_label).view(1,-1).to(dev).float())
with iba.restrict_flow():
  print(torch.sigmoid(model(img)))
plt.show()

"""## New: Search for a proper threshold to generate mask"""

length_data = 146
_label = get_label(LABEL)

def compare_threshold(index=0, new_method = False):

  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])


  iba = IBA(model.features.denseblock3)
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

  if new_method: 
      print("New method")
  else:
      print("Old method")
      
  iba.reverse_lambda = new_method
  iba.beta = 6
  #model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()
  heatmap = iba.analyze(img, model_loss_closure)  
  # show the heatmap
  ax = plot_saliency_map(heatmap, np_img)
  _ = ax.set_title(idx)
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

  heatmap = heatmap / np.linalg.norm(heatmap)


  size = 4
  rows = 5
  cols = 5
  fig, axes = plt.subplots(5, 5, figsize=(cols*(size+1), rows*size))


  threshold = 0.0015
  for ax in axes.flatten():
    mask = heatmap.copy()
    mask[mask < threshold] = 0
    mask[mask >= threshold] = 1
    
    ax.matshow(mask)
    ax.set_title("thresh={}".format(threshold),pad = 20)
    threshold += 0.00025
    

  plt.subplots_adjust(top=1.2)    
  plt.show()

"""### old method"""

compare_threshold(0)

compare_threshold(1)

compare_threshold(2)

"""### New method"""

compare_threshold(index=0, new_method = True)

compare_threshold(index=1, new_method = True)

"""## New: Initialzation of IB after block1 with mask generated from IB after block3"""

def new_initit(position = "1", threshold = 0.006, index=0, more_iter = False, new_method = False):

  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  iba = IBA(model.features.denseblock3)
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

  if new_method: 
      print("New method")
  else:
      print("Old method")
      
  iba.reverse_lambda = new_method
  iba.beta = 6
  heatmap = iba.analyze(img, model_loss_closure)  

  size = 4
  rows = 1
  cols = 2
  fig, (ax, axe1) = plt.subplots(1, 2, figsize=(cols*(size+3), rows*size)) 

  # show the heatmap
  ax = plot_saliency_map(heatmap, np_img, ax = ax )
  _ = ax.set_title(idx)
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  heatmap = heatmap / np.linalg.norm(heatmap)

  mask = heatmap.copy()
  mask[mask < threshold] = -5
  mask[mask >= threshold] = 5    
  axe1.matshow(mask)
  axe1.set_title("thresh={}".format(threshold),pad = 20)
 

  plt.subplots_adjust(top=1.2)    
  plt.show()
  
  #alter and save the old mask
  alpha = torch.tensor(mask).to(dev)
  alpha = alpha.unsqueeze(0)
  alpha = alpha.unsqueeze(0) 

  iba = IBA(eval("model.features.denseblock{0}".format(*position)))
  iba.reset_estimate()
  iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)
  iba.reverse_lambda = new_method
  iba.beta = 6

  new_alpha = nn.functional.interpolate(alpha, iba.estimator.shape[1:])
  new_alpha = new_alpha.squeeze(0)
  new_alpha = new_alpha.expand(iba.estimator.shape[0],-1,-1).clone()

  iba.old_alpha= torch.nn.Parameter(new_alpha,  requires_grad=True) 
  print("With old mask generated from IB after denseblock3 for new IB after denseblock " + position)

  if more_iter:
    fig, axes = plt.subplots(3, 4, figsize=(20,15))

    iter_set = [ 1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
    for iter_value, ax in zip(iter_set, axes.flatten()) : 
        heatmap = iba.analyze(img, model_loss_closure, optimization_steps = iter_value) 
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title("optimization steps={}".format(iter_value))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
  else: 
    fig, axes = plt.subplots(2, 4, figsize=(20,10))

    iter_set = [ 1, 2, 4, 6, 8, 10, 12, 14]
    for iter_value, ax in zip(iter_set, axes.flatten()) : 
        heatmap = iba.analyze(img, model_loss_closure, optimization_steps = iter_value) 
        ax = plot_saliency_map(heatmap, np_img, ax=ax)
        ax.set_title("optimization steps={}".format(iter_value))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))



  plt.show()

"""### old method"""

new_initit()

new_initit(position = "1", threshold = 0.003, index=0)

new_initit(position = "1", threshold = 0.003, index=1)

"""### New method"""

new_initit(more_iter = True, new_method = True)

new_initit(index=1, more_iter = True, new_method = True)

"""## New: Initialzation of IB after block2 with mask generated from IB after block3"""

new_initit(position = "2", threshold = 0.006, index=0, more_iter = True)

new_initit(position = "2", threshold = 0.003, index=0)

"""## New: Initialzation of IB after block3 with mask generated from IB after block3"""

new_initit(position = "3", threshold = 0.003, index=0)

#wrong initialization, highlighted shoulder area
# new_initit(position = "1", shift = 2, index=0)

"""## New: Compare IB in different layers of Block4"""

def compare_layers(length, ind=0, beta = 6, more_layer = False, new_method = False):

  index = min (ind, length) 
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])
  if new_method: 
    print("New method")
  else:
    print("Old method")

  if more_layer:
    size = 4
    rows = 4
    cols = 4
    fig, axes = plt.subplots(4, 4, figsize=(cols*(size+1), rows*size))

    positions = [["4","1"], ["4","2"], ["4","3"],["4","4"],["4","5"],["4","6"],["4","7"],["4","8"],\
                 ["4","9"],["4","10"],["4","11"],["4","12"],["4","13"],["4","14"],["4","15"],["4","16"]]

    for position, ax in zip(positions, axes.flatten()):
        current_pos = "model.features.denseblock{0}.denselayer{1}".format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)    
        iba.reverse_lambda = new_method
        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)  
        ax.set_title("denseblock{0}.denselayer{1}".format(*position))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

    plt.show()

  else: 
    size = 4
    rows = 2
    cols = 4
    fig, axes = plt.subplots(2, 4, figsize=(cols*(size+1), rows*size))

    positions = [["4","1"], ["4","2"], ["4","3"],["4","4"],["4","5"],["4","6"],["4","7"],["4","8"]]

    for position, ax in zip(positions, axes.flatten()):
        current_pos = "model.features.denseblock{0}.denselayer{1}".format(*position)
        iba = IBA(eval(current_pos))
        iba.reset_estimate()
        iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)          
        iba.reverse_lambda = new_method

        iba.beta = beta
        heatmap = iba.analyze(img, model_loss_closure)
        ax = plot_saliency_map(heatmap, np_img, ax=ax)  
        ax.set_title("denseblock{0}.denselayer{1}".format(*position))
        ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

    plt.show()

compare_layers(length_data,0)
compare_layers(length_data,1)
compare_layers(length_data,2)

compare_layers(length_data,0, 3)
compare_layers(length_data,0, 0.5)

# beta=0.1, beta has to be extremely small (default beta from the papaer: 10)
compare_layers(length_data,0, 0.1, more_layer = True)

"""## New: concatenated mask"""

def norm_concat_heatmaps(index=0, weight = [1, 1, 1], new_method = False):  

# normalized_concatenated_heatmaps_generator
  if new_method: 
      print("New method")
  else:
      print("Old method")
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  heatmap = []

  size = 4
  rows = 1
  cols = 3
  fig, axes = plt.subplots(1, 3, figsize=(cols*(size+1), rows*size)) 

  # show the original heatmaps
  for i in range(3):

   
    current_pos = "model.features.denseblock{0}".format(i+1)
    iba = IBA(eval(current_pos))

    iba.reset_estimate()
    iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

    iba.reverse_lambda = new_method
    iba.beta = 6
    hm = iba.analyze(img, model_loss_closure)
    hm = hm / np.linalg.norm(hm)
    heatmap.append (hm)
    # heatmap.append (iba.analyze(img, model_loss_closure))

    ax = axes[i]
    ax = plot_saliency_map(heatmap[i], np_img, ax=ax)  
    ax.set_title("normalized heatmap from mask afte block "+ str(i+1))
    ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

  heatmap = np.array(heatmap) 
  weighted_heatmap = heatmap[0] * weight[0] + heatmap[1] * weight[1] + heatmap[2] * weight[2]
  ax = plot_saliency_map(weighted_heatmap, np_img)
  _ = ax.set_title("weight=" + str(weight))
  ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))
  plt.show()

"""### Old method"""

norm_concat_heatmaps()

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

norm_concat_heatmaps(1)

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

norm_concat_heatmaps(3)

norm_concat_heatmaps(index=0, weight = [3, 2, 1])

"""### New method"""

norm_concat_heatmaps(new_method = True)

"""## Old method (try with different blocks of DenseNet) with other images

"""

def compare_blocks(length, ind=0):

  index = min (ind, length) 
  print("results for image "+str(index))
  img, target, idx, bbox = dataset[index]

  img = img[None].to(dev)
  np_img = tensor_to_np_img(img[0])

  size = 4
  rows = 1
  cols = 4
  fig, axes = plt.subplots(1, 4, figsize=(cols*(size+1), rows*size))

  positions = [["1"], ["2"], ["3"],["4"]]

  for position, ax in zip(positions, axes.flatten()):
      current_pos = "model.features.denseblock{0}".format(*position)
      iba = IBA(eval(current_pos))
      iba.reset_estimate()
      iba.estimate(model, dataloader, device=dev, n_samples=length_data, progbar=False)

      iba.beta = 6
      heatmap = iba.analyze(img, model_loss_closure)
      ax = plot_saliency_map(heatmap, np_img, ax=ax)  
      ax.set_title(current_pos)
      ax.add_patch(patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none', zorder=2))

  plt.show()

compare_blocks(length_data,0)
compare_blocks(length_data,1)
compare_blocks(length_data,2)

"""# Atelectasis"""

LABEL='Atelectasis'

length_data=samples_display(LABEL)

"""# Effusion"""

LABEL='Effusion'
length_data=samples_display(LABEL)

"""# Mass"""

LABEL='Mass'
length_data=samples_display(LABEL)

"""# Nodule"""

LABEL='Nodule'
length_data=samples_display(LABEL)

"""# Pneumonia"""

LABEL='Pneumonia'
length_data=samples_display(LABEL)

"""# Pneumothorax"""

LABEL='Pneumothorax'
length_data=samples_display(LABEL)