| | import argparse |
| | import os |
| | import ruamel_yaml as yaml |
| | import numpy as np |
| | import json |
| | import torch |
| | import torch.nn as nn |
| | from torch.utils.data import DataLoader |
| | import torch.nn.functional as F |
| | from dataset.dataset_RSNA import RSNA2018_Dataset |
| | from models.model_MeDSLIP import MeDSLIP |
| | from models.tokenization_bert import BertTokenizer |
| | from sklearn.metrics import roc_auc_score, precision_recall_curve, accuracy_score |
| | from tqdm import tqdm |
| |
|
| | 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", |
| | ] |
| |
|
| |
|
| | 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 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 main(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") |
| |
|
| | |
| | print("Creating dataset") |
| | test_dataset = RSNA2018_Dataset(config["test_file"]) |
| | test_dataloader = DataLoader( |
| | test_dataset, |
| | batch_size=config["test_batch_size"], |
| | num_workers=8, |
| | pin_memory=True, |
| | sampler=None, |
| | shuffle=False, |
| | collate_fn=None, |
| | drop_last=False, |
| | ) |
| | json_book = json.load(open(config["disease_book"], "r")) |
| | disease_book = [json_book[i] for i in json_book] |
| | ana_book = [ |
| | "It is located at " + i |
| | for i in [ |
| | "trachea", |
| | "left_hilar", |
| | "right_hilar", |
| | "hilar_unspec", |
| | "left_pleural", |
| | "right_pleural", |
| | "pleural_unspec", |
| | "heart_size", |
| | "heart_border", |
| | "left_diaphragm", |
| | "right_diaphragm", |
| | "diaphragm_unspec", |
| | "retrocardiac", |
| | "lower_left_lobe", |
| | "upper_left_lobe", |
| | "lower_right_lobe", |
| | "middle_right_lobe", |
| | "upper_right_lobe", |
| | "left_lower_lung", |
| | "left_mid_lung", |
| | "left_upper_lung", |
| | "left_apical_lung", |
| | "left_lung_unspec", |
| | "right_lower_lung", |
| | "right_mid_lung", |
| | "right_upper_lung", |
| | "right_apical_lung", |
| | "right_lung_unspec", |
| | "lung_apices", |
| | "lung_bases", |
| | "left_costophrenic", |
| | "right_costophrenic", |
| | "costophrenic_unspec", |
| | "cardiophrenic_sulcus", |
| | "mediastinal", |
| | "spine", |
| | "clavicle", |
| | "rib", |
| | "stomach", |
| | "right_atrium", |
| | "right_ventricle", |
| | "aorta", |
| | "svc", |
| | "interstitium", |
| | "parenchymal", |
| | "cavoatrial_junction", |
| | "cardiopulmonary", |
| | "pulmonary", |
| | "lung_volumes", |
| | "unspecified", |
| | "other", |
| | ] |
| | ] |
| | tokenizer = BertTokenizer.from_pretrained(config["text_encoder"]) |
| | ana_book_tokenizer = get_tokenizer(tokenizer, ana_book).to(device) |
| | 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) |
| |
|
| | checkpoint = torch.load(args.checkpoint, map_location="cpu") |
| | state_dict = checkpoint["model"] |
| | model.load_state_dict(state_dict, strict=False) |
| | print("load checkpoint from %s" % args.checkpoint) |
| |
|
| | print("Start testing") |
| | model.eval() |
| |
|
| | gt = torch.FloatTensor() |
| | gt = gt.to(device) |
| | pred = torch.FloatTensor() |
| | pred = pred.to(device) |
| | loop = tqdm(test_dataloader) |
| | for i, sample in enumerate(loop): |
| | loop.set_description(f"Testing: {i+1}/{len(test_dataloader)}") |
| | images = sample["image"].to(device) |
| | labels = sample["label"].to(device) |
| | gt = torch.cat((gt, labels), 0) |
| |
|
| | with torch.no_grad(): |
| | pred_class = model(images) |
| | pred_class = pred_class[:, original_class.index("pneumonia"), :] |
| | pred_class = 1 - F.softmax(pred_class) |
| | pred = torch.cat((pred, pred_class), 0) |
| |
|
| | AUROC = compute_AUCs(gt, pred, 1) |
| | print("The AUROC of {} is {}".format("pneumonia", AUROC[0])) |
| | max_f1s = [] |
| | accs = [] |
| | gt_np = gt[:, 0].cpu().numpy() |
| | pred_np = pred[:, 0].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_RSNA/configs/MeDSLIP_config.yaml", |
| | ) |
| | parser.add_argument("--checkpoint", 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="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 |
| |
|
| | main(args, config) |
| |
|
