train_midrc_sex_adv_cross.py

import pickle
import os
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import json
import numpy as np
import time
import copy
import torchvision
import torchvision.transforms as transforms
from torchvision import datasets, models, transforms
from torch.optim import lr_scheduler
import cv2
from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import f1_score
from sklearn.metrics import cohen_kappa_score
from sklearn.metrics import roc_auc_score
from sklearn.metrics import confusion_matrix
from data_id_sex import get_id
from dataset import Dataset
import types
path = '/prj0129/mil4012/MIDRC' 
path_r = '/prj0129/mil4012/MIDRC/result_p/'

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def get_train_test_p_id(patient_list, fold, total_num_fold):
    num = len(patient_list)
    test_num = num // total_num_fold
    
    if fold == total_num_fold:
        test_name = patient_list[((fold-1) * test_num):]
        train_name = patient_list[0:((fold-1) * test_num)]
    else:
        test_name = patient_list[((fold-1) * test_num):fold * test_num]
        train_name = np.concatenate((patient_list[0:((fold-1) * test_num)], patient_list[(fold * test_num):]), axis=0)
    

    validation_name = train_name[int(0.8*len(train_name)):]


    train_name = train_name[0:(len(train_name)-len(validation_name))]

    return train_name, validation_name, test_name



class Counter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def compute_AUCs(gt, pred):
    N_CLASSES = 1
    AUROCs = []
    gt_np = gt.cpu().numpy()
    pred_np = pred.cpu().numpy()
    for i in range(N_CLASSES):
        AUROCs.append(roc_auc_score(gt_np[:, i], pred_np[:, i]))
    return AUROCs

def cross_auc(R_a_0, R_b_1): 
    scores = np.array(list(R_a_0.cpu().numpy()) + list(R_b_1.cpu().numpy()))
    y_true = np.zeros(len(R_a_0)+len(R_b_1))
    y_true[0:len(R_a_0)] = 1 # Pr[ LEFT > RIGHT]; Y = 1 is the left (A0)
    return roc_auc_score(y_true, scores)


def group_auc(labels, outputs, groups):
    
    group0p = []
    group0n = []
    group1p = []
    group1n = []

    for i in range(len(labels)):
        if groups[i] == 0:
            if labels[i][0] == 1:
                group0p.append(i)
            if labels[i][0] == 0:
                group0n.append(i)
        if groups[i] == 1:
            if labels[i][0] == 1:
                group1p.append(i)
            if labels[i][0] == 0:
                group1n.append(i)

                
    groupp = group0p+group1p 
    groupn = group0n+group1n
    outputs_ = outputs.clone().detach().cpu()
    
    try:
        AUC = cross_auc(torch.index_select(outputs_,0,torch.tensor(groupp)), torch.index_select(outputs_,0,torch.tensor(groupn)))
    except:
        AUC = 1
    try:
        A00 = cross_auc(torch.index_select(outputs_,0,torch.tensor(group0p)), torch.index_select(outputs_,0,torch.tensor(group0n)))
    except:
        A00 = 1
    try:
        A11 = cross_auc(torch.index_select(outputs_,0,torch.tensor(group1p)), torch.index_select(outputs_,0,torch.tensor(group1n)))
    except:
        A11 = 1
    try:
        A0a = cross_auc(torch.index_select(outputs_,0,torch.tensor(group0p)), torch.index_select(outputs_,0,torch.tensor(groupn)))
    except:
        A0a = 1
    try:
        A1a = cross_auc(torch.index_select(outputs_,0,torch.tensor(group1p)), torch.index_select(outputs_,0,torch.tensor(groupn)))
    except:
        A1a = 1
    try:
        Aa0 = cross_auc(torch.index_select(outputs_,0,torch.tensor(groupp)), torch.index_select(outputs_,0,torch.tensor(group0n)))
    except:
        Aa0 = 1
    try:
        Aa1 = cross_auc(torch.index_select(outputs_,0,torch.tensor(groupp)), torch.index_select(outputs_,0,torch.tensor(group1n)))
    except:
        Aa1 = 1
                

    group_num = [len(group0p),len(group0n),len(group1p),len(group1n)]
    
    return AUC, A00, A11, A0a, A1a, Aa0, Aa1, group_num




# def train_model(dataloaders,model, criterion, optimizer, scheduler, num_epochs=25):
def train_model(dataloaders,model, criterion, optimizer, num_epochs=25):
    since = time.time()
    fopen = open("/prj0129/mil4012/AREDS/adv_age_3.txt", "w")
    best_model_wts = copy.deepcopy(model.state_dict())
    best_AUROC_avg = 0.0
    losses = Counter()
    N_CLASSES = 1
    
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 100)
        
              
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            gt = torch.FloatTensor().to(device)
            pred = torch.FloatTensor().to(device)
            losses.reset()
            
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode
                 
            
            # Iterate over data.
            t = tqdm(enumerate(dataloaders[phase]),  desc='Loss: **** ', total=len(dataloaders[phase]), bar_format='{desc}{bar}{r_bar}')
            for batch_idx, (inputs, labels, group) in t:
                
                inputs = inputs.to(device)
                labels = labels.to(device)
                #print(inputs.shape, labels.shape)
                
                groups = []
                for i in group:
                    if i == 0:
                        groups.append([1,0])
                    elif i == 1:
                        groups.append([0,1])
                    else:
                        print('Error')
                groups = torch.FloatTensor(groups).to(device)
                
                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    out_pred, out_adv = model(inputs)
                    
                    gt = torch.cat((gt, labels), 0)
                    pred = torch.cat((pred, out_pred.data), 0)
                    
                    #print(outputs.shape)
                    loss = criterion(out_pred, out_adv, labels, groups)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                losses.update(loss.data.item(), inputs.size(0))
                t.set_description('Loss: %.3f ' % (losses.avg))
            
                        
            AUCs = compute_AUCs(gt, pred)
            AUROC_avg = np.array(AUCs).mean()
            
                
            if phase == "val":
                if best_AUROC_avg < AUROC_avg:
                    best_AUROC_avg = AUROC_avg     
                    torch.save(model.state_dict(), "/prj0129/mil4012/MIDRC/weights/densenet121_sex_adv_midrc_cross.pth")
                fopen.write('\nEpoch {} \t [{}] : \t {AUROC_avg:.3f}\n'.format(epoch, phase, AUROC_avg=AUROC_avg))
                for i in range(N_CLASSES):
                    fopen.write('{} \t {}\n'.format(CLASS_NAMES[i], AUCs[i]))
                fopen.write('-' * 100)
                    
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
                   epoch, batch_idx + 1, len(dataloaders[phase]), loss=losses))
            print('{} : \t {AUROC_avg:.3f}'.format(phase, AUROC_avg=AUROC_avg))
            
            #if phase == "val":
            #    scheduler.step(losses.avg)
                
            fopen.flush()
    fopen.close()
    return model


def test_model(test_loader,model):
    model.eval()
    gt = torch.FloatTensor().to(device)
    pred = torch.FloatTensor().to(device)
    groups = []
    with torch.no_grad():
        for batch_idx, (inputs, labels, group) in enumerate(test_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            out_pred, out_adv = model(inputs)
            gt = torch.cat((gt, labels), 0)
            pred = torch.cat((pred, out_pred.data), 0)
            groups += group
    AUCs = compute_AUCs(gt, pred)
    AUC, A00, A11, A0a, A1a, Aa0, Aa1, group_num = group_auc(gt, pred, groups)
    print('AUCs',AUCs)
    print('AUC',AUC)
    print('A00',A00)
    print('A11',A11)
    print('A0a',A0a)
    print('A1a',A1a)
    print('Aa0',Aa0)
    print('Aa1',Aa1)
    print('Group Num',group_num)
    pred1 = pred.cpu()
    pred2 = pred1.numpy()
    np.savetxt('/prj0129/mil4012/MIDRC/Result/densenet121_sex_adv_midrc_cross.txt', pred2)
            
        

if __name__ == '__main__':
    
    fold = 1
    total_num_fold = 5
    
    train_sampler = None
    batch_size = 128
    workers = 4
    N_CLASSES = 1
    CLASS_NAMES = ['MIDRC']
    
    data_path=os.path.join(path,'data_resize/')
    label_path=os.path.join(path,'filtered_final1.csv')
    
    label_path1 = os.path.join(path,'Patient_list.csv')
    tmp = np.loadtxt(label_path1, dtype=np.str_, delimiter=",")
    
    tmp = tmp[1:] 
    train_name, validation_name, test_name = get_train_test_p_id(tmp, fold, total_num_fold)
    
    
    train_path, train_labels, train_groups = get_id(data_path,label_path,data_id = train_name)
    val_path, val_labels, val_groups = get_id(data_path,label_path,data_id = validation_name)
    test_path, test_labels, test_groups = get_id(data_path,label_path,data_id = test_name)
    
  
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomRotation(10),
            # transforms.ToPILImage(),
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.RandomHorizontalFlip(0.5),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            # transforms.ToPILImage(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }
    
    train_dataset = Dataset(train_path,train_labels,groups=train_groups,transform = data_transforms["train"])
    val_dataset = Dataset(val_path,val_labels,groups=val_groups,transform = data_transforms["val"])
    test_dataset = Dataset(test_path,test_labels,groups=test_groups,transform = data_transforms["val"])
    
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=(train_sampler is None), 
                                           num_workers=workers, pin_memory=True, sampler=train_sampler)
        
#     train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, 
#                                            num_workers=workers, pin_memory=True, sampler=train_sampler)
    
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=(train_sampler is None), 
                                           num_workers=workers, pin_memory=True, sampler=train_sampler)
    
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, 
                                           num_workers=workers, pin_memory=True, sampler=train_sampler)

    dataloaders = {"train": train_loader, "val": val_loader}
    
    # #used imagenet
    # model_ft = models.densenet201(pretrained=True)
    # num_ftrs = model_ft.classifier.in_features
    # model_ft.classifier = nn.Sequential(
    #             nn.Linear(num_ftrs, N_CLASSES),
    #             nn.Sigmoid()
    #         )
    # # model_ft.classifier = nn.Linear(num_ftrs, N_CLASSES)
    # model_ft = model_ft.to(device)
    
    #used chexpert
    model_ft = models.densenet121(pretrained=True)
    num_ftrs = model_ft.classifier.in_features
    model_ft.classifier = nn.Sequential(
                nn.Linear(num_ftrs, 14),
                nn.Sigmoid()
            )
    checkpoint = torch.load("/prj0129/mil4012/glaucoma/NIH-chest-x-ray/CXR8/SupCon/m-30012020-104001.pth.tar", map_location=torch.device('cpu'))
    for i in list(checkpoint['state_dict'].keys()):
        checkpoint['state_dict'][i[12:]] = checkpoint['state_dict'].pop(i)
    model_ft.load_state_dict(checkpoint['state_dict'])

    model_ft.classifier_pred = nn.Sequential(
                nn.Linear(num_ftrs, N_CLASSES),
                nn.Sigmoid()
            )

    model_ft.classifier_adv = nn.Sequential(
                nn.Linear(num_ftrs, 2),
                nn.Sigmoid()
            )


    def forward(self, x):
        features = self.features(x)
        out = F.relu(features, inplace=True)
        out = F.adaptive_avg_pool2d(out, (1, 1))
        out = torch.flatten(out, 1)
        out_pred = self.classifier_pred(out)
        out_adv = self.classifier_adv(out)
        return out_pred, out_adv


    funcType = types.MethodType
    model_ft.forward = funcType(forward, model_ft)

    model_ft = model_ft.to(device)
    # print(model_ft)

    def criterion(out_pred, out_adv, labels, groups):

        loss_pred = nn.BCELoss()
        loss_adv = nn.BCELoss()
    
        return loss_pred(out_pred, labels) - 0.3*loss_adv(out_adv, groups)

    # Observe that all parameters are being optimized
    optimizer_ft = optim.Adam(model_ft.parameters(), lr=0.0001)
    
    # optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
    
    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = lr_scheduler.ReduceLROnPlateau(optimizer_ft, 'min', patience=2, eps=1e-08, verbose=True)

    # model_ft = train_model(dataloaders, model_ft, criterion, optimizer_ft, exp_lr_scheduler,
    #                        num_epochs=20)
#     model_ft = train_model(dataloaders, model_ft, criterion, optimizer_ft,num_epochs=20)
    model_ft.load_state_dict(torch.load("/prj0129/mil4012/MIDRC/weights/densenet121_sex_adv_midrc_cross.pth"))
    test_model(test_loader,model_ft)