import pandas as pd
from mil_data_generator import *
from mil_models_pytorch import*
from mil_trainer_torch import *
from sklearn.utils.class_weight import compute_class_weight
import torch
from torchvision import transforms
import argparse
np.random.seed(42)
random.seed(42)
torch.manual_seed(42)
def main(args):
# INPUTS #
dir_images = '../data/SICAP_MIL/patches/'
dir_data_frame = '../data/SICAP_MIL/dataframes/gt_global_slides.xlsx'
dir_data_frame_test = '../data/SICAP_MIL/dataframes/gt_test_patches.xlsx'
dir_experiment = '../data/results/' + args.experiment_name + '/'
classes = ['G3', 'G4', 'G5']
proportions = ['pG3', 'pG4', 'pG5']
input_shape = (3, 224, 224)
images_on_ram = True
data_augmentation = True
pMIL = False
aggregation = 'max' # 'max', 'mean', 'attentionMIL', 'mcAttentionMIL'
mode = 'instance' # 'embedding', 'instance', 'mixed'
include_background = True
iterations = 3
df = pd.read_excel(dir_data_frame)
metrics = []
for ii_iteration in np.arange(0, iterations):
# Set data generators
dataset_train = MILDataset(dir_images, df[df['Partition'] == 'train'], classes, bag_id='slide_name',
input_shape=input_shape, data_augmentation=False, images_on_ram=images_on_ram,
pMIL=pMIL, proportions=proportions)
data_generator_train = MILDataGenerator(dataset_train, batch_size=1, shuffle=True, max_instances=512)
dataset_test = MILDataset(dir_images, df[df['Partition'] == 'test'], classes, bag_id='slide_name',
input_shape=input_shape, data_augmentation=False, images_on_ram=images_on_ram,
pMIL=pMIL, proportions=proportions, dataframe_instances=pd.read_excel(dir_data_frame_test))
data_generator_test = MILDataGenerator(dataset_test, batch_size=1, shuffle=False, max_instances=512)
# Test at instance level
X_test = data_generator_test.dataset.X[data_generator_test.dataset.y_instances[:, 0] != -1, :, :, :]
Y_test = data_generator_test.dataset.y_instances[data_generator_test.dataset.y_instances[:, 0] != -1, :]
# Load network
network = torch.load(dir_experiment + str(ii_iteration) + '_network_weights_best.pth')
# Pseudolabels on training set
labels = []
yhat_one_hot = []
Yglobal = data_generator_train.dataset.Yglobal
X = data_generator_train.dataset.X
for i in np.arange(0, X.shape[0]):
print(str(i + 1) + '/' + str(X.shape[0]), end='\r')
# Tensorize input
x = torch.tensor(X[i, :, :, :]).cuda().float()
x = x.unsqueeze(0)
features = network.bb(x)
yhat = torch.softmax(network.classifier(torch.squeeze(features)), 0)
yhat = yhat.detach().cpu().numpy()
yhat_one_hot.append(yhat)
if np.max(Yglobal[i, 1:]) == 0:
labels.append(0)
else:
if np.argmax(yhat) > 0:
if Yglobal[i, np.argmax(yhat)] == 1 and yhat[np.argmax(yhat)] > 0.5:
labels.append(np.argmax(yhat))
else:
labels.append(10)
else:
labels.append(10)
labels = np.array(labels)
yhat_one_hot = np.array(yhat_one_hot)
X = X[labels != 10, :, :, :]
Y = labels[labels != 10]
images_id = np.array(dataset_train.images)[labels != 10]
class_weights = compute_class_weight('balanced', [0, 1, 2, 3], Y)
# Set student network architecture
lr = 1e-2
network = MILArchitecture(classes, mode=mode, aggregation=aggregation, backbone='vgg19',
include_background=include_background).cuda()
opt = torch.optim.SGD(network.parameters(), lr=lr)
tranf = torch.nn.Sequential(transforms.RandomHorizontalFlip(),
transforms.RandomRotation(degrees=(-45, 45)),
transforms.GaussianBlur(3, sigma=(0.1, 2.0)),
transforms.ColorJitter(brightness=.5, hue=.3)).cuda()
training_data = CustomImageDataset(X, Y, transform=False)
train_dataloader = CustomGenerator(training_data, bs=32, shuffle=True)
def test_instances(X, Y, network, dir_out, i_iteration):
network.eval()
Yhat = []
for i in np.arange(0, X.shape[0]):
print(str(i + 1) + '/' + str(X.shape[0]), end='\r')
# Tensorize input
x = torch.tensor(X[i, :, :, :]).cuda().float()
x = x.unsqueeze(0)
features = network.bb(x)
yhat = torch.softmax(network.classifier(torch.squeeze(features)), 0)
yhat = torch.argmax(yhat).detach().cpu().numpy()
Yhat.append(yhat)
Yhat = np.array(Yhat)
Y = np.argmax(Y, 1)
cr = classification_report(Y, Yhat, target_names=['NC'] + classes, digits=4)
cm = confusion_matrix(Y, Yhat)
k2 = cohen_kappa_score(Y, Yhat, weights='quadratic')
f = open(dir_out + str(i_iteration) + '_report_student.txt', 'w')
f.write('Title\n\nClassification Report\n\n{}\n\nConfusion Matrix\n\n{}\n\nKappa\n\n{}\n'.format(cr, cm, k2))
f.close()
return k2
# STUDENT TRAINING
epochs = 60
dropout_rate = 0.2
for i_epoch in np.arange(0, epochs):
l_epoch = 0
if (i_epoch + 1) % 25 == 0:
for g in opt.param_groups:
g['lr'] = g['lr'] / 2
for i_iteration, (X, Y) in enumerate(train_dataloader):
# Set model to training mode and clear gradients
network.train()
opt.zero_grad()
X = X.cuda().float()
X = tranf(X)
logits = network.classifier(torch.nn.Dropout(dropout_rate)(torch.squeeze(network.bb(X))))
L = torch.nn.CrossEntropyLoss(weight=torch.tensor(class_weights).cuda().float())(logits, Y.type(torch.LongTensor).cuda())
L.backward()
opt.step()
L_iteration = L.detach().cpu().numpy()
l_epoch += L_iteration
info = "[INFO] Epoch {}/{} -- Step {}/{}: Lce={:.6f}".format(
i_epoch + 1, epochs, i_iteration + 1, len(train_dataloader), L_iteration)
print(info, end='\r')
l_epoch = l_epoch/len(train_dataloader)
k2 = test_instances(X_test, Y_test, network, dir_experiment, ii_iteration)
info = "[INFO] Epoch {}/{} -- Step {}/{}: Lce={:.6f}; k2={:.6f}".format(
i_epoch+1, epochs, i_iteration, len(train_dataloader), l_epoch, k2)
print(info, end='\n')
torch.save(network, dir_experiment + str(ii_iteration) + '_student_network_weights.pth')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--experiment_name", default="test_test_test", type=str)
args = parser.parse_args()
main(args)
Datasets
Models
