# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="tDCyoAQIoeMd"
# # 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:

# + id="RWSJpsyKqHjH" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1611561628936, "user_tz": -480, "elapsed": 29596, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="3dff8ccc-c2a4-472b-9d29-13b6fcafbaba"
from google.colab import drive
import sys
drive.mount('/content/drive', force_remount=True)


# + id="CUpJJ2JIEycL" executionInfo={"status": "ok", "timestamp": 1611561628939, "user_tz": -480, "elapsed": 29582, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
sys.path.append('/content/drive/MyDrive/Prak_MLMI')
sys.path.append('/content/drive/MyDrive')
sys.path.append('/content/drive/MyDrive/Prak_MLMI/model')

# + id="SSHBVC1nI7bi" executionInfo={"status": "ok", "timestamp": 1611561628940, "user_tz": -480, "elapsed": 29577, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
import warnings
warnings.filterwarnings('ignore')

# + id="i3dkynIRac4_" executionInfo={"status": "ok", "timestamp": 1611561628941, "user_tz": -480, "elapsed": 29572, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
import matplotlib.patches as patches

# + id="mhGuVUtYoeMd" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1611561636594, "user_tz": -480, "elapsed": 37182, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="b4ccdcd6-f303-4283-83d6-d34b8994bc6c"
# to set you cuda device
# %env CUDA_VISIBLE_DEVICES=0

# %load_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 import IBA, tensor_to_np_img
from IBA.utils import plot_saliency_map

import cxr_dataset as CXR


# + id="XcxZDbsAIpJ7" executionInfo={"status": "ok", "timestamp": 1611561739634, "user_tz": -480, "elapsed": 140216, "user": {"displayName": "yezi Tsai", "photoUrl": "", "userId": "05768761827968141093"}}
# #!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/

# + [markdown] id="0I8XEzIyoeMe"
# ## Loading Data and Model

# + id="2hqke6AdoeMe" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1611561755689, "user_tz": -480, "elapsed": 156249, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="be2eb469-3e6c-4397-aa97-4c303d88666e"
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)

# + id="4ofQxeOwYoSa" executionInfo={"status": "ok", "timestamp": 1611561755690, "user_tz": -480, "elapsed": 156243, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
# 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)


 

# + id="0tVtsdaaO7pe" executionInfo={"status": "ok", "timestamp": 1611561755691, "user_tz": -480, "elapsed": 156234, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
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)


# + id="SJ460epLPRtD" executionInfo={"status": "ok", "timestamp": 1611561755692, "user_tz": -480, "elapsed": 156228, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
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


# + id="-yuq0D_1Pcw1" executionInfo={"status": "ok", "timestamp": 1611561755693, "user_tz": -480, "elapsed": 156222, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
# %matplotlib inline

# + id="eEoQ2PZVRJpb" executionInfo={"status": "ok", "timestamp": 1611561755693, "user_tz": -480, "elapsed": 155926, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}}
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, 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      


# + [markdown] id="Nc9mbi2DTkUV"
# # Cardiomegaly

# + id="002jwG_gRkZl" colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1Mp5jRfNWIl_z9RHW_GES7Stbhd5-GqKs"} executionInfo={"status": "ok", "timestamp": 1611551723772, "user_tz": -480, "elapsed": 191339, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="a005e30e-78e1-4bce-8b2b-11285f82589a"
LABEL='Cardiomegaly'

length_data=samples_display(LABEL)

# + id="JqzCJWq17dYm"
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]

# + [markdown] id="y30THSog-_AO"
# ## Compare different betas (So I chose beta=6 instead of 10, the default)

# + colab={"base_uri": "https://localhost:8080/", "height": 457} id="Vy7rhjtJ73-C" executionInfo={"status": "ok", "timestamp": 1611551969552, "user_tz": -480, "elapsed": 14246, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="fc18f782-35c7-4fa9-c79e-3ade35971a81"
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, 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()

# + [markdown] id="X5_lYEAUCJI_"
# ## Compare New and Old method with IB in output of denseblock3 with another image as well as compare masked and unmasked prediction probability

# + id="W73KiRZyhWq-" colab={"base_uri": "https://localhost:8080/", "height": 1000} executionInfo={"status": "ok", "timestamp": 1611550256069, "user_tz": -480, "elapsed": 6959, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="b77c781b-6a1c-4179-e714-8ceaafcae67d"
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, 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)

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()


# + [markdown] id="0eytuLQd2CHm"
# ## Old method (try with different blocks of DenseNet) with other images
#

# + id="tl0KUzfgblO_"
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, 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()


# + id="XFBEflqmcpeR" colab={"base_uri": "https://localhost:8080/", "height": 675} executionInfo={"status": "ok", "timestamp": 1611550303799, "user_tz": -480, "elapsed": 43254, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="2ef5dc73-2c5b-42cb-9056-09f5a7af935a"
compare_blocks(length_data,0)
compare_blocks(length_data,1)
compare_blocks(length_data,2)

# + [markdown] id="KA60-q--dqtC"
# # Atelectasis

# + id="tXc9ehk3dpjf" colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1EJHgN-p8qgzrXhypFHZ450tlNoUKo5ll"} executionInfo={"status": "ok", "timestamp": 1611550512835, "user_tz": -480, "elapsed": 209022, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="63df8f4e-89cc-401e-e2c9-7239140a723c"
LABEL='Atelectasis'

length_data=samples_display(LABEL)

# + [markdown] id="yYZIH6iD5bx1"
# # Effusion

# + id="zd027g-Udy-W" colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1yDgcop-DHApdFD7nxR8X-T9hSNn84I4d"} executionInfo={"status": "ok", "timestamp": 1611550751943, "user_tz": -480, "elapsed": 192682, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="c42c4537-5bf1-4c4e-f4c2-ac05b3f8616a"
LABEL='Effusion'
length_data=samples_display(LABEL)

# + [markdown] id="lP73hTO55jPV"
# # Mass

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1nQuBo7eUGorFRdfOU0HXQQtmEIO7Cohq"} id="AwIyCv8r5gvR" executionInfo={"status": "ok", "timestamp": 1611551028748, "user_tz": -480, "elapsed": 157033, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="500ec248-94e4-4d80-d6bf-264654ed22d4"
LABEL='Mass'
length_data=samples_display(LABEL)

# + [markdown] id="TuQQeiJ254po"
# # Nodule

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1zDTFxondbCXaaXDLCw8kSvptmhbvJqt2"} id="4Qo-NPIS6JNA" executionInfo={"status": "ok", "timestamp": 1611551183483, "user_tz": -480, "elapsed": 154664, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="f86d7760-f48b-4b72-bbc9-49c0f39f44b7"
LABEL='Nodule'
length_data=samples_display(LABEL)

# + [markdown] id="eQ71Kiz5jykc"
# # Pneumonia

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1dVYPcL0T1rgjKmqnuBgoJ8Vw32qo0U85"} id="S88RIup-kg2N" executionInfo={"status": "ok", "timestamp": 1611562042654, "user_tz": -480, "elapsed": 182889, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="fa7d8846-87e8-4973-9518-29b0db1f25ce"
LABEL='Pneumonia'
length_data=samples_display(LABEL)

# + [markdown] id="Qvs_ODWMjyV7"
# # Pneumothorax

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "output_embedded_package_id": "1f-5DJoK1LZ86LnclfRobr52CDGs_PDyr"} id="JMJRm3MHjxJG" executionInfo={"status": "ok", "timestamp": 1611562210050, "user_tz": -480, "elapsed": 348131, "user": {"displayName": "<NAME>", "photoUrl": "", "userId": "05768761827968141093"}} outputId="1c530e6b-3d72-4966-fba7-7218133a9a54"
LABEL='Pneumothorax'
length_data=samples_display(LABEL)

StarcoderdataJupyterScriptsDedupFiltered