evaluate/evaluator.py

import os
import time

import torch
from torch.cuda.amp import autocast as autocast
from tqdm import tqdm
from einops import rearrange, repeat, reduce
import numpy as np
import pandas as pd
from pathlib import Path
import nibabel as nib
import shutil
import pickle
from scipy.ndimage import gaussian_filter
import torch.distributed as dist

from evaluate.metric import calculate_metric_percase
from evaluate.merge_after_evaluate import merge
from train.dist import is_master

def compute_gaussian(tile_size, sigma_scale: float = 1. / 8, value_scaling_factor: float = 10, dtype=np.float16):
    tmp = np.zeros(tile_size)
    center_coords = [i // 2 for i in tile_size]
    sigmas = [i * sigma_scale for i in tile_size]
    tmp[tuple(center_coords)] = 1
    gaussian_importance_map = gaussian_filter(tmp, sigmas, 0, mode='constant', cval=0)

    # gaussian_importance_map = torch.from_numpy(gaussian_importance_map)

    gaussian_importance_map = gaussian_importance_map / np.max(gaussian_importance_map) * value_scaling_factor
    gaussian_importance_map = gaussian_importance_map.astype(dtype)

    # gaussian_importance_map cannot be 0, otherwise we may end up with nans!
    gaussian_importance_map[gaussian_importance_map == 0] = np.min(
        gaussian_importance_map[gaussian_importance_map != 0])

    return gaussian_importance_map

def evaluate(model, 
             text_encoder, 
             device, 
             testset, 
             testloader, 
             dice_score,
             nsd_score,
             csv_path, 
             resume,
             save_interval,
             visualization):
    
    # if to store pred、gt、img (as nii.gz
    if visualization:
        nib_dir = csv_path.replace('.csv', '')
        
    # collate in master process
    if is_master():
        # datasets --> labels --> metrics
        datasets_labels_metrics = {}   # {'COVID19':{'covid19_infection':{'dice':[0.8, 0.9, ...], ...} ...}, ...}
    
        # datasets --> samples --> labels --> metrics
        samples_labels_metrics = {}   # {'COVID19':{'0.npy':{'covid19_infection':{'dice':0.8, ...} ...}, ...} 记录每个dataset里的sample（行）
        
        # datsets --> labels
        datasets_labels_sets = {}    # {'COVID19':set('covid19_infection', ...), ...}  记录每个dataset里的label种类（列）
        
    # accumulate scores of each sample in each process
    results_of_samples = [] # each element : [dataset_name, modality, sample_id, scores_of_labels(dict), label_names] 
    
    # load results from an interrupted eval (only in master process)
    if resume and is_master():
        root_dir = os.path.dirname(csv_path)
        prefix = os.path.basename(csv_path).replace('.csv', '_tmp_rank')  # xxx/test/step_xxx.csv --> step_xxx_tmp_rank
        pkl_to_del = []
        for f in os.listdir(root_dir):
            if prefix in f:
                # load list of results
                pkl_path = f'{root_dir}/{f}'
                with open(pkl_path, 'rb') as f:
                    results_of_samples += pickle.load(f)
                print(f'Load results from {pkl_path}')   
                pkl_to_del.append(pkl_path)
                
        # there may be duplication? We leave the deduplication to the final merge
        # merge all the loaded samples, del the tmp pickle files in previous evaluation task
        for pkl_path in pkl_to_del:
            os.remove(pkl_path)
            print(f'Del {pkl_path}')
        merge_pkl = csv_path.replace('.csv', f'_tmp_rank0.pkl')
        with open(merge_pkl, 'wb') as f:
            pickle.dump(results_of_samples, f)  
        print(f'Load results of {len(results_of_samples)} samples, Merge into {merge_pkl}')
                        
    model.eval()
    text_encoder.eval()
        
    with torch.no_grad():
        
        data_time = 0
        pred_time = 0
        metric_time = 0
        
        avg_patch_batch_num = 0
        avg_query_batch_num = 0
        
        # in ddp, only master process display the progress bar
        if is_master():
            testloader = tqdm(testloader, disable=False)
        else:
            testloader = tqdm(testloader, disable=True)  
        
        # gaussian kernel to accumulate predcition
        gaussian = torch.tensor(compute_gaussian((288, 288, 96))).to(device)    # hwd

        end_time = time.time()
        for sample in testloader:    # in evaluation/inference, a "batch" in loader is a volume
            # data loading
            dataset_name = sample['dataset_name']
            sample_id = sample['sample_id'] 
            batched_patches = sample['batched_patches']
            batched_y1y2_x1x2_z1z2 = sample['batched_y1y2_x1x2_z1z2']
            labels = sample['labels'] 
            gt_segmentation = sample['gt_segmentation'].numpy()  # n h w d
            modality = sample['modality']
            image_path = sample['image_path']

            n,h,w,d = gt_segmentation.shape
            prediction = torch.zeros((n, h, w, d))
            accumulation = torch.zeros((n, h, w, d))
            
            data_time += (time.time()-end_time) 
            end_time = time.time()
            
            with autocast():
                
                queries = text_encoder(labels, modality)
                      
                # for each batch of patches, query with all labels
                for patches, y1y2_x1x2_z1z2_ls in zip(batched_patches, batched_y1y2_x1x2_z1z2):   # [b, c, h, w, d]
                    patches = patches.to(device=device)
                    prediction_patch = model(queries=queries, image_input=patches, train_mode=False)
                    prediction_patch = torch.sigmoid(prediction_patch)  # bnhwd
                    prediction_patch = prediction_patch.detach() # .cpu().numpy()
                    
                    # fill in 
                    for b in range(len(y1y2_x1x2_z1z2_ls)):
                        y1, y2, x1, x2, z1, z2 = y1y2_x1x2_z1z2_ls[b]

                        # gaussian accumulation
                        tmp = prediction_patch[b, :, :y2-y1, :x2-x1, :z2-z1] * gaussian[:y2-y1, :x2-x1, :z2-z1] # on gpu
                        prediction[:, y1:y2, x1:x2, z1:z2] += tmp.cpu()
                        accumulation[:, y1:y2, x1:x2, z1:z2] += gaussian[:y2-y1, :x2-x1, :z2-z1].cpu()
            
            pred_time += (time.time()-end_time)
            end_time = time.time()
                            
            # avg            
            prediction = prediction / accumulation
            prediction = torch.where(prediction>0.5, 1.0, 0.0)
            prediction = prediction.numpy()
            
            # cal metrics : [{'dice':x, ...}, ...] 
            scores = []
            for j in range(len(labels)):
                scores.append(calculate_metric_percase(prediction[j, :, :, :], gt_segmentation[j, :, :, :], dice_score, nsd_score))    # {'dice':0.9, 'nsd':0.8} 每个label一个dict
            
            # visualization  
            if visualization:
                Path(f'{nib_dir}/{dataset_name}').mkdir(exist_ok=True, parents=True)
                # 将image、gt和prediction保存下来
                results = np.zeros((h, w, d)) # hwd
                for j, label in enumerate(labels):
                    results += prediction[j, :, :, :] * (j+1)   # 0 --> 1 (skip background)
                    Path(f'{nib_dir}/{dataset_name}/seg_{sample_id}').mkdir(exist_ok=True, parents=True)
                    # 每个label单独一个nii.gz
                    segobj = nib.nifti2.Nifti1Image(prediction[j, :, :, :], np.eye(4))
                    nib.save(segobj, f'{nib_dir}/{dataset_name}/seg_{sample_id}/{label}.nii.gz')
                segobj = nib.nifti2.Nifti1Image(results, np.eye(4))
                nib.save(segobj, f'{nib_dir}/{dataset_name}/seg_{sample_id}.nii.gz')
                
                image = testset.load_image(image_path)
                image = np.squeeze(image)
                imgobj = nib.nifti2.Nifti1Image(image, np.eye(4))
                nib.save(imgobj, f'{nib_dir}/{dataset_name}/img_{sample_id}.nii.gz')
                
                gt = np.zeros((h, w, d)) # hwd
                for j, label in enumerate(labels):
                    gt += gt_segmentation[j, :, :, :] * (j+1)   # 0 --> 1 (skip background)
                    Path(f'{nib_dir}/{dataset_name}/gt_{sample_id}').mkdir(exist_ok=True, parents=True)
                    # 每个label单独一个nii.gz
                    segobj = nib.nifti2.Nifti1Image(gt_segmentation[j, :, :, :], np.eye(4))
                    nib.save(segobj, f'{nib_dir}/{dataset_name}/gt_{sample_id}/{label}.nii.gz')
                gtobj = nib.nifti2.Nifti1Image(gt, np.eye(4))
                nib.save(gtobj, f'{nib_dir}/{dataset_name}/gt_{sample_id}.nii.gz')
                
            metric_time += (time.time()-end_time)
            end_time = time.time()
            
            # accumulate
            results_of_samples.append([dataset_name, modality, sample_id, scores, labels])
            
            # save in each process regularly in case of interruption
            if len(results_of_samples) % save_interval == 0:
                with open(csv_path.replace('.csv', f'_tmp_rank{dist.get_rank()}.pkl'), 'wb') as f:
                    pickle.dump(results_of_samples, f)
            
        """
        # gather results from all device to rank-0 (solution 1)
        gather_results = [None for i in range(dist.get_world_size())]
        dist.gather_object(
            results_of_samples, 
            gather_results if dist.get_rank() == 0 else None,
            dst = 0
            )
        
        if int(dist.get_rank()) == 0:
            results_of_samples = [tmp for ls in results_of_samples for tmp in ls]
        """         
                
        avg_patch_batch_num /= len(testloader)
        avg_query_batch_num /= len(testloader)
        data_time /= len(testloader)
        pred_time /= len(testloader)
        metric_time /= len(testloader)
        print(f'On Rank {dist.get_rank()}, each sample has {avg_patch_batch_num} batch of patches and {avg_query_batch_num} batch of queries, Data Time: {data_time}, Pred Time: {pred_time}, Dice Time: {metric_time}')
          
        torch.cuda.empty_cache()
        
        # save in each process (to a fnl pickle, also denoting this process ends)
        with open(csv_path.replace('.csv', f'_fnl_rank{dist.get_rank()}.pkl'), 'wb') as f:
            pickle.dump(results_of_samples, f)
                    
        # gather and record in rank 0 (solution 2)
        if is_master():
            
            # detect the finish of each process
            while True:
                all_process_finished = True
                for rank_id in range(torch.distributed.get_world_size()):
                    if not os.path.exists(csv_path.replace('.csv', f'_fnl_rank{rank_id}.pkl')): # xxx_tmp_rankx.pkl
                        all_process_finished = False
                        break
                if all_process_finished:
                    break
                else:
                    time.sleep(10)
            
            # read results of each process (samples may be duplicated due to the even distribution of ddp, check)
            results_of_samples = []    
            for rank_id in range(torch.distributed.get_world_size()):
                fnl_results_file = csv_path.replace('.csv', f'_fnl_rank{rank_id}.pkl')
                tmp_results_file = csv_path.replace('.csv', f'_tmp_rank{rank_id}.pkl')
                with open(fnl_results_file, 'rb') as f:
                    results_of_samples += pickle.load(f)
                os.remove(fnl_results_file)
                if os.path.exists(tmp_results_file):
                    os.remove(tmp_results_file)
                
            # check duplication
            unique_set = set()
            deduplicated_results_of_samples = []
            for dataset_name, modality, sample_id, scores, labels in results_of_samples:
                if f'{dataset_name}/{sample_id}' not in unique_set:
                    unique_set.add(f'{dataset_name}/{sample_id}')
                    deduplicated_results_of_samples.append([dataset_name, modality, sample_id, scores, labels])
            results_of_samples = deduplicated_results_of_samples
            
            # save for tmp
            with open(csv_path.replace('.csv', '.pkl'), 'wb') as f:
                pickle.dump(results_of_samples, f)

            # collate results
            for dataset_name, modality, sample_id, scores, labels in results_of_samples:    #  [[dataset_name, modality, sample_id, scores_of_labels(dict), label_names], ...]
                dataset_name = f'{dataset_name}({modality})'
                
                if dataset_name not in datasets_labels_metrics:
                    datasets_labels_metrics[dataset_name] = {}  # {'COVID19(CT)':{}}
                if dataset_name not in datasets_labels_sets:
                    datasets_labels_sets[dataset_name] = set()  # {'COVID19(CT)':set()}
                if dataset_name not in samples_labels_metrics:
                    samples_labels_metrics[dataset_name] = {}
                samples_labels_metrics[dataset_name][sample_id] = {}   # {'COVID19(CT)':{'0':{}}}
                
                for metric_dict, label in zip(scores, labels):
                    # accumulate metrics （for per dataset per class
                    # {'COVID19(CT)':{'covid19_infection':{'dice':[0.8, 0.9, ...], 'nsd':[0.8, 0.9, ...], ...} ...}, ...}
                    if label not in datasets_labels_metrics[dataset_name]:
                        datasets_labels_metrics[dataset_name][label] = {k:[v] for k,v in metric_dict.items()}
                    else:
                        for k,v in metric_dict.items():
                            datasets_labels_metrics[dataset_name][label][k].append(v)
                    
                    # statistic labels
                    # {'COVID19(CT)':set('covid19_infection', ...)}
                    if label not in datasets_labels_sets[dataset_name]:
                        datasets_labels_sets[dataset_name].add(label)
                    
                    # record metrics （for per dataset per sample per class
                    # {'COVID19':{'0.npy':{'covid19_infection':{'dice':0.8, 'nsd':0.9, ...} ...}, ...}
                    samples_labels_metrics[dataset_name][sample_id][label] = {k:v for k,v in metric_dict.items()}
                        
            # average and log (列为metrics，例如dice，nsd...)
            # create a df like:
            # {
            #   'TotalSegmentator': [0.xx, 0.xx, ...]    # 在T之前，这是一列
            #   'TotalSegmentator, Lung': [0.68, 0.72, ...]
            # }
            # by defult, print the dice (1st metric) of each dataset 
            info = 'Metrics of Each Dataset:\n'
            avg_df = {}
            for dataset in datasets_labels_metrics.keys():
                avg_df[dataset] = {k:[] for k in metric_dict.keys()}    # 'TotalSegmentator(CT)': {'dice':[0.8, ...] 'nsd':[0.5, ...], ...}
                for label in datasets_labels_metrics[dataset].keys():
                    avg_df[f'{dataset}, {label}'] = []
                    for metric in datasets_labels_metrics[dataset][label].keys():
                        label_metric = np.average(datasets_labels_metrics[dataset][label][metric])
                        avg_df[f'{dataset}, {label}'].append(label_metric)  # 'TotalSegmentator, Lung': [0.68, 0.72, ...] list of num_metrics
                        avg_df[dataset][metric].append(label_metric)
                avg_df[dataset] = {k:np.average(v) for k,v in avg_df[dataset].items()} # 'TotalSegmentator': {'dice':[0.8, ...] 'nsd':[0.5, ...], ...} --> 'TotalSegmentator': {'dice':0.x, 'nsd':0.x, ...}
                info += f'{dataset}  |  ' 
                for k ,v in avg_df[dataset].items():
                    info += f'{v}({k})  |  '
                info += '\n'
                avg_df[dataset] = list(avg_df[dataset].values())
            avg_df = pd.DataFrame(avg_df).T
            avg_df.columns = list(metric_dict.keys())   # ['dice', 'nsd']
            avg_df.to_csv(csv_path)        
            print(info)
            
            # detailed log （nsd和dice，列为class label
            # multi-sheet, two for each dataset
            df_list = [['summary', avg_df]]
            for dataset, label_set in datasets_labels_sets.items():
                metric_df ={}
                if dice_score:
                    metric_df['dice'] = {}
                if nsd_score:
                    metric_df['nsd'] = {}

                # create dfs like:
                # {
                #   '0.npy': [0.xx, 0.xx, ...]
                #   ......
                # }
                
                # {'COVID19':{'0.npy':{'covid19_infection':{'dice':0.8, ...} ...}, ...}
                for image_id, label_dict in samples_labels_metrics[dataset].items():
                    for metric in metric_df:
                        tmp = []    # one dice for each label in this dataset
                        for label in label_set:
                            score = label_dict[label][metric] if label in label_dict else -1
                            tmp.append(score)
                        metric_df[metric][image_id] = tmp   
                
                for metric, metric_df in metric_df.items():
                    metric_df = pd.DataFrame(metric_df).T
                    metric_df.columns = list(label_set)
                    df_list.append([dataset+f'({metric})', metric_df])
                
            xlsx_path = csv_path.replace('.csv', '.xlsx')
            with pd.ExcelWriter(xlsx_path) as writer:
                for name, df in df_list:
                    # 将每个 DataFrame 写入一个 sheet(sheet name must be < 31)
                    if len(name) > 31:
                        name = name[len(name)-31:]
                    df.to_excel(writer, sheet_name=name, index=True)    
                    
            # avg_dice_over_merged_labels, avg_nsd_over_merged_labels = merge(region_split_json, label_statistic_json, xlsx_path, xlsx_path)  
            
            os.remove(csv_path.replace('.csv', '.pkl'))
            
        else:
            
            pass
        
            # avg_dice_over_merged_labels = avg_nsd_over_merged_labels = 0
            
        return # avg_dice_over_merged_labels, avg_nsd_over_merged_labels