Spaces:
Runtime error
Runtime error
import os | |
import torch | |
import numpy as np | |
from PIL import Image | |
from utils import * | |
from tqdm import tqdm | |
from gradcam import GradCAM, GradCAMpp | |
from overlay_image import overlay_numpy | |
DISEASES = [ 'Atelectasis', | |
'Cardiomegaly', | |
'Effusion', | |
'Infiltration', | |
'Mass', | |
'Nodule', | |
'Pneumonia', | |
'Pneumothorax', | |
'Consolidation', | |
'Edema', | |
'Emphysema', | |
'Fibrosis', | |
'Pleural Thickening', | |
'Hernia' ] | |
class GradCamGenerator: | |
def __init__(self, model_path, layer, overlay=False): | |
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
self.model = self.load_model(model_path) | |
self.layer = layer | |
self.overlay = overlay # Overlay GC heatmaps with image | |
self.layer_module = self.model.get_submodule(layer) | |
self.gc_model2 = GradCAM(self.model, self.layer_module) # , device_ids=self.device) | |
def load_model(self, model_path, print_net=False): | |
checkpoint = torch.load(model_path, map_location=self.device) | |
model = checkpoint['model'] | |
self.set_inplace_False(model) | |
if print_net: | |
print(model) | |
return model | |
def set_inplace_False(self, module): | |
for layer in module._modules.values(): | |
if isinstance(layer, nn.ReLU): | |
layer.inplace = False | |
self.set_inplace_False(layer) | |
def generate_grad_cam(self, path): | |
img = self.pil_loader(path, 3) | |
input_image = self.transform_pil_to_tensor(img) | |
tclass = self.target_from_path(path) | |
#tmp_pred = self.model(input_image) | |
grayscale_cams = self.gc_model2(input=input_image, class_idx=tclass) | |
attribution = 255*grayscale_cams[0].detach().cpu().numpy().squeeze() | |
attribution /= attribution.max() | |
if self.overlay: | |
overlay_numpy(img, attribution, path) | |
#print() | |
return attribution | |
def target_from_path(self, path): | |
disease = path.split('/')[-2] | |
indx = DISEASES.index(disease) if disease!='No Finding' else 0 | |
return torch.tensor(indx, device=self.device) | |
def save_img(self, image, input_path): | |
gc_filename = input_path[:-4]+'_gc'+input_path[-4:] | |
image_PIL = Image.fromarray(image).convert('L') | |
image_PIL.save(gc_filename) | |
def transform_pil_to_tensor(self, pil_image): | |
# ImageNet mean and std | |
mean = [0.485, 0.456, 0.406] | |
std = [0.229, 0.224, 0.225] | |
# transformation to be applied | |
transform = transforms.Compose([ | |
transforms.Resize(224), | |
transforms.ToTensor(), | |
transforms.Normalize(mean, std) | |
]) | |
tensor = transform(pil_image).to(self.device) | |
return tensor.unsqueeze(0) | |
def pil_loader(self, path, n_channels): | |
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) | |
with open(path, 'rb') as f: | |
img = Image.open(f) | |
if n_channels == 1: | |
return img.convert('L') | |
elif n_channels == 3: | |
return img.convert('RGB') | |
else: | |
raise ValueError('Invalid value for parameter n_channels!') | |
def create_GC_from_folder(path, classifier='checkpoint', layer_name='features.norm5', overlay=True, override_gc=True): | |
GC = GradCamGenerator(classifier, layer_name, overlay=overlay) | |
folds = ['data/' + i for i in os.listdir(path) if 'No Finding' not in i][10:] | |
for cf in folds: | |
files = [cf+'/'+f for f in os.listdir(cf) if ('_gc' not in f and 'overlay' not in f and | |
(not os.path.exists(cf+'/'+f[:-4]+'_gc.png') or override_gc) and | |
(not os.path.exists(cf+'/'+f[:-4]+'_overlay.png') or not overlay))] | |
for cfil in tqdm(files): | |
GC.generate_grad_cam(cfil) | |
#GCmap = GC.generate_grad_cam(cfil) | |
#GC.save_img(GCmap, cfil) | |
if __name__=='__main__': | |
import matplotlib.pyplot as plt | |
path = 'data/' | |
create_GC_from_folder(path) | |