A-Eval: A Benchmark for Cross-Dataset Evaluation of Abdominal Multi-Organ Segmentation
from collections import defaultdict
import os
import glob
import nibabel as nib
import numpy as np
import json
import csv
from tqdm import tqdm
from SurfaceDice import compute_surface_distances, compute_surface_dice_at_tolerance, compute_dice_coefficient


# label_tolerance = OrderedDict({'Liver': 5, 'RK':3, 'Spleen':3, 'Pancreas':5, 
#                    'Aorta': 2, 'IVC':2, 'RAG':2, 'LAG':2, 'Gallbladder': 2,
#                    'Esophagus':3, 'Stomach': 5, 'Duodenum': 7, 'LK':3})
label_tolerance = [0,5,3,3,5,2,2,2,2,2,3,5,7,3]

def process_results(prediction_folder, gt_folder, class_list):
    print(prediction_folder)
    print(gt_folder)
    total_dsc = defaultdict(int)
    total_nsd = defaultdict(int)
    count = defaultdict(int)
    case_dsc = defaultdict(list)
    case_nsd = defaultdict(list)
    
    prediction_files = glob.glob(os.path.join(prediction_folder, '*.nii.gz'))
    for prediction_file in tqdm(prediction_files, desc='Processing', unit=' file'):
        filename = os.path.basename(prediction_file)
        gt_file = os.path.join(gt_folder, filename)
        gt_nii = nib.load(gt_file)
        case_spacing = gt_nii.header.get_zooms()

        if os.path.exists(gt_file):
            prediction = nib.load(prediction_file).get_fdata().astype(np.uint8)
            gt = nib.load(gt_file).get_fdata().astype(np.uint8)


            for i in class_list:
                if np.sum(gt == i) == 0 and np.sum(prediction == i) == 0:
                    DSC_i = 1
                    NSD_i = 1
                elif np.sum(gt == i) == 0 and np.sum(prediction == i) > 0:
                    DSC_i = 0
                    NSD_i = 0
                else:
                    if i == 5 or i == 6 or i == 10:  # for Aorta, IVC, and Esophagus, only evaluate the labelled slices in ground truth
                        z_lower, z_upper = find_lower_upper_zbound(gt == i)
                        organ_i_gt, organ_i_seg = gt[:, :, z_lower:z_upper] == i, prediction[:, :, z_lower:z_upper] == i
                    else:
                        organ_i_gt, organ_i_seg = gt == i, prediction == i

                    # Check if both arrays are not empty
                    if organ_i_gt.size > 0 and organ_i_seg.size > 0:
                        surface_distances = compute_surface_distances(organ_i_gt, organ_i_seg, case_spacing)
                        DSC_i = compute_dice_coefficient(organ_i_gt, organ_i_seg)
                        NSD_i = compute_surface_dice_at_tolerance(surface_distances, label_tolerance[i])
                    else:
                        DSC_i = 0
                        NSD_i = 0

                total_dsc[i] += DSC_i
                total_nsd[i] += NSD_i
                count[i] += 1
                case_dsc[filename].append(DSC_i)
                case_nsd[filename].append(NSD_i)

    # Calculate the averages
    avg_dsc = {i: total_dsc[i] / count[i] if count[i] > 0 else 0 for i in class_list}
    avg_nsd = {i: total_nsd[i] / count[i] if count[i] > 0 else 0 for i in class_list}

    
    avg_dsc_overall = sum(total_dsc.values()) / sum(count.values())
    avg_nsd_overall = sum(total_nsd.values()) / sum(count.values())

    return avg_dsc, avg_nsd, avg_dsc_overall, avg_nsd_overall, case_dsc, case_nsd

def save_to_csv(data, filename):
    with open(filename, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        for row in data:
            writer.writerow(row)

def store_and_print_results(prediction_file, gt_file, dice_list, all_dice, case_dice):
    print(f'Dice coefficients for {prediction_file} and {gt_file}:', dice_list)
    all_dice[prediction_file] = dice_list
    avg_dice = np.nanmean(dice_list)
    case_dice[prediction_file] = {'dice_list': dice_list, 'avg_dice': avg_dice}


def find_lower_upper_zbound(organ_mask):
    organ_mask = np.uint8(organ_mask)
    assert np.max(organ_mask) ==1, print('mask label error!')
    z_index = np.where(organ_mask>0)[2]
    z_lower = np.min(z_index)
    z_upper = np.max(z_index)
    
    return z_lower, z_upper

def save_summary(avg_dsc, avg_nsd, avg_dsc_overall, avg_nsd_overall, summary_filename, dsc_csv_filename, nsd_csv_filename, case_dsc, case_nsd, class_list):
    summary = {
        'avg_dice_by_class': avg_dsc,
        'avg_nsd_by_class': avg_nsd,
        'avg_dice_overall': avg_dsc_overall,
        'avg_nsd_overall': avg_nsd_overall,
    }

    with open(summary_filename, 'w') as f:
        json.dump(summary, f, indent=4)
        
    dsc_rows = [["Case Name"] + [f"Class_{i}" for i in class_list]]
    for case, values in case_dsc.items():
        dsc_rows.append([case] + values)

    nsd_rows = [["Case Name"] + [f"Class_{i}" for i in class_list]]
    for case, values in case_nsd.items():
        nsd_rows.append([case] + values)
        
    save_to_csv(dsc_rows, dsc_csv_filename)
    save_to_csv(nsd_rows, nsd_csv_filename)


if __name__ == "__main__":
    if not os.path.exists('summarys'):
        os.makedirs('summarys')

    prediction_folders_path = 'predict_folders_converted_to_overlap'
    prediction_folders = os.listdir(prediction_folders_path)
    gt_folder_path = 'gtsVal_overlap_label'
    
    for prediction_folder in prediction_folders:
        if os.path.exists(os.path.join('summarys', prediction_folder+'_summary.json')):
            continue
        
        start_index = prediction_folder.index('_2_') + len('_2_')
        gt_folder = os.path.join(gt_folder_path, prediction_folder[start_index:]+'_gt')
        class_list = [1, 2, 3, 4, 5, 6, 7, 8]
        
        avg_dsc, avg_nsd, avg_dsc_overall, avg_nsd_overall, case_dsc, case_nsd = process_results(os.path.join(prediction_folders_path, prediction_folder), gt_folder, class_list)
        
        summary_filename = os.path.join('summarys', f"{prediction_folder}_summary.json")
        dsc_csv_filename = os.path.join('summarys', f"{prediction_folder}_dsc.csv")
        nsd_csv_filename = os.path.join('summarys', f"{prediction_folder}_nsd.csv")
        
        save_summary(avg_dsc, avg_nsd, avg_dsc_overall, avg_nsd_overall, summary_filename, dsc_csv_filename, nsd_csv_filename, case_dsc, case_nsd, class_list)
Evaluation/calculate_dsc_nsd.py
uni-medical/A-Eval
import os
import SimpleITK as sitk
import numpy as np
from tqdm import tqdm
from label_systems import FLARE_label, AMOS_label, WORD_label, TotalSeg_label, BTCV_label, overlap_label

def generate_label_map(label_dict, target_label_dict):
    target_value = {v: k for k, v in target_label_dict.items()}
    return {k: target_value[v] for k, v in label_dict.items() if v in target_value}

def convert_folder(original_seg_folder, prediction_folders_raw_path, prediction_folders_converted_path, label_map):
    if not os.path.exists(os.path.join(prediction_folders_converted_path, original_seg_folder)):
        os.makedirs(os.path.join(prediction_folders_converted_path, original_seg_folder))

    for seg_file in tqdm(os.listdir(os.path.join(prediction_folders_raw_path, original_seg_folder))):
        if '.nii.gz' not in seg_file:
            continue
        seg_sitk = sitk.ReadImage(os.path.join(prediction_folders_raw_path, original_seg_folder, seg_file))
        seg = sitk.GetArrayFromImage(seg_sitk)
        converted = np.zeros_like(seg, dtype=seg.dtype)

        for label in label_map:
            converted[seg == int(label)] = int(label_map[label])
        
        converted_sitk = sitk.GetImageFromArray(converted.astype(np.uint8))
        converted_sitk.SetSpacing(seg_sitk.GetSpacing())
        converted_sitk.SetOrigin(seg_sitk.GetOrigin())
        converted_sitk.SetDirection(seg_sitk.GetDirection())
        sitk.WriteImage(converted_sitk, os.path.join(prediction_folders_converted_path, original_seg_folder, seg_file))

def main():
    prediction_folders_raw_path = 'predict_results'
    prediction_folders_converted_path = 'predict_folders_converted_to_overlap'
    # prediction_folders_raw_path = 'gtsVal_raw'
    # prediction_folders_converted_path = 'gtsVal_overlap_label'
    
    print(f"Converting all segmentation results from subfolders in {prediction_folders_raw_path} to subfolders in {prediction_folders_converted_path}.")

    label_maps = {
        'FLARE': generate_label_map(FLARE_label, overlap_label),
        'AMOS': generate_label_map(AMOS_label, overlap_label),
        'WORD': generate_label_map(WORD_label, overlap_label),
        'TotalSeg': generate_label_map(TotalSeg_label, overlap_label),
        'BTCV': generate_label_map(BTCV_label, overlap_label)
    }

    prediction_folders = os.listdir(prediction_folders_raw_path)

    for original_seg_folder in prediction_folders:
        
        for prefix, label_map in label_maps.items():
            if original_seg_folder.startswith(prefix):
                print(f"Converting label system of {prefix} to overlap label system")
                for k, v in label_map.items():
                    print(f"{FLARE_label[v]}: {k} -> {v}")
                
                convert_folder(original_seg_folder, prediction_folders_raw_path, prediction_folders_converted_path, label_map)
                break
        else:
            print(f"The label system of {original_seg_folder} is not known. Skipping conversion.")

if __name__ == '__main__':
    main()

Evaluation/convert_label_2_overlap_label.py
uni-medical/A-Eval
overlap_label = {
    "0": "background",
    "1": "liver",
    "2": "kidney right",
    "3": "spleen",
    "4": "pancreas",
    "5": "gallbladder",
    "6": "esophagus",
    "7": "stomach",
    "8": "kidney left"
}


FLARE_label = {
    "0": "background",
    "1": "liver",
    "2": "kidney right",
    "3": "spleen",
    "4": "pancreas",
    "5": "aorta",
    "6": "inferior vena cava",
    "7": "adrenal gland right",
    "8": "adrenal gland left",
    "9": "gallbladder",
    "10": "esophagus",
    "11": "stomach",
    "12": "duodenum",
    "13": "kidney left"
}

AMOS_label = {
    "0": "background",
    "1": "spleen",
    "2": "kidney right",
    "3": "kidney left",
    "4": "gallbladder",
    "5": "esophagus",
    "6": "liver",
    "7": "stomach",
    "8": "aorta",
    "9": "inferior vena cava",
    "10": "pancreas",
    "11": "adrenal gland right",
    "12": "adrenal gland left",
    "13": "duodenum",
    "14": "bladder",
    "15": "prostate/uterus"
}

WORD_label = {
    "0": "background",
    "1": "liver",
    "2": "spleen",
    "3": "kidney left",
    "4": "kidney right",
    "5": "stomach",
    "6": "gallbladder",
    "7": "esophagus",
    "8": "pancreas",
    "9": "duodenum",
    "10": "colon",
    "11": "intestine",
    "12": "adrenal",
    "13": "rectum",
    "14": "bladder",
    "15": "femur left",
    "16": "femur right",
}

TotalSeg_label = {
    "0": "background",
    "1": "adrenal gland left",
    "2": "adrenal gland right",
    "3": "aorta",
    "9": "colon",
    "10": "duodenum",
    "11": "esophagus",
    "15": "gallbladder",
    "37": "inferior vena cava",
    "38": "kidney left",
    "39": "kidney right",
    "40": "liver",
    "46": "pancreas",
    "47": "portal vein and splenic vein",
    "77": "spleen",
    "78": "stomach",
    "80": "bladder",
    "105": "prostate/uterus",
    "106": "intestine",
    "107": "rectum",
}

BTCV_label = {
    "0": "background",
    "1": "spleen",
    "2": "kidney right",
    "3": "kidney left",
    "4": "gallbladder",
    "5": "esophagus",
    "6": "liver",
    "7": "stomach",
    "8": "aorta",
    "9": "inferior vena cava",
    "10": "portal vein and splenic vein",
    "11": "pancreas",
    "12": "adrenal gland right",
    "13": "adrenal gland left"
}


Evaluation/label_systems.py
uni-medical/A-Eval
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
import torch
from torch import nn

class STUNetTrainer(nnUNetTrainer):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 1000
        self.initial_lr = 1e-2

    @staticmethod
    def build_network_architecture(plans_manager,
                                   dataset_json,
                                   configuration_manager,
                                   num_input_channels,
                                   enable_deep_supervision: bool = True) -> nn.Module: 
        label_manager = plans_manager.get_label_manager(dataset_json)
        num_classes=label_manager.num_segmentation_heads
        kernel_sizes = [[3,3,3]] * 6
        strides=configuration_manager.pool_op_kernel_sizes[1:]
        if len(strides)>5:
            strides = strides[:5]
        while len(strides)<5:
            strides.append([1,1,1])
        return STUNet(num_input_channels, num_classes, depth=[1]*6, dims= [32 * x for x in [1, 2, 4, 8, 16, 16]], 
                      pool_op_kernel_sizes=strides, conv_kernel_sizes=kernel_sizes, enable_deep_supervision=enable_deep_supervision)

class STUNetTrainer_small(STUNetTrainer):
    @staticmethod
    def build_network_architecture(plans_manager,
                                   dataset_json,
                                   configuration_manager,
                                   num_input_channels,
                                   enable_deep_supervision: bool = True) -> nn.Module: 
        label_manager = plans_manager.get_label_manager(dataset_json)
        num_classes=label_manager.num_segmentation_heads
        kernel_sizes = [[3,3,3]] * 6
        strides=configuration_manager.pool_op_kernel_sizes[1:]
        if len(strides)>5:
            strides = strides[:5]
        while len(strides)<5:
            strides.append([1,1,1])
        return STUNet(num_input_channels, num_classes, depth=[1]*6, dims= [16 * x for x in [1, 2, 4, 8, 16, 16]], 
                      pool_op_kernel_sizes=strides, conv_kernel_sizes=kernel_sizes, enable_deep_supervision=enable_deep_supervision)

class STUNetTrainer_small_ep2k(STUNetTrainer_small):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 2000

class STUNetTrainer_small_ep4k(STUNetTrainer_small):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 4000

class STUNetTrainer_base(STUNetTrainer):
    @staticmethod
    def build_network_architecture(plans_manager,
                                   dataset_json,
                                   configuration_manager,
                                   num_input_channels,
                                   enable_deep_supervision: bool = True) -> nn.Module: 
        label_manager = plans_manager.get_label_manager(dataset_json)
        num_classes=label_manager.num_segmentation_heads
        kernel_sizes = [[3,3,3]] * 6
        strides=configuration_manager.pool_op_kernel_sizes[1:]
        if len(strides)>5:
            strides = strides[:5]
        while len(strides)<5:
            strides.append([1,1,1])
        return STUNet(num_input_channels, num_classes, depth=[1]*6, dims= [32 * x for x in [1, 2, 4, 8, 16, 16]], 
                      pool_op_kernel_sizes=strides, conv_kernel_sizes=kernel_sizes, enable_deep_supervision=enable_deep_supervision)

class STUNetTrainer_base_ep2k(STUNetTrainer_base):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 2000

class STUNetTrainer_base_ep4k(STUNetTrainer_base):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 4000

class STUNetTrainer_large(STUNetTrainer):
    @staticmethod
    def build_network_architecture(plans_manager,
                                   dataset_json,
                                   configuration_manager,
                                   num_input_channels,
                                   enable_deep_supervision: bool = True) -> nn.Module: 
        label_manager = plans_manager.get_label_manager(dataset_json)
        num_classes=label_manager.num_segmentation_heads
        kernel_sizes = [[3,3,3]] * 6
        strides=configuration_manager.pool_op_kernel_sizes[1:]
        if len(strides)>5:
            strides = strides[:5]
        while len(strides)<5:
            strides.append([1,1,1])
        return STUNet(num_input_channels, num_classes, depth=[2]*6, dims= [64 * x for x in [1, 2, 4, 8, 16, 16]], 
                      pool_op_kernel_sizes=strides, conv_kernel_sizes=kernel_sizes, enable_deep_supervision=enable_deep_supervision)

class STUNetTrainer_large_exp2(STUNetTrainer_large):
    pass

class STUNetTrainer_large_exp3(STUNetTrainer_large):
    pass

class STUNetTrainer_large_lr0001(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.initial_lr = 1e-3

class STUNetTrainer_large_bs6(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_iterations_per_epoch = 80

class STUNetTrainer_large_ep500(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 500

class STUNetTrainer_large_ep1500(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 1500

class STUNetTrainer_large_ep2k(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 2000

class STUNetTrainer_large_ep3k(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 3000

class STUNetTrainer_large_ep4k(STUNetTrainer_large):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 4000

class STUNetTrainer_huge(STUNetTrainer):
    @staticmethod
    def build_network_architecture(plans_manager,
                                   dataset_json,
                                   configuration_manager,
                                   num_input_channels,
                                   enable_deep_supervision: bool = True) -> nn.Module: 
        label_manager = plans_manager.get_label_manager(dataset_json)
        num_classes=label_manager.num_segmentation_heads
        kernel_sizes = [[3,3,3]] * 6
        strides=configuration_manager.pool_op_kernel_sizes[1:]
        if len(strides)>5:
            strides = strides[:5]
        while len(strides)<5:
            strides.append([1,1,1])
        return STUNet(num_input_channels, num_classes, depth=[3]*6, dims= [96 * x for x in [1, 2, 4, 8, 16, 16]], 
                      pool_op_kernel_sizes=strides, conv_kernel_sizes=kernel_sizes, enable_deep_supervision=enable_deep_supervision)

class STUNetTrainer_huge_ep4k(STUNetTrainer_huge):
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
                 device: torch.device = torch.device('cuda')):
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
        self.num_epochs = 4000

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.deep_supervision = True

class STUNet(nn.Module):
    def __init__(self, input_channels, num_classes, depth=[1,1,1,1,1,1], dims=[32, 64, 128, 256, 512, 512],
                 pool_op_kernel_sizes=None, conv_kernel_sizes=None, enable_deep_supervision=True):
        super().__init__()
        self.conv_op = nn.Conv3d
        self.input_channels = input_channels
        self.num_classes = num_classes
        
        self.final_nonlin = lambda x:x 
        self.decoder = Decoder()
        self.decoder.deep_supervision = enable_deep_supervision
        self.upscale_logits = False

        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes
        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([i // 2 for i in krnl])

        
        num_pool  = len(pool_op_kernel_sizes)
        
        assert num_pool == len(dims) - 1
        
        # encoder
        self.conv_blocks_context = nn.ModuleList()
        stage = nn.Sequential(BasicResBlock(input_channels, dims[0], self.conv_kernel_sizes[0], self.conv_pad_sizes[0], use_1x1conv=True), 
                              *[BasicResBlock(dims[0], dims[0], self.conv_kernel_sizes[0], self.conv_pad_sizes[0]) for _ in range(depth[0]-1)])
        self.conv_blocks_context.append(stage)
        for d in range(1, num_pool+1):
            stage = nn.Sequential(BasicResBlock(dims[d-1], dims[d], self.conv_kernel_sizes[d], self.conv_pad_sizes[d], stride=self.pool_op_kernel_sizes[d-1], use_1x1conv=True),
                *[BasicResBlock(dims[d], dims[d], self.conv_kernel_sizes[d], self.conv_pad_sizes[d]) for _ in range(depth[d]-1)])
            self.conv_blocks_context.append(stage)

        # upsample_layers
        self.upsample_layers = nn.ModuleList()
        for u in range(num_pool):
            upsample_layer = Upsample_Layer_nearest(dims[-1-u], dims[-2-u], pool_op_kernel_sizes[-1-u])
            self.upsample_layers.append(upsample_layer)

        # decoder
        self.conv_blocks_localization = nn.ModuleList()
        for u in range(num_pool):
            stage = nn.Sequential(BasicResBlock(dims[-2-u] * 2, dims[-2-u], self.conv_kernel_sizes[-2-u], self.conv_pad_sizes[-2-u], use_1x1conv=True),
                *[BasicResBlock(dims[-2-u], dims[-2-u], self.conv_kernel_sizes[-2-u], self.conv_pad_sizes[-2-u]) for _ in range(depth[-2-u]-1)])
            self.conv_blocks_localization.append(stage)
            
        # outputs    
        self.seg_outputs = nn.ModuleList()
        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(nn.Conv3d(dims[-2-ds], num_classes, kernel_size=1))

        self.upscale_logits_ops = []
        for usl in range(num_pool - 1):
            self.upscale_logits_ops.append(lambda x: x)
        

    def forward(self, x):
        skips = []
        seg_outputs = []
        
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)
            skips.append(x)

        x = self.conv_blocks_context[-1](x)

        for u in range(len(self.conv_blocks_localization)):
            x = self.upsample_layers[u](x)
            x = torch.cat((x, skips[-(u + 1)]), dim=1) 
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self.decoder.deep_supervision:
            return tuple([seg_outputs[-1]] + [i(j) for i, j in
                                              zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])])
        else:
            return seg_outputs[-1]


class BasicResBlock(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size=3, padding=1, stride=1, use_1x1conv=False):
        super().__init__()
        self.conv1 = nn.Conv3d(input_channels, output_channels, kernel_size, stride=stride, padding=padding)
        self.norm1 = nn.InstanceNorm3d(output_channels, affine=True)
        self.act1 = nn.LeakyReLU(inplace=True)
        
        self.conv2 = nn.Conv3d(output_channels, output_channels, kernel_size, padding=padding)
        self.norm2 = nn.InstanceNorm3d(output_channels, affine=True)
        self.act2 = nn.LeakyReLU(inplace=True)
        
        if use_1x1conv:
            self.conv3 = nn.Conv3d(input_channels, output_channels, kernel_size=1, stride=stride)
        else:
            self.conv3 = None
                  
    def forward(self, x):
        y = self.conv1(x)
        y = self.act1(self.norm1(y))  
        y = self.norm2(self.conv2(y))
        if self.conv3:
            x = self.conv3(x)
        y += x
        return self.act2(y)

class Upsample_Layer_nearest(nn.Module):
    def __init__(self, input_channels, output_channels, pool_op_kernel_size, mode='nearest'):
        super().__init__()
        self.conv = nn.Conv3d(input_channels, output_channels, kernel_size=1)
        self.pool_op_kernel_size = pool_op_kernel_size
        self.mode = mode
        
    def forward(self, x):
        x = nn.functional.interpolate(x, scale_factor=self.pool_op_kernel_size, mode=self.mode)
        x = self.conv(x)
        return x
nnunetv2/training/nnUNetTrainer/STUNetTrainer.py
uni-medical/A-Eval