| import argparse |
| import os |
| import yaml as yaml |
| import numpy as np |
| import random |
| import time |
| import datetime |
| import json |
| from pathlib import Path |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| from sklearn.metrics import roc_auc_score, precision_recall_curve, accuracy_score |
|
|
| from models.model_MeDSLIP import MeDSLIP |
| from dataset.dataset import Chestxray14_Dataset |
| from models.tokenization_bert import BertTokenizer |
|
|
| from tqdm import tqdm |
|
|
| chexray14_cls = [ |
| "atelectasis", |
| "cardiomegaly", |
| "effusion", |
| "infiltrate", |
| "mass", |
| "nodule", |
| "pneumonia", |
| "pneumothorax", |
| "consolidation", |
| "edema", |
| "emphysema", |
| "tail_abnorm_obs", |
| "thicken", |
| "hernia", |
| ] |
|
|
| original_class = [ |
| "normal", |
| "clear", |
| "sharp", |
| "sharply", |
| "unremarkable", |
| "intact", |
| "stable", |
| "free", |
| "effusion", |
| "opacity", |
| "pneumothorax", |
| "edema", |
| "atelectasis", |
| "tube", |
| "consolidation", |
| "process", |
| "abnormality", |
| "enlarge", |
| "tip", |
| "low", |
| "pneumonia", |
| "line", |
| "congestion", |
| "catheter", |
| "cardiomegaly", |
| "fracture", |
| "air", |
| "tortuous", |
| "lead", |
| "disease", |
| "calcification", |
| "prominence", |
| "device", |
| "engorgement", |
| "picc", |
| "clip", |
| "elevation", |
| "expand", |
| "nodule", |
| "wire", |
| "fluid", |
| "degenerative", |
| "pacemaker", |
| "thicken", |
| "marking", |
| "scar", |
| "hyperinflate", |
| "blunt", |
| "loss", |
| "widen", |
| "collapse", |
| "density", |
| "emphysema", |
| "aerate", |
| "mass", |
| "crowd", |
| "infiltrate", |
| "obscure", |
| "deformity", |
| "hernia", |
| "drainage", |
| "distention", |
| "shift", |
| "stent", |
| "pressure", |
| "lesion", |
| "finding", |
| "borderline", |
| "hardware", |
| "dilation", |
| "chf", |
| "redistribution", |
| "aspiration", |
| "tail_abnorm_obs", |
| "excluded_obs", |
| ] |
|
|
| mapping = [] |
| for disease in chexray14_cls: |
| if disease in original_class: |
| mapping.append(original_class.index(disease)) |
| else: |
| mapping.append(-1) |
| MIMIC_mapping = [_ for i, _ in enumerate(mapping) if _ != -1] |
| chexray14_mapping = [i for i, _ in enumerate(mapping) if _ != -1] |
| target_class = [chexray14_cls[i] for i in chexray14_mapping] |
|
|
|
|
| def compute_AUCs(gt, pred, n_class): |
| """Computes Area Under the Curve (AUC) from prediction scores. |
| Args: |
| gt: Pytorch tensor on GPU, shape = [n_samples, n_classes] |
| true binary labels. |
| pred: Pytorch tensor on GPU, shape = [n_samples, n_classes] |
| can either be probability estimates of the positive class, |
| confidence values, or binary decisions. |
| Returns: |
| List of AUROCs of all classes. |
| """ |
| AUROCs = [] |
| gt_np = gt.cpu().numpy() |
| pred_np = pred.cpu().numpy() |
| for i in range(n_class): |
| AUROCs.append(roc_auc_score(gt_np[:, i], pred_np[:, i])) |
| return AUROCs |
|
|
|
|
| def get_tokenizer(tokenizer, target_text): |
|
|
| target_tokenizer = tokenizer( |
| list(target_text), |
| padding="max_length", |
| truncation=True, |
| max_length=64, |
| return_tensors="pt", |
| ) |
|
|
| return target_tokenizer |
|
|
|
|
| def test(args, config): |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print("Total CUDA devices: ", torch.cuda.device_count()) |
| torch.set_default_tensor_type("torch.FloatTensor") |
|
|
| test_dataset = Chestxray14_Dataset(config["test_file"], is_train=False) |
| test_dataloader = DataLoader( |
| test_dataset, |
| batch_size=config["test_batch_size"], |
| num_workers=30, |
| pin_memory=True, |
| sampler=None, |
| shuffle=True, |
| collate_fn=None, |
| drop_last=False, |
| ) |
|
|
| print("Creating book") |
| json_book = json.load(open(config["disease_book"], "r")) |
| disease_book = [json_book[i] for i in json_book] |
| tokenizer = BertTokenizer.from_pretrained(config["text_encoder"]) |
| disease_book_tokenizer = get_tokenizer(tokenizer, disease_book).to(device) |
|
|
| print("Creating model") |
| model = MeDSLIP(config, disease_book_tokenizer) |
| if args.ddp: |
| model = nn.DataParallel( |
| model, device_ids=[i for i in range(torch.cuda.device_count())] |
| ) |
| model = model.to(device) |
|
|
| print("Load model from checkpoint:", args.model_path) |
| checkpoint = torch.load(args.model_path, map_location="cpu") |
| state_dict = checkpoint["model"] |
| model.load_state_dict(state_dict) |
|
|
| |
| gt = torch.FloatTensor() |
| gt = gt.to(device) |
| pred = torch.FloatTensor() |
| pred = pred.to(device) |
|
|
| print("Start testing") |
| model.eval() |
| loop = tqdm(test_dataloader) |
| for i, sample in enumerate(loop): |
| loop.set_description(f"Testing: {i+1}/{len(test_dataloader)}") |
| image = sample["image"] |
| label = sample["label"][:, chexray14_mapping].float().to(device) |
| gt = torch.cat((gt, label), 0) |
| input_image = image.to(device, non_blocking=True) |
| with torch.no_grad(): |
| pred_class = model(input_image) |
| pred_class = F.softmax(pred_class.reshape(-1, 2)).reshape( |
| -1, len(original_class), 2 |
| ) |
| pred_class = pred_class[:, MIMIC_mapping, 1] |
| pred = torch.cat((pred, pred_class), 0) |
|
|
| AUROCs = compute_AUCs(gt, pred, len(target_class)) |
| AUROC_avg = np.array(AUROCs).mean() |
| print("The average AUROC is {AUROC_avg:.4f}".format(AUROC_avg=AUROC_avg)) |
| for i in range(len(target_class)): |
| print("The AUROC of {} is {}".format(target_class[i], AUROCs[i])) |
| max_f1s = [] |
| accs = [] |
| for i in range(len(target_class)): |
| gt_np = gt[:, i].cpu().numpy() |
| pred_np = pred[:, i].cpu().numpy() |
| precision, recall, thresholds = precision_recall_curve(gt_np, pred_np) |
| numerator = 2 * recall * precision |
| denom = recall + precision |
| f1_scores = np.divide( |
| numerator, denom, out=np.zeros_like(denom), where=(denom != 0) |
| ) |
| max_f1 = np.max(f1_scores) |
| max_f1_thresh = thresholds[np.argmax(f1_scores)] |
| max_f1s.append(max_f1) |
| accs.append(accuracy_score(gt_np, pred_np > max_f1_thresh)) |
|
|
| f1_avg = np.array(max_f1s).mean() |
| acc_avg = np.array(accs).mean() |
| print("The average f1 is {F1_avg:.4f}".format(F1_avg=f1_avg)) |
| print("The average ACC is {ACC_avg:.4f}".format(ACC_avg=acc_avg)) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--config", |
| default="Sample_zero-shot_Classification_CXR14/configs/MeDSLIP_config.yaml", |
| ) |
|
|
| parser.add_argument("--model_path", default="MeDSLIP_resnet50.pth") |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--gpu", type=str, default="0", help="gpu") |
| parser.add_argument("--ddp", action="store_true", help="whether to use ddp") |
| args = parser.parse_args() |
|
|
| config = yaml.load(open(args.config, "r"), Loader=yaml.Loader) |
|
|
| os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu |
| if args.gpu != "-1": |
| torch.cuda.current_device() |
| torch.cuda._initialized = True |
|
|
| test(args, config) |
|
|
