text_plain_output_1.png
from PIL import Image

from torch.utils.data import Dataset,DataLoader

from torchvision import transforms,datasets

import cv2

import matplotlib.pyplot as plt

import numpy as np

import torch

import torch.nn as nn

torch.manual_seed(42)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_transforms = transforms.Compose([transforms.Resize((224, 224)), transforms.Grayscale(), transforms.RandomHorizontalFlip(p=0.5), transforms.ToTensor()])
test_transforms = transforms.Compose([transforms.Resize((224, 224)), transforms.Grayscale(), transforms.ToTensor()])
dataset = datasets.ImageFolder('../input/brain-mri-images-for-brain-tumor-detection/brain_tumor_dataset', train_transforms)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

dataiter = iter(dataloader)
X,Y = dataiter.next()
classes = {0:'Negative',1:'Positive'}
fig,axes = plt.subplots(3,3,figsize=(14,14))
for i in range(3):
    for j in range(3):
        plt.sca(axes[i,j])
        idx = np.random.randint(0,31)
        image = np.moveaxis(X[idx].numpy(),0,2)
        plt.title(classes[Y[idx].item()])
        plt.imshow(image,cmap='gray')
        plt.axis('off');

def convBlock(ni, no):
    return nn.Sequential(nn.Dropout(0.2), nn.Conv2d(ni, no, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.BatchNorm2d(no), nn.MaxPool2d(2))

class BrainCancerClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(convBlock(1, 64), convBlock(64, 64), convBlock(64, 128), convBlock(128, 256), convBlock(256, 512), convBlock(512, 64), nn.Flatten(), nn.Linear(576, 256), nn.Dropout(0.3), nn.ReLU(inplace=True), nn.Linear(256, 128), nn.Dropout(0.2), nn.ReLU(inplace=True), nn.Linear(128, 64), nn.Dropout(0.2), nn.ReLU(inplace=True), nn.Linear(64, 2))
    def forward(self, x):
        return self.model(x)

def compute_metrics(preds, targets):
    loss = nn.CrossEntropyLoss()
    acc = (torch.max(preds, 1)[1] == targets).float().mean()
    return (loss(preds, targets), acc)

def train_batch(model, data, optimizer):
    model.train()
    images, labels = data
    images = images.to(device)
    labels = labels.to(device)
    optimizer.zero_grad()
    loss, acc = compute_metrics(model(images), labels)
    loss.backward()
    optimizer.step()
    return (loss.item(), acc.item())

model = BrainCancerClassifier().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

im2fmap = nn.Sequential(*list(model.model[:5].children()) + list(model.model[5][:2].children()))

def load_image(path):
    image = Image.open(path)
    image = test_transforms(image)
    image = image.unsqueeze(0)
    return image.to(device)

def get_heatmap(image_path):
    model.eval()
    for module in model.modules():
        if isinstance(module, nn.BatchNorm2d):
            module.track_running_stats = False
    image = load_image(image_path)
    logits = model(image)
    activations = im2fmap(image)
    prediction = logits.argmax(axis=1)
    model.zero_grad()
    logits[0, prediction].backward(retain_graph=True)
    pooled_grads = model.model[5][1].weight.grad.data.mean((0, 2, 3))
    for i in range(activations.shape[1]):
        activations[:, i, :, :] *= -pooled_grads[i]
    heatmap = torch.mean(activations, dim=1)[0].detach().to('cpu')
    return (heatmap, prediction.item())

size = 224
def upsampleHeatmap(map, img):
    m, M = (map.min(), map.max())
    map = 255 * ((map - m) / (M - m))
    map = np.uint8(map)
    map = cv2.resize(map, (size, size))
    map = cv2.applyColorMap(255 - map, cv2.COLORMAP_JET)
    map = np.uint8(map)
    map = np.uint8(map * 0.5 + img * 0.5)
    return map

image_paths = ['../input/brain-mri-images-for-brain-tumor-detection/no/19 no.jpg', '../input/brain-mri-images-for-brain-tumor-detection/yes/Y1.jpg']
cols = 2
fig, axes = plt.subplots(1, cols, figsize=(8, 10))
for i in range(cols):
    image = cv2.imread(image_paths[i])
    image = cv2.resize(image, (224, 224))
    heatmap, pred = get_heatmap(image_paths[i])
    result = upsampleHeatmap(heatmap, image)
    plt.sca(axes[i])
    plt.imshow(result, cmap='gray')
    plt.title(classes[pred])
    plt.axis('off')