inference_embed_attn.py

import os
import json
import argparse
from pickle import TRUE

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader,ConcatDataset, Subset
from torch.cuda.amp import autocast
import numpy as np

from dataset.CT_pancreas_ids import EvaPanCTDataset
from model.trans_3DUnet import get_model_dict
from loss.criterions import get_criterions

import monai
from monai.inferers import sliding_window_inference

def get_parse():
    parser = argparse.ArgumentParser()
    '''
    parser.add_argument('--dir_data', type=str,
                        default='../../data/CT_Pancreas/Sloan_data',
                        help='direction for the dataset')
    '''
    parser.add_argument('--dir_data', type=str,
                        default='/data/datasets/zheyuan/Raw_Pancreas',
                        help='direction for the dataset')
    # 12 -> 20220115-15_2
    parser.add_argument('--pretrained_dir', type=str,
                        default='./out/log/20220125-17_2', help='pretrained dir')
    parser.add_argument('--model_name', type=str,
                        default='MaskTransUnet', help='model name for training')
    parser.add_argument('--batch_size', type=int,
                        default=1, help='patient batch size')
    parser.add_argument('--depth_size', type=int,
                        default=32, help='patient depth size')
    # [16, 32, 64, 128, 256]
    # [32, 64, 64, 128, 256]
    parser.add_argument('--num_layers', type=list, 
                        default=[16, 32, 64, 128, 256], help='number of layer for each layer')
    # reference 320 160 80 40 20 [128, 80, 50, 30, 32]
    # 256-128-64-32-16 65
    parser.add_argument('--roi_size_list', type=list, 
                        default=[100, 65, 40, 25, 10], help='size of roi for each layer')
    # False, True, True, True, True
    parser.add_argument('--is_roi_list', type=list, 
                        default=[False, True, True, True, True], help='using roi for each layer')
    '''
    parser.add_argument('--num_layers', type=list, 
                        default=[16, 32, 32, 64], help='number of layer for each layer')
    '''
    parser.add_argument('--dim_input', type=int,
                        default=1, help='input dimension or modality')
    parser.add_argument('--dim_output', type=int,
                        default=2, help='output dimension or classes')
    parser.add_argument('--kernel_size', type=int,
                        default=3, help='kernel_size for convolution')

    parser.add_argument('--device', type=str,
                        default='cuda', help='device for training')
    parser.add_argument('--criterion_list', type=list,
                        default=['DiceClassLoss', 'Recall', 'Precision', 'LocalizationLoss'],
                        help='criterion')

    parser.add_argument('--is_save', type=bool,
                        default=False, help='save prediction or not')
    parser.add_argument('--saved_folder', type=str,
                        default='./prediction/test', 
                        help='saved folder dir')

    args = parser.parse_args()
    return args

def get_model(args, fold_num, device):
    model_fn = get_model_dict(args.model_name)
    model = model_fn(num_layers=args.num_layers,
                     roi_size_list=args.roi_size_list,
                     is_roi_list=args.is_roi_list,
                     dim_input=args.dim_input,
                     dim_output=args.dim_output,
                     kernel_size=args.kernel_size)
    
    pretrain_dir = os.path.join(args.pretrained_dir, f'fold_{fold_num}', 'temp_model.pt')
    # state_dict = torch.load(pretrain_dir).state_dict()
    state_dict = torch.load(pretrain_dir)

    model.load_state_dict(state_dict)
    model = torch.nn.DataParallel(model.to(device))
    return model

def main(args):
    fold_nums = 1
    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
    num_device = torch.cuda.device_count()
    root = args.dir_data
    depth_size = args.depth_size
    sw_batch_size = 4

    with open('split_dataset_8.json', 'r') as f:
        dataset_ids = json.load(f)

    criterions = get_criterions(args.criterion_list)
    final_loss_list = [0] * len(criterions)
    roi_size = 512
    center = 256
    name_list = sorted(os.listdir(os.path.join(root, 'image')))
    # up_sample = torch.nn.Upsample(scale_factor=(2, 2, 1), mode='trilinear', align_corners=True)
    # post_processing = monai.transforms.KeepLargestConnectedComponent(applied_labels=[1], connectivity=1)

    for fold_num in range(fold_nums):
        test_ids = dataset_ids[f'test_id fold_{fold_num}']
        eval_pandataset = EvaPanCTDataset(root=root,
                                           depth_size=depth_size,
                                           ids=test_ids[:-1])
        eval_panDl = DataLoader(dataset=eval_pandataset, batch_size=args.batch_size,
                                num_workers=12, shuffle=False)

        device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
        model = get_model(args, fold_num, device=device)
        
        model.eval()
        summary_patient_loss = []
        total_loss_list = [0] * len(criterions)
        threshold = 0.5

        if not os.path.exists(args.saved_folder):
            os.makedirs(args.saved_folder)

        for i, (images, masks) in enumerate(eval_panDl):
            name = name_list[test_ids[i]]
            print(name)

            images, masks = images.to(device), masks.to(device).long()

            # predict = torch.zeros((masks.size(0), 2, masks.size(2), masks.size(3), masks.size(4)), device=masks.device)
            patient_loss_list = [0] * len(criterions)

            with torch.no_grad():
                with autocast():
                    predict = sliding_window_inference(images, (roi_size, roi_size, depth_size), sw_batch_size, model, overlap=0.6, sigma_scale=0)
                    # print(predict_center.shape)
                    '''
                    predict2 = predict
                    '''
                    predict2 = (predict >= threshold).float().squeeze(0)
                    # predict2 = post_processing(predict2)
                    predict2 = predict2.unsqueeze(0)
                    
                    loss_list = [l(predict2, masks.long()).item() for l in criterions.values()]
                    
            if args.is_save:
                # predict = up_sample(predict)
                temp_out = predict2[:, 1]
                # temp_out = predict
                temp_out = temp_out.squeeze_().permute((2, 0, 1)).cpu().numpy()

                np.save(os.path.join(args.saved_folder, '{:0>4}'.format(name)), temp_out)
            
            for loss_name, loss_value in zip(criterions.keys(), loss_list):
                print(f'eval patient average {loss_name}', loss_value)

            for index, loss_value in enumerate(patient_loss_list):
                patient_loss_list[index] = loss_list[index]
                total_loss_list[index] += patient_loss_list[index]
            summary_patient_loss.append(patient_loss_list)

        for index, loss_value in enumerate(total_loss_list):
            total_loss_list[index] = loss_value / (i+1)
            final_loss_list[index] += total_loss_list[index]

        for loss_name, loss_value in zip(criterions.keys(), total_loss_list):
            print(f'eval total average {loss_name} loss', loss_value)

        out_dict = {f'patient_{fold_num}': summary_patient_loss,
                    f'summary_{fold_num}': total_loss_list}

    for index, loss_value in enumerate(final_loss_list):
        final_loss_list[index] = loss_value / (fold_num+1)

    for loss_name, loss_value in zip(criterions.keys(), final_loss_list):
        print(f'eval final average {loss_name} loss', loss_value)

    with open('summary_4_fold.json', 'w') as f:
            json.dump(out_dict, f, indent=4)


if __name__ =='__main__':
    args = get_parse()
    main(args)