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): | |
| mean = [0.485, 0.456, 0.406] | |
| std = [0.229, 0.224, 0.225] | |
| 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): | |
| 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) | |