!pip install medmnist
from tqdm import tqdm
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
import torchvision.transforms as transforms
import medmnist
from medmnist import INFO, Evaluator
import matplotlib.pyplot as plt
import scipy

## Using the code from https://github.com/MedMNIST/MedMNIST/blob/main/examples/getting_started.ipynb
data_flag = 'dermamnist'

## This defines our NN parameters
NUM_EPOCHS = 10
BATCH_SIZE = 128
lr = 0.001

info = INFO[data_flag]
task = info['task']
n_channels = info['n_channels']
n_classes = len(info['label'])

data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5], std=[.5])
])


DataClass = getattr(medmnist, info['python_class'])


# load the data
train_dataset = DataClass(split = 'train', transform = data_transform, download = True)
test_dataset  = DataClass(split = 'test' , transform = data_transform, download = True)
pil_dataset   = DataClass(split = 'train',                             download = True)

train_loader = data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
train_loader_at_eval = data.DataLoader(dataset=train_dataset, batch_size= 2 * BATCH_SIZE,  shuffle=False)
test_loader = data.DataLoader(dataset=test_dataset, batch_size=2*BATCH_SIZE, shuffle=False)

train_dataset.montage(length=20)

info['label']

train_targets = []
for i, t in train_loader:
    train_targets = np.append(train_targets, t.numpy())

import pandas as pd
pd.DataFrame({'target' : train_targets }).value_counts(normalize = True)

class Net(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(Net, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels, 16, kernel_size=3),
            nn.BatchNorm2d(16),
            nn.ReLU())

        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 16, kernel_size=3),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))

        self.layer3 = nn.Sequential(
            nn.Conv2d(16, 64, kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU())

        self.layer4 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU())

        self.layer5 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))

        self.fc = nn.Sequential(
            nn.Linear(64 * 4 * 4, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes))

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = Net(in_channels=n_channels, num_classes=n_classes)

criterion = nn.CrossEntropyLoss()

optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

for epoch in range(NUM_EPOCHS):
    train_correct = 0
    train_total = 0
    test_correct = 0
    test_total = 0

    model.train()
    for inputs, targets in train_loader:
        # forward + backward + optimize
        optimizer.zero_grad()
        outputs = model(inputs)

        if task == 'multi-label, binary-class':
            targets = targets.to(torch.float32)
            loss = criterion(outputs, targets)
        else:
            targets = targets.squeeze().long()
            loss = criterion(outputs, targets)

        loss.backward()
        optimizer.step()

from pathlib import Path

checkpoint_path = Path("book/.cache/convnet_example_dermamnist.pt")
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), checkpoint_path)
print(f"Saved checkpoint to {checkpoint_path}")

def test(split):
    model.eval()
    y_true = torch.tensor([])
    y_score = torch.tensor([])

    data_loader = train_loader_at_eval if split == 'train' else test_loader

    with torch.no_grad():
        for inputs, targets in data_loader:
            outputs = model(inputs)
            outputs = outputs.softmax(dim=-1)

            if task == 'multi-label, binary-class':
                targets = targets.to(torch.float32)
            else:
                targets = targets.squeeze().long()
                targets = targets.float().resize_(len(targets), 1)

            y_true = torch.cat((y_true, targets), 0)
            y_score = torch.cat((y_score, outputs), 0)

        y_true = y_true.numpy()
        y_score = y_score.detach().numpy()

        evaluator = Evaluator(data_flag, split)
        metrics = evaluator.evaluate(y_score)

        print('%s  auc: %.3f  acc:%.3f' % (split, *metrics))


print('==> Evaluating ...')
test('train')
test('test')

from pathlib import Path
import pandas as pd

import medmnist
import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision.transforms as transforms
from medmnist import INFO

data_flag = 'dermamnist'
info = INFO[data_flag]
n_channels = info['n_channels']
n_classes = len(info['label'])
batch_size = globals().get('BATCH_SIZE', 128)

data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5], std=[.5])
])

DataClass = getattr(medmnist, info['python_class'])
test_dataset = DataClass(split='test', transform=data_transform, download=True)
test_loader = data.DataLoader(dataset=test_dataset, batch_size=2 * batch_size, shuffle=False)

if 'Net' not in globals():
    class Net(nn.Module):
        def __init__(self, in_channels, num_classes):
            super(Net, self).__init__()

            self.layer1 = nn.Sequential(
                nn.Conv2d(in_channels, 16, kernel_size=3),
                nn.BatchNorm2d(16),
                nn.ReLU())

            self.layer2 = nn.Sequential(
                nn.Conv2d(16, 16, kernel_size=3),
                nn.BatchNorm2d(16),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2, stride=2))

            self.layer3 = nn.Sequential(
                nn.Conv2d(16, 64, kernel_size=3),
                nn.BatchNorm2d(64),
                nn.ReLU())

            self.layer4 = nn.Sequential(
                nn.Conv2d(64, 64, kernel_size=3),
                nn.BatchNorm2d(64),
                nn.ReLU())

            self.layer5 = nn.Sequential(
                nn.Conv2d(64, 64, kernel_size=3, padding=1),
                nn.BatchNorm2d(64),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2, stride=2))

            self.fc = nn.Sequential(
                nn.Linear(64 * 4 * 4, 128),
                nn.ReLU(),
                nn.Linear(128, 128),
                nn.ReLU(),
                nn.Linear(128, num_classes))

        def forward(self, x):
            x = self.layer1(x)
            x = self.layer2(x)
            x = self.layer3(x)
            x = self.layer4(x)
            x = self.layer5(x)
            x = x.view(x.size(0), -1)
            x = self.fc(x)
            return x

checkpoint_path = Path('book/.cache/convnet_example_dermamnist.pt')
if 'model' not in globals():
    model = Net(in_channels=n_channels, num_classes=n_classes)
    if checkpoint_path.exists():
        model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
    else:
        print(f'Checkpoint not found at {checkpoint_path}; using a fresh model.')

model.eval()
inputs, targets = next(iter(test_loader))

with torch.no_grad():
    outputs = model(inputs)
    outputs = outputs.softmax(dim=-1)

I, J = outputs.shape

plt.figure(figsize=(12, 3))
for i in range(8):
    plt.subplot(1,8,i+1)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(outputs.detach().numpy()[int(i * I / 8 ) : int((i + 1) * I / 8), :]);

batch = outputs.detach().numpy()


#grab 16 images to plot
indices = np.random.permutation(np.array(range(I)))[0 : 16]

## The associated
actual = targets.squeeze(-1).numpy()[indices]
target_pred = batch.argmax(axis = 1)[indices]
target_prob = np.round(batch.max(axis = 1) * 100)[indices]
images = inputs.numpy()[indices,:,:,:]

plt.figure(figsize=(10,10))

for i in range(16):
    plt.subplot(4,4,i+1)
    plt.xticks([])
    plt.yticks([])
    img = np.transpose(images[i,:,:,:], (1, 2, 0));
    img = ((img + 1)* 255/2).astype(np.uint8)
    plt.title("P="+str(target_prob[i])+",Yhat="+str(target_pred[i])+",Y="+ str(actual[i]))
    plt.imshow(img);

targets_pred, targets_actual = [], []

for i,t in test_loader:
    targets_pred.append(model(i).softmax(dim=-1))
    targets_actual.append(t)

targets_pred = torch.cat(targets_pred, dim = 0).detach().numpy().argmax(axis = 1)
targets_actual = torch.cat(targets_actual, dim=0).numpy().squeeze()


## This is the confusion  matrix data
confusion = pd.DataFrame({'y' : targets_actual, 'yhat' : targets_pred}).value_counts()
confusion.head(10)

### here's a list of the states that the model has kept track of
state_dictionary = list(model.state_dict())
state_dictionary[1 : 5]
### let's get the first round of weights
layer1_kernel = model.get_parameter('layer1.0.weight').detach().numpy()

layer1_kernel = layer1_kernel - layer1_kernel.min()
layer1_kernel = layer1_kernel / layer1_kernel.max()
layer1_kernel = layer1_kernel * 255

plt.figure(figsize=(5,5))
for h in range(16):
    plt.subplot(4,4,h+1)
    plt.xticks([])
    plt.yticks([])
    img = np.transpose(layer1_kernel[h,:,:,:],(1,2,0)).astype(np.uint8);
    plt.imshow(img);

l1 = model.layer1(inputs)

## plot one
images = l1.detach().numpy()

i = 0

plt.figure(figsize=(10,10))
for h in range(16):
    plt.subplot(4,4,h+1)
    plt.xticks([])
    plt.yticks([])
    img = images[i,h,:,:];
    plt.imshow(img);

# Re-obtain batch and SHAP values
batch_new = next(iter(test_loader))
images_new, _ = batch_new

background_new = images_new[:100]
test_images_new = images_new[100:105]

e_new = shap.DeepExplainer(model, background_new)
shap_values_new = e_new.shap_values(test_images_new)

def to_image_plot_format_new(values):
    if isinstance(values, list):
        return [np.moveaxis(v, 1, -1) for v in values]
    values = np.asarray(values)
    n_samples = test_images_new.shape[0]
    if values.shape[0] == n_samples:
        return [np.moveaxis(values[..., i], 1, -1) for i in range(values.shape[-1])]
    if values.shape[1] == n_samples:
        return [np.moveaxis(values[i], 1, -1) for i in range(values.shape[0])]
    raise ValueError(f"Could not interpret SHAP output shape: {values.shape}")

def min_max_scale_new(arr):
    min_val = arr.min()
    max_val = arr.max()
    if max_val == min_val:
        return np.zeros_like(arr)
    return (arr - min_val) / (max_val - min_val)

shap_numpy_raw = to_image_plot_format_new(shap_values_new)
test_numpy_original = np.moveaxis(test_images_new.cpu().numpy(), 1, -1)

# Scale the original images for background, but do NOT scale shap_numpy_raw
test_numpy_scaled_for_bg = min_max_scale_new(test_numpy_original)

# Plot the SHAP values using the raw shap_numpy_raw and the scaled original images as background
# The raw shap_numpy_raw allows shap.image_plot to correctly apply diverging red/blue colormap
shap.image_plot(shap_numpy_raw, test_numpy_scaled_for_bg)