import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.model_mil import MIL_fc, MIL_fc_mc
from models.model_clam import CLAM
from models.model_attention_mil import MIL_Attention_fc
import pdb
import os
import pandas as pd
from utils.utils import *
from utils.core_utils import EarlyStopping, Accuracy_Logger
from utils.file_utils import save_pkl, load_pkl
from sklearn.metrics import roc_auc_score, roc_curve, auc
import h5py
from models.resnet_custom import resnet50_baseline
import math
from sklearn.preprocessing import label_binarize
def initiate_model(args, ckpt_path=None):
print('Init Model')
model_dict = {"dropout": args.drop_out, 'n_classes': args.n_classes}
if args.model_size is not None and args.model_type in ['clam', 'attention_mil', 'clam_new']:
model_dict.update({"size_arg": args.model_size})
if args.model_type =='clam':
model = CLAM(**model_dict)
elif args.model_type == 'attention_mil':
model = MIL_Attention_fc(**model_dict)
else: # args.model_type == 'mil'
if args.n_classes > 2:
model = MIL_fc_mc(**model_dict)
else:
model = MIL_fc(**model_dict)
model.relocate()
#print_network(model)
if ckpt_path is not None:
ckpt = torch.load(ckpt_path)
model.load_state_dict(ckpt, strict=False)
model.eval()
return model
def eval(dataset, args, ckpt_path):
model = initiate_model(args, ckpt_path)
print('Init Loaders')
loader = get_simple_loader(dataset)
patient_results, test_error, auc, aucs, df, _ = summary(model, loader, args)
print('test_error: ', test_error)
print('auc: ', auc)
for cls_idx in range(len(aucs)):
print('class {} auc: {}'.format(cls_idx, aucs[cls_idx]))
return model, patient_results, test_error, auc, aucs, df
def infer(dataset, args, ckpt_path, class_labels):
model = initiate_model(args, ckpt_path)
df = infer_dataset(model, dataset, args, class_labels)
return model, df
# Code taken from pytorch/examples for evaluating topk classification on on ImageNet
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(1.0 / batch_size))
return res
def summary(model, loader, args):
acc_logger = Accuracy_Logger(n_classes=args.n_classes)
model.eval()
test_loss = 0.
test_error = 0.
all_probs = np.zeros((len(loader), args.n_classes))
all_labels = np.zeros(len(loader))
all_preds = np.zeros(len(loader))
slide_ids = loader.dataset.slide_data['slide_id']
patient_results = {}
for batch_idx, (data, label) in enumerate(loader):
data, label = data.to(device), label.to(device)
slide_id = slide_ids.iloc[batch_idx]
with torch.no_grad():
logits, Y_prob, Y_hat, _, results_dict = model(data)
acc_logger.log(Y_hat, label)
probs = Y_prob.cpu().numpy()
all_probs[batch_idx] = probs
all_labels[batch_idx] = label.item()
all_preds[batch_idx] = Y_hat.item()
patient_results.update({slide_id: {'slide_id': np.array(slide_id), 'prob': probs, 'label': label.item()}})
error = calculate_error(Y_hat, label)
test_error += error
del data
test_error /= len(loader)
if args.n_classes > 2:
acc1, acc3 = accuracy(torch.from_numpy(all_probs), torch.from_numpy(all_labels), topk=(1, 3))
print('top1 acc: {:.3f}, top3 acc: {:.3f}'.format(acc1.item(), acc3.item()))
if len(np.unique(all_labels)) == 1:
auc_score = -1
else:
if args.n_classes == 2:
auc_score = roc_auc_score(all_labels, all_probs[:, 1])
aucs = []
else:
aucs = []
binary_labels = label_binarize(all_labels, classes=[i for i in range(args.n_classes)])
for class_idx in range(args.n_classes):
if class_idx in all_labels:
fpr, tpr, _ = roc_curve(binary_labels[:, class_idx], all_probs[:, class_idx])
aucs.append(auc(fpr, tpr))
else:
aucs.append(float('nan'))
if args.micro_average:
binary_labels = label_binarize(all_labels, classes=[i for i in range(args.n_classes)])
fpr, tpr, _ = roc_curve(binary_labels.ravel(), all_probs.ravel())
auc_score = auc(fpr, tpr)
else:
auc_score = np.nanmean(np.array(aucs))
results_dict = {'slide_id': slide_ids, 'Y': all_labels, 'Y_hat': all_preds}
for c in range(args.n_classes):
results_dict.update({'p_{}'.format(c): all_probs[:,c]})
df = pd.DataFrame(results_dict)
return patient_results, test_error, auc_score, aucs, df, acc_logger
def infer_dataset(model, dataset, args, class_labels, k=3):
model.eval()
all_probs = np.zeros((len(dataset), k))
all_preds = np.zeros((len(dataset), k))
all_preds_str = np.full((len(dataset), k), ' ', dtype=object)
slide_ids = dataset.slide_data
for batch_idx, data in enumerate(dataset):
data = data.to(device)
with torch.no_grad():
logits, Y_prob, Y_hat, _, results_dict = model(data)
probs, ids = torch.topk(Y_prob, k)
probs = probs.cpu().numpy()
ids = ids.cpu().numpy()
all_probs[batch_idx] = probs
all_preds[batch_idx] = ids
all_preds_str[batch_idx] = np.array(class_labels)[ids]
del data
results_dict = {'slide_id': slide_ids}
for c in range(k):
results_dict.update({'Pred_{}'.format(c): all_preds_str[:, c]})
results_dict.update({'p_{}'.format(c): all_probs[:, c]})
df = pd.DataFrame(results_dict)
return df
# def infer_dataset(model, dataset, args, class_labels, k=3):
# model.eval()
# all_probs = np.zeros((len(dataset), args.n_classes))
# all_preds = np.zeros(len(dataset))
# all_str_preds = np.full(len(dataset), ' ', dtype=object)
# slide_ids = dataset.slide_data
# for batch_idx, data in enumerate(dataset):
# data = data.to(device)
# with torch.no_grad():
# logits, Y_prob, Y_hat, _, results_dict = model(data)
# probs = Y_prob.cpu().numpy()
# all_probs[batch_idx] = probs
# all_preds[batch_idx] = Y_hat.item()
# all_str_preds[batch_idx] = class_labels[Y_hat.item()]
# del data
# results_dict = {'slide_id': slide_ids, 'Prediction': all_str_preds, 'Y_hat': all_preds}
# for c in range(args.n_classes):
# results_dict.update({'p_{}_{}'.format(c, class_labels[c]): all_probs[:,c]})
# df = pd.DataFrame(results_dict)
# return df
def compute_features(dataset, args, ckpt_path, save_dir, model=None, feature_dim=512):
if model is None:
model = initiate_model(args, ckpt_path)
names = dataset.get_list(np.arange(len(dataset))).values
file_path = os.path.join(save_dir, 'features.h5')
initialize_features_hdf5_file(file_path, len(dataset), feature_dim=feature_dim, names=names)
for i in range(len(dataset)):
print("Progress: {}/{}".format(i, len(dataset)))
save_features(dataset, i, model, args, file_path)
def save_features(dataset, idx, model, args, save_file_path):
name = dataset.get_list(idx)
print(name)
features, label = dataset[idx]
features = features.to(device)
with torch.no_grad():
if type(model) == CLAM:
_, Y_prob, Y_hat, _, results_dict = model(features, instance_eval=False, return_features=True)
bag_feat = results_dict['features'][Y_hat.item()]
else:
_, Y_prob, Y_hat, _, results_dict = model(features, return_features=True)
bag_feat = results_dict['features']
del features
Y_hat = Y_hat.item()
Y_prob = Y_prob.view(-1).cpu().numpy()
bag_feat = bag_feat.view(1, -1).cpu().numpy()
with h5py.File(save_file_path, 'r+') as file:
print('label', label)
file['features'][idx, :] = bag_feat
file['label'][idx] = label
file['Y_hat'][idx] = Y_hat
file['Y_prob'][idx] = Y_prob[Y_hat]
def initialize_features_hdf5_file(file_path, length, feature_dim=512, names = None):
file = h5py.File(file_path, "w")
dset = file.create_dataset('features',
shape=(length, feature_dim), chunks=(1, feature_dim), dtype=np.float32)
# if names is not None:
# names = np.array(names, dtype='S')
# dset.attrs['names'] = names
if names is not None:
dt = h5py.string_dtype()
label_dset = file.create_dataset('names',
shape=(length, ), chunks=(1, ), dtype=dt)
label_dset = file.create_dataset('label',
shape=(length, ), chunks=(1, ), dtype=np.int32)
pred_dset = file.create_dataset('Y_hat',
shape=(length, ), chunks=(1, ), dtype=np.int32)
prob_dset = file.create_dataset('Y_prob',
shape=(length, ), chunks=(1, ), dtype=np.float32)
file.close()
return file_path
def eval2(datasets: tuple, cur: int, args: Namespace):
"""
train for a single fold
"""
print('\nTraining Fold {}!'.format(cur))
writer_dir = os.path.join(args.results_dir, str(cur))
if not os.path.isdir(writer_dir):
os.mkdir(writer_dir)
if args.log_data:
from tensorboardX import SummaryWriter
writer = SummaryWriter(writer_dir, flush_secs=15)
else:
writer = None
if args.pretrain_VAE:
print("Initializing VAE")
VAE = GenomicVAE(input_dim=args.omic_input_dim, hidden=[1024, 256, 128])
ckpt = torch.load('./VAE/logs/tcga_base/000-all/%d/%d/%d_best.ckpt' % (cur, cur, cur))
state_dict = ckpt['state_dict']
state_dict = OrderedDict((k[6:], v) for k, v in state_dict.items())
VAE.load_state_dict(state_dict)
args.omic_input_dim = 128
VAE.relocate()
dfs_freeze(VAE)
VAE.eval()
else:
VAE = None
print('\nInit train/val/test splits...', end=' ')
train_split, val_split, test_split = datasets
save_splits(datasets, ['train', 'val', 'test'], os.path.join(args.results_dir, 'splits_{}.csv'.format(cur)))
print('Done!')
print("Training on {} samples".format(len(train_split)))
print("Validating on {} samples".format(len(val_split)))
print("Testing on {} samples".format(len(test_split)))
print('\nInit loss function...', end=' ')
if args.task_type == 'survival':
if args.bag_loss == 'ce_surv':
loss_fn = CrossEntropySurvLoss(alpha=args.alpha_surv)
elif args.bag_loss == 'nll_surv':
loss_fn = NLLSurvLoss(alpha=args.alpha_surv)
elif args.bag_loss == 'cox_surv':
loss_fn = CoxSurvLoss()
else:
raise NotImplementedError
else:
if args.bag_loss == 'svm':
from topk import SmoothTop1SVM
loss_fn = SmoothTop1SVM(n_classes = args.n_classes)
if device.type == 'cuda':
loss_fn = loss_fn.cuda()
elif args.bag_loss == 'ce':
loss_fn = nn.CrossEntropyLoss()
else:
raise NotImplementedError
if args.reg_type == 'omic':
reg_fn = l1_reg_all
elif args.reg_type == 'pathomic':
reg_fn = l1_reg_modules
else:
reg_fn = None
print('Done!')
print('\nInit Model...', end=' ')
model_dict = {"dropout": args.drop_out, 'n_classes': args.n_classes}
if args.model_type in ['clam', 'clam_simple'] and args.subtyping:
model_dict.update({'subtyping': True})
if args.model_size is not None:
model_dict.update({"size_arg": args.model_size})
if args.model_type in ['clam', 'clam_simple']:
if args.task_type == 'survival':
raise NotImplementedError
else:
if args.inst_loss == 'svm':
from topk import SmoothTop1SVM
instance_loss_fn = SmoothTop1SVM(n_classes = 2)
if device.type == 'cuda':
instance_loss_fn = instance_loss_fn.cuda()
else:
instance_loss_fn = nn.CrossEntropyLoss()
if args.model_type =='clam':
model = CLAM(**model_dict, instance_loss_fn=instance_loss_fn)
else:
model = CLAM_Simple(**model_dict, instance_loss_fn=instance_loss_fn)
elif args.model_type =='attention_mil':
if args.task_type == 'survival':
model = MIL_Attention_fc_surv(**model_dict)
# model.alpha.requires_grad = False
else:
model = MIL_Attention_fc(**model_dict)
elif args.model_type =='mm_attention_mil':
model_dict.update({'input_dim': args.omic_input_dim, 'meta_dim': args.meta_dim,
'fusion': args.fusion, 'model_size_wsi':args.model_size_wsi, 'model_size_omic':args.model_size_omic,
'gate_path': args.gate_path, 'gate_omic': args.gate_omic, 'n_classes': args.n_classes,
'pretrain': args.pretrain, 'tcga_proj': '_'.join(args.task.split('_')[:2]), 'split_idx': cur})
if args.task_type == 'survival':
model = MM_MIL_Attention_fc_surv(**model_dict)
# model.alpha.requires_grad = False
else:
model = MM_MIL_Attention_fc(**model_dict)
elif args.model_type =='max_net':
model_dict = {'input_dim': args.omic_input_dim, 'meta_dim': args.meta_dim, 'model_size_omic': args.model_size_omic, 'n_classes': args.n_classes}
if args.task_type == 'survival':
model = MaxNet(**model_dict)
# model.alpha.requires_grad = False
else:
raise NotImplementedError
else: # args.model_type == 'mil'
if args.task_type == 'survival':
raise NotImplementedError
else:
if args.n_classes > 2:
model = MIL_fc_mc(**model_dict)
else:
model = MIL_fc(**model_dict)
model.relocate()
print('Done!')
print_network(model)
ckpt = torch.load(os.path.join(args.results_dir, "s_{}_checkpoint.pt".format(cur)))
model.load_state_dict(ckpt, strict=False)
model.eval()
print('\nInit Loaders...', end=' ')
train_loader = get_split_loader(train_split, training=True, testing = args.testing,
weighted = args.weighted_sample, task_type=args.task_type, batch_size=args.batch_size)
val_loader = get_split_loader(val_split, testing = args.testing, task_type=args.task_type, batch_size=args.batch_size)
test_loader = get_split_loader(test_split, testing = args.testing, task_type=args.task_type, batch_size=args.batch_size)
print('Done!')
if args.task_type == 'survival':
results_val_dict, val_c_index = summary_survival(model, val_loader, args.n_classes, VAE)
print('Val c-index: {:.4f}'.format(val_c_index))
results_test_dict, test_c_index = summary_survival(model, test_loader, args.n_classes, VAE)
print('Test c-index: {:.4f}'.format(test_c_index))
if writer:
writer.add_scalar('final/val_c_index', val_c_index, 0)
writer.add_scalar('final/test_c_index', test_c_index, 0)
writer.close()
return results_val_dict, results_test_dict, val_c_index, test_c_index
elif args.task_type == 'classification':
pass
Datasets
Models
