Upload 12 files

utils/data_preparation.py

@@ -0,0 +1,241 @@
+import os
+from sklearn.model_selection import train_test_split
+import monai
+from monai.data import Dataset, DataLoader
+from data_transforms import define_transforms, define_transforms_loadonly
+import torch 
+import numpy as np
+from visualization import visualize_patient
+from monai.data import list_data_collate
+import pandas as pd
+
+
+def prepare_clinical_data(data_file, predictors):
+    
+    # read data file 
+    info = pd.read_excel(data_file, sheet_name=0)
+    
+    # convert to numerical 
+    info['CPS'] = info['CPS'].map({'A': 1, 'B': 2, 'C': 3})
+    info['T_involvment'] = info['T_involvment'].map({'< or = 50%': 1, '>50%': 2})
+    info['CLIP_Score'] = info['CLIP_Score'].map({'Stage_0': 0, 'Stage_1': 1, 'Stage_2': 2, 'Stage_3': 3, 'Stage_4': 4, 'Stage_5': 5, 'Stage_6': 6})
+    info['Okuda'] = info['Okuda'].map({'Stage I': 1, 'Stage II': 2, 'Stage III': 3})
+    info['TNM'] = info['TNM'].map({'Stage-I': 1, 'Stage-II': 2, 'Stage-IIIA': 3, 'Stage-IIIB': 4, 'Stage-IIIC': 5, 'Stage-IVA': 6, 'Stage-IVB': 7})
+    info['BCLC'] = info['BCLC'].map({'0': 0, 'Stage-A': 1, 'Stage-B': 2, 'Stage-C': 3, 'Stage-D': 4})
+    
+    # remove duplicates 
+    info.groupby("TCIA_ID").first() 
+    
+    # select columns 
+    info = info[['TCIA_ID'] + predictors].rename(columns={'TCIA_ID': "patient_id"})
+    
+    
+    return info
+    
+
+
+def preparare_train_test_txt(data_dir, test_patient_ratio=0.2, seed=1):
+    """
+    From a list of patients, split them into train and test and export list to .txt files 
+    """
+    
+    # split based on seed, write to txt files
+    patients = os.listdir(data_dir)
+    patients.remove("HCC-TACE-Seg_clinical_data-V2.xlsx")
+    patients = list(set(patients))
+    
+    # remove one patient with wrong labels
+    try:
+        patients.remove("HCC_017")
+        print("The patient HCC_017 is removed due to label issues including necrosis.")
+    except Exception as e:
+        pass 
+    
+    print("Total patients:", len(patients))
+    patients_train, patients_test = train_test_split(patients, test_size=test_patient_ratio, random_state=seed)
+    print("   There are", len(patients_train), "patients in training")
+    print("   There are", len(patients_test), "patients in test")
+
+    # export a copy
+    if not os.path.exists('train-test-split-seed' + str(seed)):
+        os.makedirs('train-test-split-seed' + str(seed))
+    with open(r'train-test-split-seed' + str(seed) + '/train.txt', 'w') as f:
+        f.write(','.join(patient for patient in patients_train))
+    with open(r'train-test-split-seed' + str(seed) + '/test.txt', 'w') as f:
+        f.write(','.join(patient for patient in patients_test))
+    
+    print("Files saved to", 'train-test-split-seed' + str(seed) + '/train.txt and train-test-split-seed' + str(seed) + '/test.txt')
+    return
+
+
+    
+
+def extract_file_path(patient_id, data_folder):
+    """
+    Given one patient's ID, obtain the file path of the image and mask data. 
+    If patient has multiple images, they are labeled as pre1, pre2, etc. 
+    """
+    path = os.path.join(data_folder, patient_id)
+    files = os.listdir(path)
+    patient_files = {}
+    count = 1
+    for file in files:
+      if "seg" in file or "Segmentation" in file:
+        patient_files["mask"] = os.path.join(path, file)
+      else:
+        patient_files["pre_" + str(count)] = os.path.join(path, file)
+        count += 1
+    return patient_files
+    
+    
+    
+def get_patient_dictionaries(txt_file, data_dir):
+    """
+    From .txt file that stores list of patients, look through data folders and extract a dictionary of patient data 
+    """
+    assert os.path.isfile(txt_file), "The file " + txt_file + " was not found. Please check your file directory."
+        
+    file = open(txt_file, "r")
+    patients = file.read().split(',')
+
+    data_dict = []
+
+    for patient_id in patients:
+
+      # get directories for mask and images
+      patient_files = extract_file_path(patient_id, data_dir)
+
+      # pair up each image with the mask
+      for key, value in patient_files.items():
+        if key != "mask":
+          data_dict.append(
+              {
+                "patient_id": patient_id,
+                "image": patient_files[key],
+                "mask": patient_files["mask"]
+              }
+          )
+
+    print("   There are", len(data_dict), "image-masks in this dataset.")
+    return data_dict
+    
+    
+
+
+def build_dataset(config, get_clinical=False):
+
+    def custom_collate_fn(batch):
+        """
+        Custom collate function to stack samples along the first dimension.
+
+        Args:
+            batch (list): List of dictionaries with keys "image" and "mask",
+                          where values are tensors of shape (N, 1, 512, 512).
+
+        Returns:
+            tuple: Tuple containing two tensors:
+                  - Stacked images of shape (B, 1, 512, 512)
+                  - Stacked masks of shape (B, 1, 512, 512)
+                  where B is the total number of samples in the batch.
+        """
+        # torch.manual_seed(1)
+        num_samples_to_select = config['BATCH_SIZE']
+
+        # Extract images and masks from the batch
+        images, masks = [], []
+        for sample in batch:
+            num_samples = min(sample["image"].shape[0], sample["mask"].shape[0])
+            random_indices = torch.randperm(num_samples)[:num_samples_to_select]
+            if "3D" in config['MODEL_NAME']: # 3D image
+                images.append(sample["image"][:,:512,:512,:]) # ensure image and mask same size
+                masks.append(sample["mask"][:,:512,:512,:])
+            else:
+                images.append(sample["image"][random_indices,:,:512,:512]) # ensure image and mask same size
+                masks.append(sample["mask"][random_indices,:,:512,:512])
+                #images.append(sample["image"][:,:,:512,:512]) # ensure image and mask same size
+                #masks.append(sample["mask"][:,:,:512,:512])
+
+        # Stack images and masks along the first dimension
+        try:
+            if "3D" not in config['MODEL_NAME']: # 3D image
+                concatenated_images = torch.cat(images, dim=0)
+                concatenated_masks = torch.cat(masks, dim=0)
+            else:
+                concatenated_images = torch.stack(images, dim=0)
+                concatenated_masks = torch.stack(masks, dim=0)
+        except Exception as e:
+            print("WARNING: not all images/masks are 512 by 512. Please check. ", images[0].shape, images[1].shape, masks[0].shape, masks[1].shape)
+            return None, None
+
+        # Return stacked images and masks as tensors
+        return {"image": concatenated_images, "mask": concatenated_masks}
+
+    # get list of training and test patient files
+    train_data_dict = get_patient_dictionaries(config['TRAIN_PATIENTS_FILE'], config['DATA_DIR'])
+    test_data_dict = get_patient_dictionaries(config['TEST_PATIENTS_FILE'], config['DATA_DIR'])
+    if config['ONESAMPLETESTRUN']: train_data_dict = train_data_dict[:2]
+    ttrain_data_dict, valid_data_dict = train_test_split(train_data_dict, test_size=config['VALID_PATIENT_RATIO'], shuffle=False, random_state=1) # must be false to match with linical data 
+    print("   Training patients:", len(ttrain_data_dict), " Validation patients:", len(valid_data_dict))
+    print("   Test patients:", len(test_data_dict))
+
+    # define data transformations
+    preprocessing_transforms_train, preprocessing_transforms_test, postprocessing_transforms = define_transforms(config)
+
+    # create data loaders
+    train_ds = Dataset(ttrain_data_dict, transform=preprocessing_transforms_train)
+    valid_ds = Dataset(valid_data_dict, transform=preprocessing_transforms_test)
+    test_ds = Dataset(test_data_dict, transform=preprocessing_transforms_test)
+
+    if "3D" in config['MODEL_NAME']:
+        train_loader = DataLoader(train_ds, batch_size=config['BATCH_SIZE'], collate_fn=custom_collate_fn, shuffle=False, num_workers=config['NUM_WORKERS']) 
+        valid_loader = DataLoader(valid_ds, batch_size=config['BATCH_SIZE'], collate_fn=custom_collate_fn, shuffle=False, num_workers=config['NUM_WORKERS'])
+        test_loader = DataLoader(test_ds, batch_size=config['BATCH_SIZE'], collate_fn=custom_collate_fn, shuffle=False, num_workers=config['NUM_WORKERS'])
+    else:
+        train_loader = DataLoader(train_ds, batch_size=1, shuffle=False, collate_fn=custom_collate_fn, num_workers=config['NUM_WORKERS']) #, pin_memory=torch.cuda.is_available())
+        valid_loader = DataLoader(valid_ds, batch_size=1, shuffle=False, collate_fn=custom_collate_fn, num_workers=config['NUM_WORKERS']) #, pin_memory=torch.cuda.is_available())
+        test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, collate_fn=custom_collate_fn, num_workers=config['NUM_WORKERS']) #, pin_memory=torch.cuda.is_available())
+
+    # get clinical data 
+    df_clinical_train = pd.DataFrame()
+    if get_clinical: 
+        # define transforms 
+        simple_transforms = define_transforms_loadonly()
+        simple_train_ds = Dataset(train_data_dict, transform=simple_transforms)
+        simple_train_loader = DataLoader(simple_train_ds, batch_size=config['BATCH_SIZE'], collate_fn=list_data_collate, shuffle=False, num_workers=config['NUM_WORKERS']) #, pin_memory=torch.cuda.is_available())
+        
+        # compute tumor ratio within liver 
+        df_clinical_train['patient_id'] = [p["patient_id"] for p in train_data_dict] 
+        ratios_train, ratios_test = [], []
+        for batch_data in simple_train_loader:
+            labels = batch_data["mask"]
+            ratio = torch.sum(labels == 2, dim=(1, 2, 3, 4)) / torch.sum(labels > 0, dim=(1, 2, 3, 4))
+            ratios_train.append(ratio.cpu().numpy()[0]) # [metatensor()]
+        df_clinical_train['tumor_ratio'] = ratios_train
+        
+        # get clinical features 
+        info = prepare_clinical_data(config['CLINICAL_DATA_FILE'], config['CLINICAL_PREDICTORS'])
+        df_clinical_train = pd.merge(df_clinical_train, info, on='patient_id', how="left")
+        df_clinical_train.fillna(df_clinical_train.median(), inplace=True)
+        df_clinical_train.set_index("patient_id", inplace=True)
+        
+    # visualize the data loader for one image to ensure correct formatting
+    print("Example data transformations:")
+    while True:
+        sample = preprocessing_transforms_train(train_data_dict[0])
+        if isinstance(sample, list): # depending on preprocessing, one sample may be [sample] or sample 
+            sample = sample[0]
+        if torch.sum(sample['mask'][-1]) == 0: continue
+        print(f"  image shape: {sample['image'].shape}")
+        print(f"  mask shape: {sample['mask'].shape}")
+        print(f"  mask values: {np.unique(sample['mask'])}")
+        #print(f"  image affine:\n{sample['image'].meta['affine']}")
+        print(f"  image min max: {np.min(sample['image']), np.max(sample['image'])}")
+        visualize_patient(sample['image'], sample['mask'], n_slices=3, z_dim_last="3D" in config['MODEL_NAME'], mask_channel=-1)
+        break
+
+    temp = monai.utils.first(test_loader)
+    print("Test loader shapes:", temp['image'].shape, temp['mask'].shape)
+        
+    return train_loader, valid_loader, test_loader, postprocessing_transforms, df_clinical_train
+    
+    

utils/data_transforms.py

@@ -0,0 +1,267 @@
+import monai 
+import cv2
+from monai.transforms import MapTransform
+import math 
+import numpy as np
+import torch
+import morphsnakes as ms
+import monai
+import nrrd
+import torchvision.transforms as transforms
+from monai.transforms import (
+    Activations, AsDiscreteD, AsDiscrete, Compose, CastToTypeD, RandSpatialCropd,
+    ToTensorD, CropForegroundD, Resized, GaussianSmoothD, 
+    LoadImageD, TransposeD, OrientationD, ScaleIntensityRangeD,
+    RandAffineD, ResizeWithPadOrCropd, ToTensor,
+    FillHoles, KeepLargestConnectedComponent, HistogramNormalizeD, NormalizeIntensityD
+)
+
+
+
+def define_transforms_loadonly():
+    transformations = Compose([
+        LoadImageD(keys=["mask"], reader="NrrdReader", ensure_channel_first=True),
+        ConvertMaskValues(keys=["mask"], keep_classes=["liver", "tumor"]),
+        ToTensor()
+    ])
+    return transformations
+
+
+def define_post_processing(config):
+      # Post-processing transforms
+      post_processing = [
+          # Apply softmax activation to convert logits to probabilities
+          Activations(sigmoid=True),
+          # Convert predicted probabilities to discrete values (0 or 1)
+          AsDiscrete(argmax=True, to_onehot=None if len(config['KEEP_CLASSES']) <= 2 else len(config['KEEP_CLASSES'])),
+          # Remove small connected components for 1=liver and 2=tumor
+          KeepLargestConnectedComponent(applied_labels=[1]),
+          # Fill holes in the binary mask for 1=liver and 2=tumor
+          FillHoles(applied_labels=[1]),
+          ToTensor()
+      ]
+
+      return Compose(post_processing)
+
+def define_transforms(config):
+
+      transformations_test = [
+              LoadImageD(keys=["image", "mask"], reader="NrrdReader", ensure_channel_first=True),
+              # Orient up and down
+              OrientationD(keys=["image", "mask"], axcodes="PLI"),
+              ToTensorD(keys=["image", "mask"])
+              # histogram equilization or normalization
+              # HistogramNormalizeD(keys=["image"], num_bins=256, min=0, max=1),
+              # Intensity normalization 
+              # NormalizeIntensityD(keys=["image"]),
+              #CastToTypeD(keys=["image"], dtype=torch.float32),
+              #CastToTypeD(keys=["mask"], dtype=torch.int32),
+          ]
+    
+      if config['MASKNONLIVER']:
+          transformations_test.extend(
+              [
+                MaskOutNonliver(mask_key="mask"),
+                CropForegroundD(keys=["image", "mask"], source_key="image", allow_smaller=True), 
+              ]
+          )
+          
+      transformations_test.append(
+          # Windowing based on liver parameters
+          ScaleIntensityRangeD(keys=["image"],
+            a_min=config['HU_RANGE'][0],
+            a_max=config['HU_RANGE'][1],
+            b_min=0.0, b_max=1.0, clip=True
+          )
+      )
+      
+      if config['PREPROCESSING'] == "clihe":
+          transformations_test.append(CLIHE(keys=["image"]))
+      
+      elif config['PREPROCESSING'] == "gaussian":
+          transformations_test.append(GaussianSmoothD(keys=["image"], sigma=0.5))
+
+      # convert labels to 0,1,2 instead of 0,1,2,3,4
+      transformations_test.append(ConvertMaskValues(keys=["mask"], keep_classes=config['KEEP_CLASSES']))
+
+      if len(config['KEEP_CLASSES']) > 2: # NEEDED FOR MULTICLASS  https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb
+          transformations_test.append(AsDiscreteD(keys=["mask"], to_onehot=len(config['KEEP_CLASSES']))) # (N, C, H, W) 2d; (1, C, H, W, Z)
+      
+      if "3D" not in config['MODEL_NAME']:
+          transformations_test.append(TransposeD(keys=["image", "mask"], indices=(3,0,1,2)))
+
+      # training transforms include data augmentation
+      transformations_train = transformations_test.copy()
+      if config['MASKNONLIVER']: transformations_test = transformations_test[:4] + transformations_test[5:] # do not crop to liver foregroudn 
+
+      if config['DATA_AUGMENTATION']:
+          if "3D" in config["MODEL_NAME"]:
+              transformations_train.append(
+                  RandAffineD(keys=["image", "mask"], prob=0.2, padding_mode="border",
+                              mode="bilinear", spatial_size=config['ROI_SIZE'],
+                              rotate_range=(0.15,0.15,0.15), #translate_range=(30,30,30), 
+                              scale_range=(0.1,0.1,0.1)))
+          else:
+              transformations_train.append(
+                  RandAffineD(keys=["image", "mask"], prob=0.2, padding_mode="border",
+                              mode="bilinear", #spatial_size=(512, 512),
+                              rotate_range=(0.15,0.15), #translate_range=(30,30), 
+                              scale_range=(0.1,0.1)))
+
+      transformations_train.extend(
+          [
+            RandSpatialCropd(keys=["image", "mask"], roi_size=config['ROI_SIZE'], random_size=False),
+            ResizeWithPadOrCropd(keys=["image", "mask"], spatial_size=config['ROI_SIZE'], method="end", mode='constant', value=0)
+          ]
+      )
+
+      postprocessing_transforms = define_post_processing(config)
+      preprocessing_transforms_test = Compose(transformations_test)
+      preprocessing_transforms_train = Compose(transformations_train)
+      preprocessing_transforms_train.set_random_state(seed=1)
+      preprocessing_transforms_test.set_random_state(seed=1)
+
+      return preprocessing_transforms_train, preprocessing_transforms_test, postprocessing_transforms
+      
+      
+
+class CLIHE(MapTransform):
+    def __init__(self, keys, allow_missing_keys=False):
+        super().__init__(allow_missing_keys)
+        self.keys = keys
+
+    def __call__(self, data):
+        for key in self.keys:
+            if len(data['image'].shape) > 3: # 3D image
+                data[key] = self.apply_clahe_3d(data[key]) # [B, 1, H, W, Z]
+            else:
+                data[key] = self.apply_clahe_2d(data[key]) # [B, 1, H, W, Z]
+        return data
+
+    def apply_clahe_3d(self, image):
+        image = np.asarray(image)
+        clahe_slices = []
+        for slice_idx in range(image.shape[-1]):
+            # Extract the current slice
+            slice_2d = image[0, :, :, slice_idx]
+
+            # Apply CLAHE to the current slice
+            # slice_2d = cv2.medianBlur(slice_2d, 5)
+            # slice_2d = cv2.anisotropicDiffusion(slice_2d, alpha=0.1, K=1, iterations=50)
+            # slice_2d = anisotropic_diffusion(slice_2d)
+            # slice_2d = cv2.Sobel(slice_2d, cv2.CV_64F, dx=1, dy=1, ksize=5)
+            clahe = cv2.createCLAHE(clipLimit=1, tileGridSize=(16,16))
+            slice_2d = clahe.apply(slice_2d.astype(np.uint8))
+            #cv2.threshold(clahe_slice, 155, 255, cv2.THRESH_BINARY)
+            kernel = np.ones((2,2), np.float32)/4
+            slice_2d = cv2.filter2D(slice_2d, -1, kernel)
+            #t = anisodiff2D(delta_t=0.2,kappa=50)
+            #slice_2d = t.fit(slice_2d)
+
+            # Append the CLAHE enhanced slice to the list
+            clahe_slices.append(slice_2d)
+
+            # Stack the CLAHE enhanced slices along the slice axis to form the 3D image
+            clahe_image = np.stack(clahe_slices, axis=-1)
+
+        return torch.from_numpy(clahe_image[None,:])
+
+    def apply_clahe_2d(self, image):
+        image = np.asarray(image)
+
+        clahe = cv2.createCLAHE(clipLimit=5)
+        clahe_slice = clahe.apply(image[0].astype(np.uint8))
+
+        return torch.from_numpy(clahe_slice)
+
+
+
+class GaussianFilter(MapTransform):
+    def __init__(self, keys, allow_missing_keys=False):
+        super().__init__(allow_missing_keys)
+        self.keys = keys
+
+    def __call__(self, data):
+        for key in self.keys:
+            if len(data['image'].shape) > 3: # 3D image
+                data[key] = self.apply_clahe_3d(data[key]) # [B, 1, H, W, Z]
+            else:
+                data[key] = self.apply_clahe_2d(data[key]) # [B, 1, H, W, Z]
+        return data
+
+    def apply_clahe_3d(self, image):
+        image = np.asarray(image)
+        clahe_slices = []
+        for slice_idx in range(image.shape[-1]):
+            # Extract the current slice
+            slice_2d = image[0, :, :, slice_idx]
+
+            # Apply CLAHE to the current slice
+            kernel = np.ones((3,3), np.float32)/9
+            slice_2d = cv2.filter2D(slice_2d, -1, kernel)
+            
+            # Append the CLAHE enhanced slice to the list
+            clahe_slices.append(slice_2d)
+
+            # Stack the CLAHE enhanced slices along the slice axis to form the 3D image
+            clahe_image = np.stack(clahe_slices, axis=-1)
+
+        return torch.from_numpy(clahe_image[None,:])
+
+    def apply_clahe_2d(self, image):
+        image = np.asarray(image)
+
+        kernel = np.ones((3,3), np.float32)/9
+        slice_2d = cv2.filter2D(image, -1, kernel)
+
+        return torch.from_numpy(slice_2d)
+
+
+class Morphsnakes(MapTransform):
+    # https://github.com/pmneila/morphsnakes/blob/master/morphsnakes.py
+    def __init__(self, allow_missing_keys=False):
+        super().__init__(allow_missing_keys)
+
+    def __call__(self, data):
+        if np.sum(data['mask'][-1]) > 0:
+            res = ms.morphological_chan_vese(data['image'][0], iterations=2, init_level_set=data['mask'][-1])
+            data['mask'] = res
+        return data
+        
+        
+class MaskOutNonliver(MapTransform):
+    def __init__(self, allow_missing_keys=False, mask_key="mask"):
+        super().__init__(allow_missing_keys)
+        self.mask_key = mask_key
+
+    def __call__(self, data):
+        # mask out non-liver regions of an image
+        # non-liver regions are liver, tumor, or portal vein
+        if data[self.mask_key].shape != data['image'].shape:
+            return data
+        data['image'][data[self.mask_key] >= 4] = -1000
+        data['image'][data[self.mask_key] <= 0] = -1000
+        return data
+        
+        
+class ConvertMaskValues(MapTransform):
+    def __init__(self, keys, allow_missing_keys=False, keep_classes=["normal", "liver", "tumor"]):
+        super().__init__(keys, allow_missing_keys)
+        self.keep_classes = keep_classes
+
+    def __call__(self, data):
+        # original labels: 0 for normal region, 1 for liver, 2 for tumor mass, 3 for portal vein, and 4 for abdominal aorta.
+        # converted labels: 0 for normal region and abdominal aorta, 1 for liver and portal vein, 2 for tumor mass
+
+        for key in self.keys:
+          data[key][data[key] > 4] = 4 # one patient had class label = 5, converted to 4
+          if key in data:
+            if "liver" not in self.keep_classes:
+                data[key][data[key] == 1] = 0
+            if "tumor" not in self.keep_classes:
+                data[key][data[key] == 2] = 1
+            if "portal vein" not in self.keep_classes:
+                data[key][data[key] == 3] = 1
+            if "abdominal aorta" not in self.keep_classes:
+                data[key][data[key] >= 4] = 0
+        return data

utils/inference.py

@@ -0,0 +1,155 @@
+import numpy as np
+import torch
+from monai.transforms import (
+    Activations, AsDiscreteD, AsDiscrete, Compose, ToTensorD, 
+    GaussianSmoothD, LoadImageD, TransposeD, OrientationD, ScaleIntensityRangeD,
+    ToTensor, FillHoles, KeepLargestConnectedComponent, NormalizeIntensityD
+)
+from nrrd import read 
+from visualization import visualize_results
+from data_preparation import get_patient_dictionaries
+from monai.data import Dataset, DataLoader
+import os
+from data_transforms import ConvertMaskValues, MaskOutNonliver
+from pipeline import build_model, evaluate 
+    
+def run_sequential_inference(txt_file, config_liver, config_tumor, eval_metrics, output_dir, only_tumor=False, export=True):
+
+    def custom_collate_fn(batch):
+        num_samples_to_select = config_liver['BATCH_SIZE']
+
+        # Extract images and masks from the batch,  ensure image and mask same size
+        images, masks, pred_liver = [], [], []
+        for sample in batch:
+            num_samples = min(sample["image"].shape[0], sample["mask"].shape[0])
+            random_indices = torch.randperm(num_samples)[:num_samples_to_select]
+            images.append(sample["image"][:,:512,:512,:]) 
+            masks.append(sample["mask"][:,:512,:512,:])
+
+        # Stack images and masks along the first dimension
+        try:
+            concatenated_images = torch.stack(images, dim=0)
+            concatenated_masks = torch.stack(masks, dim=0)
+        except Exception as e:
+            print("WARNING: not all images/masks are 512 by 512. Please check. ", images[0].shape, images[1].shape, masks[0].shape, masks[1].shape)
+            return None, None
+
+        # Return stacked images and masks as tensors
+        if "pred_liver" in sample.keys():
+            return {"image": concatenated_images, "mask": concatenated_masks, "pred_liver": sample["pred_liver"]}
+        else:
+            return {"image": concatenated_images, "mask": concatenated_masks}
+
+    ### Model preparation 
+    print("")
+    print("Loading models....")
+    liver_model = build_model(config_liver)
+    tumor_model = build_model(config_tumor)
+
+    #### Data preparation 
+    print("")
+    print("Loading test data....")
+    test_data_dict = get_patient_dictionaries(txt_file=txt_file, data_dir=config_liver['DATA_DIR'])
+    print("   Number of test patients:", len(test_data_dict))
+ 
+    # assign output file names and paths 
+    export_file_metadata = []
+    if not os.path.exists(output_dir): os.makedirs(output_dir)
+    for patient_dict in test_data_dict:
+        patient_folder = os.path.join(output_dir, patient_dict['patient_id'])
+        if not os.path.exists(patient_folder): os.makedirs(patient_folder)
+        patient_dict['pred_liver'] = os.path.join(patient_folder, "liver_segmentation.nrrd")
+        patient_dict['pred_tumor'] = os.path.join(patient_folder, "tumor_segmentation.nrrd")
+        export_file_metadata.append(read(patient_dict['image'])[1])
+    
+    #### Liver segmentation 
+    # define liver data loading and preprocessing 
+    if not only_tumor:
+        print("")
+        print("Producing liver segmentations....")
+        liver_preprocessing = Compose([
+            LoadImageD(keys=["image", "mask"], reader="NrrdReader", ensure_channel_first=True),
+            OrientationD(keys=["image", "mask"], axcodes="PLI"),
+            ScaleIntensityRangeD(keys=["image"],
+                a_min=config_liver['HU_RANGE'][0],
+                a_max=config_liver['HU_RANGE'][1],
+                b_min=0.0, b_max=1.0, clip=True
+            ),
+            ConvertMaskValues(keys=["mask"], keep_classes=["liver"]),
+            ToTensorD(keys=["image", "mask"])
+        ])
+    
+        liver_postprocessing = Compose([
+            Activations(sigmoid=True),
+            AsDiscrete(argmax=True, to_onehot=None),
+            KeepLargestConnectedComponent(applied_labels=[1]),
+            FillHoles(applied_labels=[1]),
+            ToTensor()
+        ])
+        test_ds_liver = Dataset(test_data_dict, transform=liver_preprocessing)
+        test_ds_liver = DataLoader(test_ds_liver, batch_size=config_liver['BATCH_SIZE'], collate_fn=custom_collate_fn, shuffle=False, num_workers=config_liver['NUM_WORKERS'])
+    
+        # produce liver model results 
+        test_metrics_liver, sample_output_liver = evaluate(liver_model, test_ds_liver, eval_metrics, config_liver, postprocessing_transforms=liver_postprocessing, export_filenames = [p['pred_liver'] for p in test_data_dict], export_file_metadata=export_file_metadata)
+
+        print("")
+        print("==============================")
+        print("Liver segmentation test performance ....")
+        for key, value in test_metrics_liver.items():
+            print(f'   {key.replace("_avg", "_liver")}: {value:.3f}')
+        print("==============================")
+        
+    ##### Tumor segmentation 
+    print("")
+    print("Producing tumor segmentations....")
+    
+    # define tumor loading and preprocessing
+    tumor_preprocessing = Compose([
+        LoadImageD(keys=["image", "mask", "pred_liver"], reader="NrrdReader", ensure_channel_first=True),
+        OrientationD(keys=["image", "mask"], axcodes="PLI"),
+        MaskOutNonliver(mask_key="pred_liver"), # note that liver's predicted segmentation is used to crop to the liver region 
+        ScaleIntensityRangeD(keys=["image"],
+            a_min=config_tumor['HU_RANGE'][0],
+            a_max=config_tumor['HU_RANGE'][1],
+            b_min=0.0, b_max=1.0, clip=True
+        ),
+        ConvertMaskValues(keys=["mask"], keep_classes=["liver", "tumor"]), # format mask for measuring test performance 
+        AsDiscreteD(keys=["mask"], to_onehot=3),           # format mask for measuring test performance 
+        ToTensorD(keys=["image", "mask", "pred_liver"])
+    ])
+
+    tumor_postprocessing = Compose([
+        Activations(sigmoid=True),
+        AsDiscrete(argmax=True, to_onehot=3),
+        ToTensor()
+    ])
+ 
+    test_ds_tumor = Dataset(test_data_dict, transform=tumor_preprocessing)
+    test_ds_tumor = DataLoader(test_ds_tumor, batch_size=config_tumor['BATCH_SIZE'], collate_fn=custom_collate_fn, shuffle=False, num_workers=config_tumor['NUM_WORKERS'])
+
+    test_metrics_tumor, sample_output_tumor = evaluate(tumor_model, test_ds_tumor, eval_metrics, config_tumor, tumor_postprocessing, use_liver_seg = True, export_filenames = [p['pred_tumor'] for p in test_data_dict] if export else [], export_file_metadata=export_file_metadata)
+
+    print("")
+    print("==============================")
+    print("Tumor segmentation test performance ....")
+    for key, value in test_metrics_tumor.items():
+        if "class2" in key: 
+            print(f'   {key.replace("_class2", "_tumor")}: {value:.3f}')
+    print("==============================")
+    print("")
+
+    #### Visualization 
+
+    # combine liver and tumor segmentations into one segmentation output
+    if not only_tumor: sample_output_tumor[2][0][1] = sample_output_liver[2][0][0]
+
+    # visualization 
+    print("")
+    if not only_tumor:
+      visualize_results(sample_output_liver[0][0].cpu(), sample_output_tumor[1][0].cpu(), sample_output_tumor[2][0].cpu(), n_slices=5, title="")
+    else:
+      visualize_results(sample_output_tumor[0][0].cpu(), sample_output_tumor[1][0].cpu(), sample_output_tumor[2][0].cpu(), n_slices=5, title="")
+
+    return 
+
+

utils/loss.py

@@ -0,0 +1,153 @@
+import warnings
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.loss import _Loss
+
+from monai.losses.dice import DiceLoss
+from monai.losses.focal_loss import FocalLoss
+from monai.networks import one_hot
+from monai.utils import DiceCEReduction, LossReduction, Weight, deprecated_arg, look_up_option, pytorch_after
+
+
+
+##### Adapted from Monai DiceFocalLoss 
+class WeaklyDiceFocalLoss(_Loss):
+    """
+    Compute Dice loss,  Focal Loss, and weakly supervised loss from clinical predictor, and return the weighted sum of these three losses.
+    
+    ``gamma`` and ``lambda_focal`` are only used for the focal loss.
+    ``include_background``, ``weight`` and ``reduction`` are used for both losses
+    and other parameters are only used for dice loss.
+    """
+
+    def __init__(
+        self,
+        include_background: bool = True,
+        to_onehot_y: bool = False,
+        sigmoid: bool = False,
+        softmax: bool = False,
+        other_act: Callable | None = None,
+        squared_pred: bool = False,
+        jaccard: bool = False,
+        reduction: str = "mean",
+        smooth_nr: float = 1e-5,
+        smooth_dr: float = 1e-5,
+        batch: bool = False,
+        gamma: float = 2.0,
+        focal_weight: Sequence[float] | float | int | torch.Tensor | None = None,
+        weight: Sequence[float] | float | int | torch.Tensor | None = None,
+        lambda_dice: float = 1.0,
+        lambda_focal: float = 1.0,
+        lambda_weak: float = 1.0,
+    ) -> None:
+        """
+        Args:
+            include_background: if False channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y: whether to convert the ``target`` into the one-hot format,
+                using the number of classes inferred from `input` (``input.shape[1]``). Defaults to False.
+            sigmoid: if True, apply a sigmoid function to the prediction, only used by the `DiceLoss`,
+                don't need to specify activation function for `FocalLoss`.
+            softmax: if True, apply a softmax function to the prediction, only used by the `DiceLoss`,
+                don't need to specify activation function for `FocalLoss`.
+            other_act: callable function to execute other activation layers, Defaults to ``None``.
+                for example: `other_act = torch.tanh`. only used by the `DiceLoss`, not for `FocalLoss`.
+            squared_pred: use squared versions of targets and predictions in the denominator or not.
+            jaccard: compute Jaccard Index (soft IoU) instead of dice or not.
+            reduction: {``"none"``, ``"mean"``, ``"sum"``}
+                Specifies the reduction to apply to the output. Defaults to ``"mean"``.
+
+                - ``"none"``: no reduction will be applied.
+                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
+                - ``"sum"``: the output will be summed.
+
+            smooth_nr: a small constant added to the numerator to avoid zero.
+            smooth_dr: a small constant added to the denominator to avoid nan.
+            batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
+                Defaults to False, a Dice loss value is computed independently from each item in the batch
+                before any `reduction`.
+            gamma: value of the exponent gamma in the definition of the Focal loss.
+            weight: weights to apply to the voxels of each class. If None no weights are applied.
+                The input can be a single value (same weight for all classes), a sequence of values (the length
+                of the sequence should be the same as the number of classes).
+            lambda_dice: the trade-off weight value for dice loss. The value should be no less than 0.0.
+                Defaults to 1.0.
+            lambda_focal: the trade-off weight value for focal loss. The value should be no less than 0.0.
+                Defaults to 1.0.
+            lambda_weak: the trade-off weight value for weakly supervised loss. The value should be no less than 0.0
+                Defaults to 0.2. 
+
+        """
+        super().__init__()
+        weight = focal_weight if focal_weight is not None else weight
+        self.dice = DiceLoss(
+            include_background=include_background,
+            to_onehot_y=False,
+            sigmoid=sigmoid,
+            softmax=softmax,
+            other_act=other_act,
+            squared_pred=squared_pred,
+            jaccard=jaccard,
+            reduction=reduction,
+            smooth_nr=smooth_nr,
+            smooth_dr=smooth_dr,
+            batch=batch,
+            weight=weight,
+        )
+        self.focal = FocalLoss(
+            include_background=include_background, to_onehot_y=False, gamma=gamma, weight=weight, reduction=reduction
+        )
+        if lambda_dice < 0.0:
+            raise ValueError("lambda_dice should be no less than 0.0.")
+        if lambda_focal < 0.0:
+            raise ValueError("lambda_focal should be no less than 0.0.")
+        if lambda_weak < 0.0:
+            raise ValueError("lambda_weak should be no less than 0.0.")
+        self.lambda_dice = lambda_dice
+        self.lambda_focal = lambda_focal
+        self.to_onehot_y = to_onehot_y
+        self.lambda_weak = lambda_weak
+
+
+    def compute_weakly_supervised_loss(self, input: torch.Tensor, weaktarget: torch.Tensor) -> torch.Tensor:
+        # compute ratio of tumor/liver in the predicted mask
+        tumor_pixels = torch.sum(input[:, -1, ...], dim=(1, 2, 3))
+        liver_pixels = torch.sum(input[:, -2, ...], dim=(1, 2, 3)) + tumor_pixels
+        predicted_ratio = tumor_pixels / liver_pixels 
+        loss = torch.mean((predicted_ratio - weaktarget) ** 2)
+        return loss
+        
+        
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor, weaktarget: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input: the shape should be BNH[WD]. The input should be the original logits
+                due to the restriction of ``monai.losses.FocalLoss``.
+            target: the shape should be BNH[WD] or B1H[WD].
+
+        Raises:
+            ValueError: When number of dimensions for input and target are different.
+            ValueError: When number of channels for target is neither 1 nor the same as input.
+
+        """
+        if len(input.shape) != len(target.shape):
+            raise ValueError(
+                "the number of dimensions for input and target should be the same, "
+                f"got shape {input.shape} and {target.shape}."
+            )
+        if self.to_onehot_y:
+            n_pred_ch = input.shape[1]
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
+            else:
+                target = one_hot(target, num_classes=n_pred_ch)
+        dice_loss = self.dice(input, target)
+        focal_loss = self.focal(input, target)
+        weak_loss = self.compute_weakly_supervised_loss(input, weaktarget)
+        total_loss: torch.Tensor = self.lambda_dice * dice_loss + self.lambda_focal * focal_loss + self.lambda_weak * weak_loss
+        return total_loss

utils/models.py

@@ -0,0 +1,670 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# 2D: net = UNet2D(1,2,pab_channels=64,use_batchnorm=True)
+# 3D: net = UNet3D(1,2,pab_channels=32,use_batchnorm=True)
+
+class _NonLocalBlockND(nn.Module):
+    def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):
+        super(_NonLocalBlockND, self).__init__()
+
+        assert dimension in [1, 2, 3]
+
+        self.dimension = dimension
+        self.sub_sample = sub_sample
+
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+
+        if self.inter_channels is None:
+            self.inter_channels = in_channels // 2
+            if self.inter_channels == 0:
+                self.inter_channels = 1
+
+        if dimension == 3:
+            conv_nd = nn.Conv3d
+            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
+            bn = nn.BatchNorm3d
+        elif dimension == 2:
+            conv_nd = nn.Conv2d
+            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
+            bn = nn.BatchNorm2d
+        else:
+            conv_nd = nn.Conv1d
+            max_pool_layer = nn.MaxPool1d(kernel_size=(2))
+            bn = nn.BatchNorm1d
+
+        self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                         kernel_size=1, stride=1, padding=0)
+
+        if bn_layer:
+            self.W = nn.Sequential(
+                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                        kernel_size=1, stride=1, padding=0),
+                bn(self.in_channels)
+            )
+            nn.init.constant_(self.W[1].weight, 0)
+            nn.init.constant_(self.W[1].bias, 0)
+        else:
+            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                             kernel_size=1, stride=1, padding=0)
+            nn.init.constant_(self.W.weight, 0)
+            nn.init.constant_(self.W.bias, 0)
+
+        self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                             kernel_size=1, stride=1, padding=0)
+
+        self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                           kernel_size=1, stride=1, padding=0)
+
+        if sub_sample:
+            self.g = nn.Sequential(self.g, max_pool_layer)
+            self.phi = nn.Sequential(self.phi, max_pool_layer)
+
+    def forward(self, x):
+        '''
+        :param x: (b, c, t, h, w)
+        :return:
+        '''
+
+        batch_size = x.size(0)
+
+        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
+        g_x = g_x.permute(0, 2, 1)
+
+        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
+        theta_x = theta_x.permute(0, 2, 1)
+        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
+        f = torch.matmul(theta_x, phi_x)
+        N = f.size(-1)
+        f_div_C = f / N
+
+        y = torch.matmul(f_div_C, g_x)
+        y = y.permute(0, 2, 1).contiguous()
+        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
+        W_y = self.W(y)
+        z = W_y + x
+
+        return z
+
+
+class NONLocalBlock1D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock1D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=1, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+class NONLocalBlock2D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock2D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=2, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+class NONLocalBlock3D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock3D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=3, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+
+class Conv2dReLU(nn.Sequential):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            padding=0,
+            stride=1,
+            use_batchnorm=True,
+    ):
+
+        if use_batchnorm == "inplace" and InPlaceABN is None:
+            raise RuntimeError(
+                "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
+                + "To install see: https://github.com/mapillary/inplace_abn"
+            )
+
+        conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=not (use_batchnorm),
+        )
+        relu = nn.ReLU(inplace=True)
+
+        if use_batchnorm == "inplace":
+            bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
+            relu = nn.Identity()
+
+        elif use_batchnorm and use_batchnorm != "inplace":
+            bn = nn.BatchNorm2d(out_channels)
+
+        else:
+            bn = nn.Identity()
+
+        super(Conv2dReLU, self).__init__(conv, bn, relu)
+
+class Conv3dReLU(nn.Sequential):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            padding=0,
+            stride=1,
+            use_batchnorm=True,
+    ):
+
+        if use_batchnorm == "inplace" and InPlaceABN is None:
+            raise RuntimeError(
+                "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
+                + "To install see: https://github.com/mapillary/inplace_abn"
+            )
+
+        conv = nn.Conv3d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=not (use_batchnorm),
+        )
+        relu = nn.ReLU(inplace=True)
+
+        if use_batchnorm == "inplace":
+            bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
+            relu = nn.Identity()
+
+        elif use_batchnorm and use_batchnorm != "inplace":
+            bn = nn.BatchNorm3d(out_channels)
+
+        else:
+            bn = nn.Identity()
+
+        super(Conv3dReLU, self).__init__(conv, bn, relu)
+class PAB2D(nn.Module):
+    def __init__(self, in_channels, out_channels, pab_channels=64):
+        super(PAB2D, self).__init__()
+        # Series of 1x1 conv to generate attention feature maps
+        self.pab_channels = pab_channels
+        self.in_channels = in_channels
+        self.top_conv = nn.Conv2d(in_channels, pab_channels, kernel_size=1)
+        self.center_conv = nn.Conv2d(in_channels, pab_channels, kernel_size=1)
+        self.bottom_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
+        self.map_softmax = nn.Softmax(dim=1)
+        self.out_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        bsize = x.size()[0]
+        h = x.size()[2]
+        w = x.size()[3]
+        x_top = self.top_conv(x)
+        x_center = self.center_conv(x)
+        x_bottom = self.bottom_conv(x)
+
+        x_top = x_top.flatten(2)
+        x_center = x_center.flatten(2).transpose(1, 2)
+        x_bottom = x_bottom.flatten(2).transpose(1, 2)
+
+        sp_map = torch.matmul(x_center, x_top)
+        sp_map = self.map_softmax(sp_map.view(bsize, -1)).view(bsize, h*w, h*w)
+        sp_map = torch.matmul(sp_map, x_bottom)
+        sp_map = sp_map.reshape(bsize, self.in_channels, h, w)
+        x = x + sp_map
+        x = self.out_conv(x)
+        # print('x_top',x_top.shape,'x_center',x_center.shape,'x_bottom',x_bottom.shape,'x',x.shape,'sp_map',sp_map.shape)
+        return x
+
+class MFAB2D(nn.Module):
+    def __init__(self, in_channels, skip_channels, out_channels, use_batchnorm=True, reduction=16):
+        # MFAB is just a modified version of SE-blocks, one for skip, one for input
+        super(MFAB2D, self).__init__()
+        self.hl_conv = nn.Sequential(
+            Conv2dReLU(
+                in_channels,
+                in_channels,
+                kernel_size=3,
+                padding=1,
+                use_batchnorm=use_batchnorm,
+            ),
+            Conv2dReLU(
+                in_channels,
+                skip_channels,
+                kernel_size=1,
+                use_batchnorm=use_batchnorm,
+            )
+        )
+        self.SE_ll = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(skip_channels, skip_channels // reduction, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(skip_channels // reduction, skip_channels, 1),
+            nn.Sigmoid(),
+        )
+        self.SE_hl = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(skip_channels, skip_channels // reduction, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(skip_channels // reduction, skip_channels, 1),
+            nn.Sigmoid(),
+        )
+        self.conv1 = Conv2dReLU(
+            skip_channels + skip_channels,  # we transform C-prime form high level to C from skip connection
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            use_batchnorm=use_batchnorm,
+        )
+        self.conv2 = Conv2dReLU(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            use_batchnorm=use_batchnorm,
+        )
+
+    def forward(self, x, skip=None):
+        x = self.hl_conv(x)
+        x = F.interpolate(x, scale_factor=2, mode="nearest")
+        attention_hl = self.SE_hl(x)
+        if skip is not None:
+            attention_ll = self.SE_ll(skip)
+            attention_hl = attention_hl + attention_ll
+            x = x * attention_hl
+            x = torch.cat([x, skip], dim=1)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+class PAB3D(nn.Module):
+    def __init__(self, in_channels, out_channels, pab_channels=64):
+        super(PAB3D, self).__init__()
+        # Series of 1x1 conv to generate attention feature maps
+        self.pab_channels = pab_channels
+        self.in_channels = in_channels
+        self.top_conv = nn.Conv3d(in_channels, pab_channels, kernel_size=1)
+        self.center_conv = nn.Conv3d(in_channels, pab_channels, kernel_size=1)
+        self.bottom_conv = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
+        self.map_softmax = nn.Softmax(dim=1)
+        self.out_conv = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        bsize = x.size()[0]
+        h = x.size()[2]
+        w = x.size()[3]
+        d = x.size()[4]
+        x_top = self.top_conv(x)
+        x_center = self.center_conv(x)
+        x_bottom = self.bottom_conv(x)
+
+        x_top = x_top.flatten(2)
+        x_center = x_center.flatten(2).transpose(1, 2)
+        x_bottom = x_bottom.flatten(2).transpose(1, 2)
+        sp_map = torch.matmul(x_center, x_top)
+        sp_map = self.map_softmax(sp_map.view(bsize, -1)).view(bsize, h*w*d, h*w*d)
+        sp_map = torch.matmul(sp_map, x_bottom)
+        sp_map = sp_map.reshape(bsize, self.in_channels, h, w, d)
+        x = x + sp_map
+        x = self.out_conv(x)
+        # print('x_top',x_top.shape,'x_center',x_center.shape,'x_bottom',x_bottom.shape,'x',x.shape,'sp_map',sp_map.shape)
+        return x
+
+class MFAB3D(nn.Module):
+    def __init__(self, in_channels, skip_channels, out_channels, use_batchnorm=True, reduction=16):
+        # MFAB is just a modified version of SE-blocks, one for skip, one for input
+        super(MFAB3D, self).__init__()
+        self.hl_conv = nn.Sequential(
+            Conv3dReLU(
+                in_channels,
+                in_channels,
+                kernel_size=3,
+                padding=1,
+                use_batchnorm=use_batchnorm,
+            ),
+            Conv3dReLU(
+                in_channels,
+                skip_channels,
+                kernel_size=1,
+                use_batchnorm=use_batchnorm,
+            )
+        )
+        self.SE_ll = nn.Sequential(
+            nn.AdaptiveAvgPool3d(1),
+            nn.Conv3d(skip_channels, skip_channels // reduction, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(skip_channels // reduction, skip_channels, 1),
+            nn.Sigmoid(),
+        )
+        self.SE_hl = nn.Sequential(
+            nn.AdaptiveAvgPool3d(1),
+            nn.Conv3d(skip_channels, skip_channels // reduction, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(skip_channels // reduction, skip_channels, 1),
+            nn.Sigmoid(),
+        )
+        self.conv1 = Conv3dReLU(
+            skip_channels + skip_channels,  # we transform C-prime form high level to C from skip connection
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            use_batchnorm=use_batchnorm,
+        )
+        self.conv2 = Conv3dReLU(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            use_batchnorm=use_batchnorm,
+        )
+
+    def forward(self, x, skip=None):
+        x = self.hl_conv(x)
+        x = F.interpolate(x, scale_factor=2, mode="nearest")
+        attention_hl = self.SE_hl(x)
+        if skip is not None:
+            attention_ll = self.SE_ll(skip)
+            attention_hl = attention_hl + attention_ll
+            x = x * attention_hl
+            x = torch.cat([x, skip], dim=1)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+class DoubleConv2D(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class Down2D(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            NONLocalBlock2D(in_channels),
+            DoubleConv2D(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up2D(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv2D(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv2D(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv2D(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv2D, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class UNet2D(nn.Module):
+    def __init__(self, n_channels, n_classes, bilinear=True, pab_channels=64, use_batchnorm=True, aux_classifier = False):
+        super(UNet2D, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        self.inc = DoubleConv2D(n_channels, pab_channels)
+        self.down1 = Down2D(pab_channels, 2*pab_channels)
+        self.down2 = Down2D(2*pab_channels, 4*pab_channels)
+        self.down3 = Down2D(4*pab_channels, 8*pab_channels)
+        factor = 2 if bilinear else 1
+        self.down4 = Down2D(8*pab_channels, 16*pab_channels // factor)
+        self.pab = PAB2D(8*pab_channels,8*pab_channels)
+        self.up1 = Up2D(16*pab_channels, 8*pab_channels // factor, bilinear)
+        self.up2 = Up2D(8*pab_channels, 4*pab_channels // factor, bilinear)
+        self.up3 = Up2D(4*pab_channels, 2*pab_channels // factor, bilinear)
+        self.up4 = Up2D(2*pab_channels, pab_channels, bilinear)
+
+        self.mfab1 = MFAB2D(8*pab_channels,8*pab_channels,4*pab_channels,use_batchnorm)
+        self.mfab2 = MFAB2D(4*pab_channels,4*pab_channels,2*pab_channels,use_batchnorm)
+        self.mfab3 = MFAB2D(2*pab_channels,2*pab_channels,pab_channels,use_batchnorm)
+        self.mfab4 = MFAB2D(pab_channels,pab_channels,pab_channels,use_batchnorm)
+        self.outc = OutConv2D(pab_channels, n_classes)
+        
+        if aux_classifier == False:
+          self.aux = None
+        else:
+          # customize the auxiliary classification loss
+          # self.aux = nn.Sequential(nn.AdaptiveAvgPool2d(1),
+          #                          nn.Flatten(),
+          #                          nn.Dropout(p=0.1, inplace=True),
+          #                          nn.Linear(8*pab_channels, 16, bias=True),
+          #                          nn.Dropout(p=0.1, inplace=True),
+          #                          nn.Linear(16, n_classes, bias=True),
+          #                          nn.Softmax(1))
+
+          self.aux = nn.Sequential(
+                                   NONLocalBlock2D(8*pab_channels),
+                                   nn.Conv2d(8*pab_channels,1,1),
+                                   nn.InstanceNorm2d(1),
+                                   nn.ReLU(),       
+                                   nn.Flatten(),
+                                   nn.Linear(24*24, 16, bias=True),
+                                   nn.Dropout(p=0.2, inplace=True),
+                                   nn.Linear(16, n_classes, bias=True),
+                                   nn.Softmax(1))
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x5 = self.pab(x5)
+        
+        x = self.mfab1(x5,x4)
+        x = self.mfab2(x,x3)
+        x = self.mfab3(x,x2)
+        x = self.mfab4(x,x1)
+
+        # x = self.up1(x5, x4)
+        # x = self.up2(x, x3)
+        # x = self.up3(x, x2)
+        # x = self.up4(x, x1)
+        logits = self.outc(x)
+        logits = F.softmax(logits,1)
+
+        if self.aux ==None:
+          return logits
+        else:
+          aux = self.aux(x5)
+          return logits, aux
+          
+          
+          
+          
+class DoubleConv3D(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm3d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm3d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class Down3D(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool3d(2),
+            # NONLocalBlock3D(in_channels),
+            DoubleConv3D(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+class Up3D(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv3D(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose3d(in_channels , in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv3D(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv3D(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv3D, self).__init__()
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class UNet3D(nn.Module):
+    def __init__(self, n_channels, n_classes, bilinear=True, pab_channels=64, use_batchnorm=True, aux_classifier = False):
+        super(UNet3D, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+
+        self.inc = DoubleConv3D(n_channels, pab_channels)
+        self.down1 = Down3D(pab_channels, 2*pab_channels)
+        self.nnblock2 = NONLocalBlock3D(2*pab_channels)
+        self.down2 = Down3D(2*pab_channels, 4*pab_channels)
+        self.down3 = Down3D(4*pab_channels, 8*pab_channels)
+        factor = 2 if bilinear else 1
+        self.down4 = Down3D(8*pab_channels, 16*pab_channels // factor)
+        self.pab = PAB3D(8*pab_channels,8*pab_channels)
+        self.up1 = Up3D(16*pab_channels, 8*pab_channels // factor, bilinear)
+        self.up2 = Up3D(8*pab_channels, 4*pab_channels // factor, bilinear)
+        self.up3 = Up3D(4*pab_channels, 2*pab_channels // factor, bilinear)
+        self.up4 = Up3D(2*pab_channels, pab_channels, bilinear)
+
+        self.mfab1 = MFAB3D(8*pab_channels,8*pab_channels,4*pab_channels,use_batchnorm)
+        self.mfab2 = MFAB3D(4*pab_channels,4*pab_channels,2*pab_channels,use_batchnorm)
+        self.mfab3 = MFAB3D(2*pab_channels,2*pab_channels,pab_channels,use_batchnorm)
+        self.mfab4 = MFAB3D(pab_channels,pab_channels,pab_channels,use_batchnorm)
+        self.outc = OutConv3D(pab_channels, n_classes)
+
+        if aux_classifier == False:
+          self.aux = None
+        else:
+          # customize the auxiliary classification loss
+          # self.aux = nn.Sequential(nn.AdaptiveMaxPool3d(1),
+          #                          nn.Flatten(),
+          #                          nn.Dropout(p=0.1, inplace=True),
+          #                          nn.Linear(8*pab_channels, 16, bias=True),
+          #                          nn.Dropout(p=0.1, inplace=True),
+          #                          nn.Linear(16, n_classes, bias=True),
+          #                          nn.Softmax(1))
+          
+          self.aux = nn.Sequential(nn.Conv3d(8*pab_channels,1,1),   
+                                   nn.InstanceNorm3d(1),
+                                   nn.ReLU(),       
+                                   nn.Flatten(),
+                                   nn.Linear(16*16*2, 16, bias=True),
+                                   nn.Dropout(p=0.2, inplace=True),
+                                   nn.Linear(16, n_classes, bias=True),
+                                   nn.Softmax(1))
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        # x2 = self.nnblock2(x2)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x5 = self.pab(x5)
+        
+        x = self.mfab1(x5,x4)
+        x = self.mfab2(x,x3)
+        x = self.mfab3(x,x2)
+        x = self.mfab4(x,x1)
+
+        # x = self.up1(x5, x4)
+        # x = self.up2(x, x3)
+        # x = self.up3(x, x2)
+        # x = self.up4(x, x1)
+        logits = self.outc(x)
+        logits = F.softmax(logits,1)
+
+        if self.aux ==None:
+          return logits
+        else:
+          aux = self.aux(x5)
+          return logits, aux

utils/pipeline.py

@@ -0,0 +1,501 @@
+import logging
+import sys
+import tempfile
+from glob import glob
+from torchsummary import summary
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from torch.cuda.amp import autocast, GradScaler
+import torch.nn as nn
+import torchvision
+import monai
+from monai.metrics import DiceMetric, ConfusionMatrixMetric, MeanIoU
+from monai.visualize import plot_2d_or_3d_image
+from visualization import visualize_patient 
+from sliding_window import sw_inference
+from data_preparation import build_dataset 
+from models import UNet2D, UNet3D
+from loss import WeaklyDiceFocalLoss
+from sklearn.linear_model import LinearRegression
+from nrrd import write, read 
+import morphsnakes as ms
+from monai.data import decollate_batch
+
+    
+def build_optimizer(model, config):
+
+    if config['LOSS'] == "gdice":
+        loss_function = monai.losses.GeneralizedDiceLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'],
+            reduction="mean", to_onehot_y=True, sigmoid=True) if len(config['KEEP_CLASSES'])<=2 else monai.losses.GeneralizedDiceLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'], reduction="mean", to_onehot_y=False, softmax=True)
+    elif config['LOSS'] == 'cdice':
+        loss_function = monai.losses.DiceCELoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'],
+            reduction="mean", to_onehot_y=True, sigmoid=True) if len(config['KEEP_CLASSES'])<=2 else monai.losses.DiceCELoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'], reduction="mean", to_onehot_y=False, softmax=True)
+    elif config['LOSS'] == 'mdice':
+        loss_function = monai.losses.MaskedDiceLoss()
+    elif config['LOSS'] == 'wdice':
+        # Example with 3 classes (including the background: label 0).
+        # The distance between the background class (label 0) and the other classes is the maximum, equal to 1.
+        # The distance between class 1 and class 2 is 0.5.
+        dist_mat = np.array([[0.0, 1.0, 1.0], [1.0, 0.0, 0.5], [1.0, 0.5, 0.0]], dtype=np.float32)
+        loss_function = monai.losses.GeneralizedWassersteinDiceLoss(dist_matrix=dist_mat)
+    elif config['LOSS'] == "fdice":
+        loss_function = monai.losses.DiceFocalLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'], to_onehot_y=True, sigmoid=True) if len(config['KEEP_CLASSES'])<=2 else monai.losses.DiceFocalLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'], to_onehot_y=False, softmax=True)
+    elif config['LOSS'] == "wfdice":
+        loss_function = WeaklyDiceFocalLoss(include_background=config['EVAL_INCLUDE_BACKGROUND'], to_onehot_y=True, sigmoid=True, lambda_weak=config['LAMBDA_WEAK']) if len(config['KEEP_CLASSES'])<=2 else WeaklyDiceFocalLoss(include_background=config['EVAL_INCLUDE_BACKGROUND'], to_onehot_y=False, softmax=True, lambda_weak=config['LAMBDA_WEAK'])
+    else:
+        loss_function = monai.losses.DiceLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'],
+            reduction="mean", to_onehot_y=True, sigmoid=True, squared_pred=True) if len(config['KEEP_CLASSES'])<=2 else monai.losses.DiceLoss(
+            include_background=config['EVAL_INCLUDE_BACKGROUND'], reduction="mean", to_onehot_y=False, softmax=True, squared_pred=True)
+
+    eval_metrics = [
+        ("sensitivity", ConfusionMatrixMetric(include_background=config['EVAL_INCLUDE_BACKGROUND'], metric_name='sensitivity', reduction="mean_batch")),
+        ("specificity", ConfusionMatrixMetric(include_background=config['EVAL_INCLUDE_BACKGROUND'], metric_name='specificity', reduction="mean_batch")),
+        ("accuracy", ConfusionMatrixMetric(include_background=config['EVAL_INCLUDE_BACKGROUND'], metric_name='accuracy', reduction="mean_batch")),
+        ("dice", DiceMetric(include_background=config['EVAL_INCLUDE_BACKGROUND'], reduction="mean_batch")),
+        ("IoU", MeanIoU(include_background=config['EVAL_INCLUDE_BACKGROUND'], reduction="mean_batch"))
+    ]
+
+    optimizer = torch.optim.Adam(model.parameters(), config['LEARNING_RATE'], weight_decay=1e-5, amsgrad=True)
+    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['MAX_EPOCHS'])
+    return loss_function, optimizer, lr_scheduler, eval_metrics
+
+
+    
+def load_weights(model, config):
+    try:
+        model.load_state_dict(torch.load("checkpoints/" + config['PRETRAINED_WEIGHTS'] + ".pth", map_location=torch.device(config['DEVICE'])))
+        print("Model weights from", config['PRETRAINED_WEIGHTS'], "have been loaded")
+    except Exception as e:
+        try:
+            model.load_state_dict(torch.load(config['PRETRAINED_WEIGHTS'], map_location=torch.device(config['DEVICE'])))
+            print("Model weights from", config['PRETRAINED_WEIGHTS'], "have been loaded")
+        except Exception as e: # load 
+            print("WARNING: weights were not loaded. ", e)
+            pass    
+    
+    return model 
+
+
+def build_model(config):
+    
+    config = get_defaults(config)
+    
+    dropout_prob = config['DROPOUT']
+    
+    if "SegResNetVAE" in config["MODEL_NAME"]:
+        model = monai.networks.nets.SegResNetVAE(
+            input_image_size=config['ROI_SIZE'] if "3D" in config['MODEL_NAME'] else (config['ROI_SIZE'][0], config['ROI_SIZE'][1]),
+            vae_estimate_std=False, 
+            vae_default_std=0.3, 
+            vae_nz=256, 
+            spatial_dims=3 if "3D" in config["MODEL_NAME"] else 2,
+            blocks_down=[1, 2, 2, 4],
+            blocks_up=[1, 1, 1],
+            init_filters=16,
+            in_channels=1,
+            norm='instance',
+            out_channels=len(config['KEEP_CLASSES']),
+            dropout_prob=dropout_prob,
+        ).to(config['DEVICE'])
+    
+    elif "SegResNet" in config["MODEL_NAME"]:
+        model = monai.networks.nets.SegResNet(
+            spatial_dims=3 if "3D" in config["MODEL_NAME"] else 2,
+            blocks_down=[1, 2, 2, 4],
+            blocks_up=[1, 1, 1],
+            init_filters=16,
+            in_channels=1,
+            out_channels=len(config['KEEP_CLASSES']),
+            dropout_prob=dropout_prob,
+            norm="instance"
+        ).to(config['DEVICE'])
+    
+    elif "SwinUNETR" in config["MODEL_NAME"]:
+        model = monai.networks.nets.SwinUNETR(
+          img_size=config['ROI_SIZE'],
+          in_channels=1,
+          out_channels=len(config['KEEP_CLASSES']),
+          feature_size=48,
+          drop_rate=dropout_prob,
+          attn_drop_rate=0.0,
+          dropout_path_rate=0.0,
+          use_checkpoint=True
+       ).to(config['DEVICE'])
+    
+    elif "UNETR" in config["MODEL_NAME"]:
+        model = monai.networks.nets.UNETR(
+          img_size=config['ROI_SIZE'] if "3D" in config['MODEL_NAME'] else (config['ROI_SIZE'][0], config['ROI_SIZE'][1]),
+          in_channels=1,
+          out_channels=len(config['KEEP_CLASSES']),
+          feature_size=16,
+          hidden_size=256,
+          mlp_dim=3072,
+          num_heads=8,
+          pos_embed="perceptron",
+          norm_name="instance",
+          res_block=True,
+          dropout_rate=dropout_prob,
+      ).to(config['DEVICE'])
+    
+    elif "MANet" in config["MODEL_NAME"]:
+        if "2D" in config["MODEL_NAME"]:
+            model = UNet2D(
+                1, 
+                len(config['KEEP_CLASSES']), 
+                pab_channels=64, 
+                use_batchnorm=True
+                ).to(config['DEVICE'])
+        else:
+            model = UNet3D(
+                1, 
+                len(config['KEEP_CLASSES']), 
+                pab_channels=32, 
+                use_batchnorm=True
+                ).to(config['DEVICE'])
+    
+    elif "UNetPlusPlus" in config["MODEL_NAME"]:
+        model = monai.networks.nets.BasicUNetPlusPlus(
+            spatial_dims=3 if "3D" in config["MODEL_NAME"] else 2,
+            in_channels=1,
+            out_channels=len(config['KEEP_CLASSES']),
+            features=(32, 32, 64, 128, 256, 32),
+            norm="instance",
+            dropout=dropout_prob,
+        ).to(config['DEVICE'])
+    
+    elif "UNet1" in config['MODEL_NAME']:
+        model = monai.networks.nets.UNet(
+            spatial_dims=3 if "3D" in config["MODEL_NAME"] else 2,
+            in_channels=1,
+            out_channels=len(config['KEEP_CLASSES']),
+            channels=(16, 32, 64, 128, 256),
+            strides=(2, 2, 2, 2),
+            num_res_units=2,
+            norm="instance"
+        ).to(config['DEVICE'])
+    
+    elif "UNet2" in config['MODEL_NAME']:
+        model = monai.networks.nets.UNet(
+            spatial_dims=3 if "3D" in config["MODEL_NAME"] else 2,
+            in_channels=1,
+            out_channels=len(config['KEEP_CLASSES']),
+            channels=(32, 64, 128, 256),
+            strides=(2, 2, 2, 2),
+            num_res_units=4,
+            norm="instance"
+        ).to(config['DEVICE'])
+        
+    else:
+        print(config["MODEL_NAME"], "is not a valid model name")
+        return None
+
+    try:
+        if "3D" in config['MODEL_NAME']:
+            print(summary(model, input_size=(1, config['ROI_SIZE'][0], config['ROI_SIZE'][1], config['ROI_SIZE'][2])))
+        else:
+            print(summary(model, input_size=(1, config['ROI_SIZE'][0], config['ROI_SIZE'][1])))
+    except Exception as e:
+        print("could not load model summary:", e)
+
+    if config['PRETRAINED_WEIGHTS'] is not None and config['PRETRAINED_WEIGHTS']:
+        model = load_weights(model, config)
+    return model
+
+
+def train(model, train_loader, val_loader, loss_function, eval_metrics, optimizer, config,
+          scheduler=None, writer=None, postprocessing_transforms = None, weak_labels = None):
+
+    if writer is None: writer = SummaryWriter(log_dir="runs/" + config['EXPORT_FILE_NAME'])
+    best_metric, best_metric_epoch = -1, -1
+    prev_metric, patience, patience_counter = 1, config['EARLY_STOPPING_PATIENCE'], 0
+    if config['AUTOCAST']: scaler = GradScaler() # Initialize GradScaler for mixed precision training
+
+    for epoch in range(config['MAX_EPOCHS']):
+        print("-" * 10)
+        model.train()
+        epoch_loss, step = 0, 0
+        with tqdm(train_loader) as progress_bar:
+            for batch_data in progress_bar:
+                step += 1
+                inputs, labels = batch_data["image"].to(config['DEVICE']), batch_data["mask"].to(config['DEVICE'])
+                
+                # only train with batches that have tumor; skip those without tumor  
+                if config['TYPE'] == "tumor":
+                    if torch.sum(labels[:,-1]) == 0:
+                        continue
+                    
+                # check input shapes 
+                if inputs is None or labels is None:
+                    continue
+                if inputs.shape[-1] != labels.shape[-1] or inputs.shape[0] != labels.shape[0]:
+                    print("WARNING: Batch skipped. Image and mask shape does not match:", inputs.shape[0], labels.shape[0])
+                    continue
+
+                optimizer.zero_grad()
+                if not config['AUTOCAST']:
+                    
+                    # segmentation output 
+                    outputs = model(inputs)
+                    if "SegResNetVAE" in config["MODEL_NAME"]: outputs = outputs[0]
+                    if isinstance(outputs, list): outputs = outputs[0]
+
+                    # loss 
+                    if weak_labels is not None: 
+                        weak_label = torch.tensor([weak_labels[step]]).to(config['DEVICE'])
+                    loss = loss_function(outputs, labels, weak_label) if config['LOSS'] == 'wfdice' else loss_function(outputs, labels)
+                    loss.backward()
+                    optimizer.step()
+                    
+                else:
+                    with autocast():
+                        outputs = model(inputs)
+                        if "SegResNetVAE" in config["MODEL_NAME"]: outputs = outputs[0]
+                        if isinstance(outputs, list): outputs = outputs[0]
+                        loss = loss_function(outputs, labels, [weak_labels[step]]) if config['LOSS'] == 'wfdice' else loss_function(outputs, labels)
+                        
+                    scaler.scale(loss).backward()
+                    scaler.unscale_(optimizer)
+                    if torch.isinf(loss).any():
+                        print("Detected inf in gradients.")
+                    else:
+                        scaler.step(optimizer)
+                        scaler.update()
+
+                epoch_loss += loss.item()
+                progress_bar.set_description(f'Epoch [{epoch+1}/{config["MAX_EPOCHS"]}], Loss: {epoch_loss/step:.4f}')
+
+        epoch_loss /= step
+        writer.add_scalar("train_loss_epoch", epoch_loss, epoch)
+        progress_bar.set_description(f'Epoch [{epoch+1}/{config["MAX_EPOCHS"]}], Loss: {epoch_loss:.4f}')
+
+        # validation
+        if (epoch + 1) % config['VAL_INTERVAL'] == 0:
+            
+            # get a list of validation measures, pick one to be the decision maker 
+            val_metrics, (val_images, val_labels, val_outputs) = evaluate(model, val_loader, eval_metrics, config, postprocessing_transforms)
+            if isinstance(config['EVAL_METRIC'], list):
+                cur_metric = np.mean([val_metrics[m] for m in config['EVAL_METRIC']])
+            else:
+                cur_metric = val_metrics[config['EVAL_METRIC']]
+            
+            # determine if better than previous best validation metric 
+            if cur_metric > best_metric:
+                best_metric, best_metric_epoch = cur_metric, epoch + 1
+                torch.save(model.state_dict(), "checkpoints/" + config['EXPORT_FILE_NAME'] + ".pth")
+            
+            # early stopping
+            patience_counter = patience_counter + 1 if prev_metric > cur_metric else 0
+            if patience_counter == patience or epoch - best_metric_epoch > patience:
+                print("Early stopping at epoch", epoch + 1)
+                break
+            print(f'Current epoch: {epoch + 1} current avg {config["EVAL_METRIC"]}: {cur_metric :.4f} best avg {config["EVAL_METRIC"]}: {best_metric:.4f} at epoch {best_metric_epoch}')
+            prev_metric = cur_metric
+            
+            # writer
+            for key, value in val_metrics.items():
+                writer.add_scalar("val_" + key, value, epoch)
+            plot_2d_or_3d_image(val_images, epoch + 1, writer, index=len(val_outputs)//2, tag="image",frame_dim=-1)
+            plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=len(val_outputs)//2, tag="label",frame_dim=-1)
+            plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=len(val_outputs)//2, tag="output",frame_dim=-1)
+
+        # update scheduler
+        try:
+            if scheduler is not None: scheduler.step()
+        except:
+            pass 
+
+    print(f"Train completed, best {config['EVAL_METRIC']}: {best_metric:.4f} at epoch: {best_metric_epoch}")
+    writer.close()
+    return model, writer
+
+
+
+def evaluate(model, val_loader, eval_metrics, config, postprocessing_transforms=None, use_liver_seg=False, export_filenames = [], export_file_metadata = []):
+
+  val_metrics = {}
+  model.eval()
+  with torch.no_grad():
+
+    step = 0
+    for val_data in val_loader:
+        # 3D: val_images has shape (1,C,H,W,Z)
+        # 2D: val_images has shape (B,C,H,W)
+        val_images, val_labels = val_data["image"].to(config['DEVICE']), val_data["mask"].to(config['DEVICE'])
+        if use_liver_seg: val_liver = val_data["pred_liver"].to(config['DEVICE'])
+                
+        if (val_images[0].shape[-1] != val_labels[0].shape[-1]) or (
+            "3D" not in config["MODEL_NAME"] and val_images.shape[0] != val_labels.shape[0]):
+                print("WARNING: Batch skipped. Image and mask shape does not match:", val_images.shape, val_labels.shape)
+                continue
+
+        # convert outputs to probability
+        if "3D" in config["MODEL_NAME"]:
+            val_outputs = sw_inference(model, val_images, config['ROI_SIZE'], config['AUTOCAST'], discard_second_output='SegResNetVAE' in config['MODEL_NAME'])
+        else:
+            if "SegResNetVAE" in config["MODEL_NAME"]: val_outputs, _ = model(val_images)
+            else: val_outputs = model(val_images)
+
+        # post-procesing
+        if postprocessing_transforms is not None:
+            val_outputs = [postprocessing_transforms(i) for i in decollate_batch(val_outputs)]
+
+        # remove tumor predictions outside liver 
+        for i in range(len(val_outputs)):
+            val_outputs[i][-1][torch.where(val_images[i][0] <= 1e-6)] = 0
+  
+        # apply morphological snakes algorithm 
+        if config['POSTPROCESSING_MORF']: 
+            for i in range(len(val_outputs)):
+                val_outputs[i][-1] = torch.from_numpy(ms.morphological_chan_vese(val_images[i][0].cpu(), iterations=2, init_level_set=val_outputs[i][-1].cpu())).to(config['DEVICE'])
+  
+        for i in range(len(val_outputs)):
+            if use_liver_seg: 
+                # use liver model outputs for liver channel 
+                val_outputs[i][1] = val_liver[i]
+                # if region is tumor, assign liver prediction to 0 
+                val_outputs[i][1] -= val_outputs[i][2]
+        
+        # compute metric for current iteration
+        for metric_name, metric in eval_metrics:
+            if isinstance(val_outputs[0], list):
+                val_outputs = val_outputs[0]
+            metric(val_outputs, val_labels)
+        
+        # save prediction to local folder
+        if len(export_filenames) > 0:
+            for _ in range(len(val_outputs)):
+                numpy_array = val_outputs[_].cpu().detach().numpy()
+                write(export_filenames[step], numpy_array[-1], header=export_file_metadata[step])
+                print("   Segmentation exported to", export_filenames[step]) 
+                step += 1
+            
+    # aggregate the final mean metric
+    for metric_name, metric in eval_metrics:
+        if "dice" in metric_name or "IoU" in metric_name: metric_value = metric.aggregate().tolist() 
+        else: metric_value = metric.aggregate()[0].tolist() # a list of accuracies, one per class
+        val_metrics[metric_name + "_avg"] = np.mean(metric_value)
+        if config['TYPE'] != "liver":  
+            for c in range(1, len(metric_value) + 1): # class-wise accuracies
+                val_metrics[metric_name + "_class" + str(c)] = metric_value[c-1]
+        metric.reset()
+
+    return val_metrics, (val_images, val_labels, val_outputs)
+    
+    
+    
+
+def get_defaults(config):
+    
+    if 'TRAIN' not in config.keys(): config['TRAIN'] = True
+    if 'VALID_PATIENT_RATIO' not in config.keys(): config['VALID_PATIENT_RATIO'] = 0.2
+    if 'VAL_INTERVAL' not in config.keys(): config['VAL_INTERVAL'] = 1
+    if 'VAL_INTERVAL' not in config.keys(): config['DROPOUT'] = 0.1
+    if 'EARLY_STOPPING_PATIENCE' not in config.keys(): config['EARLY_STOPPING_PATIENCE'] = 20
+    if 'AUTOCAST' not in config.keys(): config['AUTOCAST'] = False
+    if 'NUM_WORKERS' not in config.keys(): config['NUM_WORKERS'] = 0
+    if 'DROPOUT' not in config.keys(): config['DROPOUT'] = 0.1
+    if 'ONESAMPLETESTRUN' not in config.keys(): config['ONESAMPLETESTRUN'] = False
+    if 'TRAIN' not in config.keys(): config['TRAIN'] = True
+    if 'DATA_AUGMENTATION' not in config.keys(): config['DATA_AUGMENTATION'] = False
+    if 'POSTPROCESSING_MORF' not in config.keys(): config['POSTPROCESSING_MORF'] = False
+    if 'PREPROCESSING' not in config.keys(): config['PREPROCESSING'] = ""
+    if 'PRETRAINED_WEIGHTS' not in config.keys(): config['PRETRAINED_WEIGHTS'] = ""
+    
+    if 'EVAL_INCLUDE_BACKGROUND' not in config.keys(): 
+        if config['TYPE'] == "liver": config['EVAL_INCLUDE_BACKGROUND'] = True
+        else: config['EVAL_INCLUDE_BACKGROUND'] = False
+    if 'EVAL_METRIC' not in config.keys(): 
+        if config['TYPE'] == "liver": config['EVAL_METRIC'] = ["dice_avg"]
+        else: config['EVAL_METRIC'] = ["dice_class2"]
+
+    if 'CLINICAL_DATA_FILE' not in config.keys(): config['CLINICAL_DATA_FILE'] = "Dataset/HCC-TACE-Seg_clinical_data-V2.xlsx"
+    if 'CLINICAL_PREDICTORS' not in config.keys(): config['CLINICAL_PREDICTORS'] = ['T_involvment', 'CLIP_Score','Personal history of cancer', 'TNM', 'Metastasis','fhx_can', 'Alcohol', 'Smoking', 'Evidence_of_cirh', 'AFP', 'age', 'Diabetes', 'Lymphnodes', 'Interval_BL', 'TTP']
+    if 'LAMBDA_WEAK' not in config.keys(): config['LAMBDA_WEAK'] = 0.5
+    if 'MASKNONLIVER' not in config.keys(): config['MASKNONLIVER'] = False
+    
+    if config['TYPE'] == "liver": config['KEEP_CLASSES']=["normal", "liver"]
+    elif config['TYPE'] == "tumor": config['KEEP_CLASSES']=["normal", "liver", "tumor"]
+    else: config['KEEP_CLASSES'] = ["normal", "liver", "tumor", "portal vein", "abdominal aorta"]
+    
+    config['DEVICE'] = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    config['EXPORT_FILE_NAME'] = config['TYPE']+ "_" + config['MODEL_NAME'] + "_" + config['LOSS'] + "_batchsize" + str(config['BATCH_SIZE']) + "_DA" + str(config['DATA_AUGMENTATION']) + "_HU" + str(config['HU_RANGE'][0]) + "-" + str(config['HU_RANGE'][1]) + "_" + config['PREPROCESSING'] + "_" + str(config['ROI_SIZE'][0]) + "_" + str(config['ROI_SIZE'][1]) + "_" + str(config['ROI_SIZE'][2]) + "_dropout" + str(config['DROPOUT'])
+    if config['MASKNONLIVER']: config['EXPORT_FILE_NAME'] += "_wobackground"
+    if config['LOSS'] == "wfdice": config['EXPORT_FILE_NAME'] += "_weaklambda" + str(config['LAMBDA_WEAK'])
+    if config['PRETRAINED_WEIGHTS'] != "" and config['PRETRAINED_WEIGHTS'] != config['EXPORT_FILE_NAME']: config['EXPORT_FILE_NAME'] += "_pretraining"
+    if config['POSTPROCESSING_MORF']: config['EXPORT_FILE_NAME'] += "_wpostmorf"
+    if not config['EVAL_INCLUDE_BACKGROUND']: config['EXPORT_FILE_NAME'] += "_evalnobackground"
+
+    return config 
+    
+
+def train_clinical(df_clinical):
+
+    clinical_model = LinearRegression()
+    
+    # train model 
+    print("Training model using", df_clinical.loc[:, df_clinical.columns != 'tumor_ratio'].shape[1], "features")
+    print(df_clinical.head())
+    clinical_model.fit(df_clinical.loc[:, df_clinical.columns != 'tumor_ratio'], df_clinical['tumor_ratio'])
+    
+    # obtain predicted ratios 
+    pred = clinical_model.predict(df_clinical.loc[:, df_clinical.columns != 'tumor_ratio'])
+    
+    # evaluate
+    corr = np.corrcoef(pred, df_clinical['tumor_ratio'])[0][1]
+    mae = np.mean(np.abs(pred - df_clinical['tumor_ratio']))
+    print(f"The clinical model was fitted. Corr = {corr: .6f}  MAE = {mae: .6f}")
+    
+    return pred 
+    
+    
+def model_pipeline(config=None, plot=True):
+
+    torch.cuda.empty_cache()
+    config = get_defaults(config)
+    print(f"You Are Running on a: {config['DEVICE']}")
+    print("file name:", config['EXPORT_FILE_NAME'])
+
+    writer = SummaryWriter(log_dir="runs/" + config['EXPORT_FILE_NAME'])
+    
+    # prepare data
+    train_loader, valid_loader, test_loader, postprocessing_transforms, df_clinical_train = build_dataset(config, get_clinical=config['LOSS']=="wfdice")
+    
+    # train clinical model
+    if config['LOSS'] == "wfdice": weak_labels = train_clinical(df_clinical_train)
+    else: weak_labels = None
+    
+    # train segmentation model 
+    model = build_model(config)
+    loss_function, optimizer, lr_scheduler, eval_metrics = build_optimizer(model, config)
+    if config['TRAIN']:
+        train(model, train_loader, valid_loader, loss_function, eval_metrics, optimizer, config, lr_scheduler, writer, postprocessing_transforms, weak_labels)
+        model.load_state_dict(torch.load("checkpoints/" + config['EXPORT_FILE_NAME'] + ".pth", map_location=torch.device(config['DEVICE'])))
+    if config['ONESAMPLETESTRUN']:
+        return None, None, None 
+    
+    # test segmentation model 
+    test_metrics, (test_images, test_labels, test_outputs) = evaluate(model, test_loader, eval_metrics, config, postprocessing_transforms)
+    print("Test metrics")
+    for key, value in test_metrics.items():
+      print(f"   {key}: {value:.4f}")
+
+    # visualize
+    if plot:
+        if "3D" in config['MODEL_NAME']:
+          visualize_patient(test_images[0].cpu(), mask=test_labels[0].cpu(), n_slices=9, title="ground truth", z_dim_last="3D" in config['MODEL_NAME'], mask_channel=-1)
+          visualize_patient(test_images[0].cpu(), mask=test_outputs[0].cpu(), n_slices=9, title="predicted", z_dim_last="3D" in config['MODEL_NAME'], mask_channel=-1)
+        else:
+          visualize_patient(test_images.cpu(), mask=test_labels.cpu(), n_slices=9, title="ground truth", z_dim_last="3D" in config['MODEL_NAME'], mask_channel=-1)
+          visualize_patient(test_images.cpu(), mask=torch.stack(test_outputs).cpu(), n_slices=9, title="predicted", z_dim_last="3D" in config['MODEL_NAME'], mask_channel=-1)
+
+    return (test_images, test_labels, test_outputs)

utils/sliding_window.py

@@ -0,0 +1,328 @@
+from collections.abc import Callable, Sequence
+from typing import Any, Iterable
+import numpy as np
+import torch
+import torch.nn.functional as F
+from monai.data.meta_tensor import MetaTensor
+from monai.data.utils import compute_importance_map, dense_patch_slices, get_valid_patch_size
+from monai.inferers.utils import _create_buffered_slices, _compute_coords, _get_scan_interval, _flatten_struct, _pack_struct
+from monai.utils import (
+    BlendMode,
+    PytorchPadMode,
+    convert_data_type,
+    convert_to_dst_type,
+    ensure_tuple,
+    ensure_tuple_rep,
+    fall_back_tuple,
+    look_up_option,
+    optional_import,
+    pytorch_after,
+)
+from tqdm import tqdm
+
+# Adapted from monai 
+def sliding_window_inference(
+    inputs: torch.Tensor | MetaTensor,
+    roi_size: Sequence[int] | int,
+    sw_batch_size: int,
+    predictor: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]],
+    overlap: Sequence[float] | float = 0.25,
+    mode: BlendMode | str = BlendMode.CONSTANT,
+    sigma_scale: Sequence[float] | float = 0.125,
+    padding_mode: PytorchPadMode | str = PytorchPadMode.CONSTANT,
+    cval: float = 0.0,
+    sw_device: torch.device | str | None = None,
+    device: torch.device | str | None = None,
+    progress: bool = False,
+    roi_weight_map: torch.Tensor | None = None,
+    process_fn: Callable | None = None,
+    buffer_steps: int | None = None,
+    buffer_dim: int = -1,
+    with_coord: bool = False,
+    discard_second_output: bool = False,
+    *args: Any,
+    **kwargs: Any,
+) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]:
+    """
+    Sliding window inference on `inputs` with `predictor`.
+
+    The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors.
+    Each output in the tuple or dict value is allowed to have different resolutions with respect to the input.
+    e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes
+    could be ([128,64,256], [64,32,128]).
+    In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still
+    an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters
+    so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension).
+
+    When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
+    To maintain the same spatial sizes, the output image will be cropped to the original input size.
+
+    Args:
+        inputs: input image to be processed (assuming NCHW[D])
+        roi_size: the spatial window size for inferences.
+            When its components have None or non-positives, the corresponding inputs dimension will be used.
+            if the components of the `roi_size` are non-positive values, the transform will use the
+            corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
+            to `(32, 64)` if the second spatial dimension size of img is `64`.
+        sw_batch_size: the batch size to run window slices.
+        predictor: given input tensor ``patch_data`` in shape NCHW[D],
+            The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary
+            with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D'];
+            where H'W'[D'] represents the output patch's spatial size, M is the number of output channels,
+            N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128),
+            the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)).
+            In this case, the parameter `overlap` and `roi_size` need to be carefully chosen
+            to ensure the scaled output ROI sizes are still integers.
+            If the `predictor`'s input and output spatial sizes are different,
+            we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension.
+        overlap: Amount of overlap between scans along each spatial dimension, defaults to ``0.25``.
+        mode: {``"constant"``, ``"gaussian"``}
+            How to blend output of overlapping windows. Defaults to ``"constant"``.
+
+            - ``"constant``": gives equal weight to all predictions.
+            - ``"gaussian``": gives less weight to predictions on edges of windows.
+
+        sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``.
+            Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``.
+            When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding
+            spatial dimensions.
+        padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}
+            Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
+            See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
+        cval: fill value for 'constant' padding mode. Default: 0
+        sw_device: device for the window data.
+            By default the device (and accordingly the memory) of the `inputs` is used.
+            Normally `sw_device` should be consistent with the device where `predictor` is defined.
+        device: device for the stitched output prediction.
+            By default the device (and accordingly the memory) of the `inputs` is used. If for example
+            set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
+            `inputs` and `roi_size`. Output is on the `device`.
+        progress: whether to print a `tqdm` progress bar.
+        roi_weight_map: pre-computed (non-negative) weight map for each ROI.
+            If not given, and ``mode`` is not `constant`, this map will be computed on the fly.
+        process_fn: process inference output and adjust the importance map per window
+        buffer_steps: the number of sliding window iterations along the ``buffer_dim``
+            to be buffered on ``sw_device`` before writing to ``device``.
+            (Typically, ``sw_device`` is ``cuda`` and ``device`` is ``cpu``.)
+            default is None, no buffering. For the buffer dim, when spatial size is divisible by buffer_steps*roi_size,
+            (i.e. no overlapping among the buffers) non_blocking copy may be automatically enabled for efficiency.
+        buffer_dim: the spatial dimension along which the buffers are created.
+            0 indicates the first spatial dimension. Default is -1, the last spatial dimension.
+        with_coord: whether to pass the window coordinates to ``predictor``. Default is False.
+            If True, the signature of ``predictor`` should be ``predictor(patch_data, patch_coord, ...)``.
+        args: optional args to be passed to ``predictor``.
+        kwargs: optional keyword args to be passed to ``predictor``.
+
+    Note:
+        - input must be channel-first and have a batch dim, supports N-D sliding window.
+
+    """
+    buffered = buffer_steps is not None and buffer_steps > 0
+    num_spatial_dims = len(inputs.shape) - 2
+    if buffered:
+        if buffer_dim < -num_spatial_dims or buffer_dim > num_spatial_dims:
+            raise ValueError(f"buffer_dim must be in [{-num_spatial_dims}, {num_spatial_dims}], got {buffer_dim}.")
+        if buffer_dim < 0:
+            buffer_dim += num_spatial_dims
+    overlap = ensure_tuple_rep(overlap, num_spatial_dims)
+    for o in overlap:
+        if o < 0 or o >= 1:
+            raise ValueError(f"overlap must be >= 0 and < 1, got {overlap}.")
+    compute_dtype = inputs.dtype
+
+    # determine image spatial size and batch size
+    # Note: all input images must have the same image size and batch size
+    batch_size, _, *image_size_ = inputs.shape
+    device = device or inputs.device
+    sw_device = sw_device or inputs.device
+
+    temp_meta = None
+    if isinstance(inputs, MetaTensor):
+        temp_meta = MetaTensor([]).copy_meta_from(inputs, copy_attr=False)
+    inputs = convert_data_type(inputs, torch.Tensor, wrap_sequence=True)[0]
+    roi_size = fall_back_tuple(roi_size, image_size_)
+
+    # in case that image size is smaller than roi size
+    image_size = tuple(max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
+    pad_size = []
+    for k in range(len(inputs.shape) - 1, 1, -1):
+        diff = max(roi_size[k - 2] - inputs.shape[k], 0)
+        half = diff // 2
+        pad_size.extend([half, diff - half])
+    if any(pad_size):
+        inputs = F.pad(inputs, pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval)
+
+    # Store all slices
+    scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap)
+    slices = dense_patch_slices(image_size, roi_size, scan_interval, return_slice=not buffered)
+
+    num_win = len(slices)  # number of windows per image
+    total_slices = num_win * batch_size  # total number of windows
+    windows_range: Iterable
+    if not buffered:
+        non_blocking = False
+        windows_range = range(0, total_slices, sw_batch_size)
+    else:
+        slices, n_per_batch, b_slices, windows_range = _create_buffered_slices(
+            slices, batch_size, sw_batch_size, buffer_dim, buffer_steps
+        )
+        non_blocking, _ss = torch.cuda.is_available(), -1
+        for x in b_slices[:n_per_batch]:
+            if x[1] < _ss:  # detect overlapping slices
+                non_blocking = False
+                break
+            _ss = x[2]
+
+    # Create window-level importance map
+    valid_patch_size = get_valid_patch_size(image_size, roi_size)
+    if valid_patch_size == roi_size and (roi_weight_map is not None):
+        importance_map_ = roi_weight_map
+    else:
+        try:
+            valid_p_size = ensure_tuple(valid_patch_size)
+            importance_map_ = compute_importance_map(
+                valid_p_size, mode=mode, sigma_scale=sigma_scale, device=sw_device, dtype=compute_dtype
+            )
+            if len(importance_map_.shape) == num_spatial_dims and not process_fn:
+                importance_map_ = importance_map_[None, None]  # adds batch, channel dimensions
+        except Exception as e:
+            raise RuntimeError(
+                f"patch size {valid_p_size}, mode={mode}, sigma_scale={sigma_scale}, device={device}\n"
+                "Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'."
+            ) from e
+    importance_map_ = convert_data_type(importance_map_, torch.Tensor, device=sw_device, dtype=compute_dtype)[0]
+
+    # stores output and count map
+    output_image_list, count_map_list, sw_device_buffer, b_s, b_i = [], [], [], 0, 0  # type: ignore
+    # for each patch
+    for slice_g in tqdm(windows_range) if progress else windows_range:
+        slice_range = range(slice_g, min(slice_g + sw_batch_size, b_slices[b_s][0] if buffered else total_slices))
+        unravel_slice = [
+            [slice(idx // num_win, idx // num_win + 1), slice(None)] + list(slices[idx % num_win])
+            for idx in slice_range
+        ]
+        if sw_batch_size > 1:
+            win_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
+        else:
+            win_data = inputs[unravel_slice[0]].to(sw_device)
+        if with_coord:
+            seg_prob_out = predictor(win_data, unravel_slice, *args, **kwargs)
+            if discard_second_output and seg_prob_out is not None: seg_prob_out = seg_prob_out[0]
+        else:
+            seg_prob_out = predictor(win_data, *args, **kwargs)
+            if discard_second_output and seg_prob_out is not None: seg_prob_out = seg_prob_out[0]
+
+        # convert seg_prob_out to tuple seg_tuple, this does not allocate new memory.
+        dict_keys, seg_tuple = _flatten_struct(seg_prob_out)
+        if process_fn:
+            seg_tuple, w_t = process_fn(seg_tuple, win_data, importance_map_)
+        else:
+            w_t = importance_map_
+        if len(w_t.shape) == num_spatial_dims:
+            w_t = w_t[None, None]
+        w_t = w_t.to(dtype=compute_dtype, device=sw_device)
+        if buffered:
+            c_start, c_end = b_slices[b_s][1:]
+            if not sw_device_buffer:
+                k = seg_tuple[0].shape[1]  # len(seg_tuple) > 1 is currently ignored
+                sp_size = list(image_size)
+                sp_size[buffer_dim] = c_end - c_start
+                sw_device_buffer = [torch.zeros(size=[1, k, *sp_size], dtype=compute_dtype, device=sw_device)]
+            for p, s in zip(seg_tuple[0], unravel_slice):
+                offset = s[buffer_dim + 2].start - c_start
+                s[buffer_dim + 2] = slice(offset, offset + roi_size[buffer_dim])
+                s[0] = slice(0, 1)
+                sw_device_buffer[0][s] += p * w_t
+            b_i += len(unravel_slice)
+            if b_i < b_slices[b_s][0]:
+                continue
+        else:
+            sw_device_buffer = list(seg_tuple)
+
+        for ss in range(len(sw_device_buffer)):
+            b_shape = sw_device_buffer[ss].shape
+            seg_chns, seg_shape = b_shape[1], b_shape[2:]
+            z_scale = None
+            if not buffered and seg_shape != roi_size:
+                z_scale = [out_w_i / float(in_w_i) for out_w_i, in_w_i in zip(seg_shape, roi_size)]
+                w_t = F.interpolate(w_t, seg_shape, mode=_nearest_mode)
+            if len(output_image_list) <= ss:
+                output_shape = [batch_size, seg_chns]
+                output_shape += [int(_i * _z) for _i, _z in zip(image_size, z_scale)] if z_scale else list(image_size)
+                # allocate memory to store the full output and the count for overlapping parts
+                new_tensor: Callable = torch.empty if non_blocking else torch.zeros  # type: ignore
+                output_image_list.append(new_tensor(output_shape, dtype=compute_dtype, device=device))
+                count_map_list.append(torch.zeros([1, 1] + output_shape[2:], dtype=compute_dtype, device=device))
+                w_t_ = w_t.to(device)
+                for __s in slices:
+                    if z_scale is not None:
+                        __s = tuple(slice(int(_si.start * z_s), int(_si.stop * z_s)) for _si, z_s in zip(__s, z_scale))
+                    count_map_list[-1][(slice(None), slice(None), *__s)] += w_t_
+            if buffered:
+                o_slice = [slice(None)] * len(inputs.shape)
+                o_slice[buffer_dim + 2] = slice(c_start, c_end)
+                img_b = b_s // n_per_batch  # image batch index
+                o_slice[0] = slice(img_b, img_b + 1)
+                if non_blocking:
+                    output_image_list[0][o_slice].copy_(sw_device_buffer[0], non_blocking=non_blocking)
+                else:
+                    output_image_list[0][o_slice] += sw_device_buffer[0].to(device=device)
+            else:
+                sw_device_buffer[ss] *= w_t
+                sw_device_buffer[ss] = sw_device_buffer[ss].to(device)
+                _compute_coords(unravel_slice, z_scale, output_image_list[ss], sw_device_buffer[ss])
+        sw_device_buffer = []
+        if buffered:
+            b_s += 1
+
+    if non_blocking:
+        torch.cuda.current_stream().synchronize()
+
+    # account for any overlapping sections
+    for ss in range(len(output_image_list)):
+        output_image_list[ss] /= count_map_list.pop(0)
+
+    # remove padding if image_size smaller than roi_size
+    if any(pad_size):
+        for ss, output_i in enumerate(output_image_list):
+            zoom_scale = [_shape_d / _roi_size_d for _shape_d, _roi_size_d in zip(output_i.shape[2:], roi_size)]
+            final_slicing: list[slice] = []
+            for sp in range(num_spatial_dims):
+                si = num_spatial_dims - sp - 1
+                slice_dim = slice(
+                    int(round(pad_size[sp * 2] * zoom_scale[si])),
+                    int(round((pad_size[sp * 2] + image_size_[si]) * zoom_scale[si])),
+                )
+                final_slicing.insert(0, slice_dim)
+            output_image_list[ss] = output_i[(slice(None), slice(None), *final_slicing)]
+
+    final_output = _pack_struct(output_image_list, dict_keys)
+    if temp_meta is not None:
+        final_output = convert_to_dst_type(final_output, temp_meta, device=device)[0]
+    else:
+        final_output = convert_to_dst_type(final_output, inputs, device=device)[0]
+
+    return final_output  # type: ignore
+    
+    
+def sw_inference(model, input, roi_size, autocast_on, discard_second_output, overlap=0.8):
+    def _compute(input):
+        return sliding_window_inference(
+            inputs=input,
+            roi_size=roi_size,
+            sw_batch_size=1,
+            predictor=model,
+            overlap=overlap, 
+            progress=False,
+            mode="constant",
+            discard_second_output=discard_second_output
+        )
+
+    if autocast_on:
+        with torch.cuda.amp.autocast():
+            return _compute(input)
+    else:
+        return _compute(input)
+        
+        
+

utils/tumor_features.py

@@ -0,0 +1,55 @@
+from scipy.ndimage import label, find_objects
+import numpy as np
+import cv2
+
+
+IMAGE_SPACING_X = 0.7031
+IMAGE_SPACING_Y = 0.7031
+IMAGE_SPACING_Z = 2.5 
+
+
+
+def compute_largest_diameter(binary_mask):
+    
+    # Label connected components in the binary mask
+    labeled_array, num_features = label(binary_mask)
+
+    # Find the objects (tumors) in the labeled array
+    tumor_objects = find_objects(labeled_array)
+
+    # Initialize the largest diameter variable
+    largest_diameter = 0
+
+    # Iterate through each tumor object
+    for obj in tumor_objects:
+        # Calculate the dimensions of the tumor object
+        z_dim = obj[2].stop - obj[2].start
+        y_dim = obj[1].stop - obj[1].start
+        x_dim = obj[0].stop - obj[0].start
+
+        # Calculate the diameter using the longest dimension
+        diameter = max(z_dim * IMAGE_SPACING_Z, y_dim * IMAGE_SPACING_Y, x_dim * IMAGE_SPACING_X)
+
+        # Update the largest diameter if necessary
+        if diameter > largest_diameter:
+            largest_diameter = diameter
+
+    return largest_diameter / 10 # IN CM
+
+
+
+
+def generate_features(img, liver, tumor):
+    
+    contours, _ = cv2.findContours(mask_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    
+    features = {
+        "lesion size (cm)": compute_largest_diameter(tumor),
+        "lesion shape": "irregular",
+        "lesion density (HU)": np.mean(img[tumor==1]),
+        "involvement of adjacent organs:": "Yes" if np.sum(np.multiply(liver==0, tumor)) > 0 else "No"
+    }
+
+    
+    return features 

utils/visualization.py

@@ -0,0 +1,109 @@
+from matplotlib import pyplot as plt 
+import math 
+import numpy as np
+
+
+
+def visualize_results(img, mask, pred, n_slices: int=3, slices: list=None, title: str=""):
+    """
+    img: tensor [C, H, W, Z]
+    mask: tensor [C, H, W, Z]
+    pred: tensor [C, H, W, Z]
+    n_slices: number of slices to visualize
+    slices: list of slices to visualize
+    title; title of the plot 
+    """
+    if slices is not None:
+      n_slices = len(slices)
+
+    fig, ax = plt.subplots(n_slices, 3, figsize=(14, 5*n_slices))
+    inc = img.shape[-1] // n_slices
+    mask_masked = np.ma.masked_where(mask == 0, mask)
+    pred_masked = np.ma.masked_where(pred == 0, pred)
+
+    for i in range(n_slices):
+        slice_num = i*inc if slices is None else slices[i]
+        
+        # image 
+        for c in range(3):
+          ax[i,c].imshow(img[0,:,:,slice_num], cmap="gray")
+          ax[i,c].axis("off")
+          ax[i,c].set_title(f'image')
+
+        # ground truth 
+        ax[i,1].imshow(mask_masked[1,:,:,slice_num], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.5)
+        ax[i,1].imshow(mask_masked[2,:,:,slice_num], cmap='Reds', vmin=0, vmax=1.3, interpolation='none', alpha=0.8)
+        ax[i,1].set_title(f'ground truth')
+        
+        # predicted 
+        ax[i,2].imshow(pred_masked[1,:,:,slice_num], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.5)
+        ax[i,2].imshow(pred_masked[2,:,:,slice_num], cmap='Reds', vmin=0, vmax=1.3, interpolation='none', alpha=0.8)
+        ax[i,2].set_title(f'predicted')
+
+    plt.suptitle(title, size=14)
+    plt.tight_layout()
+    plt.show()
+    
+
+def visualize_patient(img, mask=None, n_slices: int=3, slices: list=None, z_dim_last=True, mask_channel=0, title: str=""):
+    """
+    img: tensor [C, H, W, Z]
+    mask: tensor [C, H, W, Z]
+    n: number of slices to visualize
+    """
+    if slices is not None:
+      n_slices = len(slices)
+
+    fig, ax = plt.subplots(math.ceil(n_slices/3), 3, figsize=(14, 5*math.ceil(n_slices/3)))
+    if z_dim_last: inc = img.shape[-1] // n_slices
+    else: inc = img.shape[0] // n_slices
+    masked = np.ma.masked_where(mask == 0, mask)
+
+    for i in range(n_slices):
+        r, c = divmod(i, 3)
+        slice_num = i*inc if slices is None else slices[i]
+        if n_slices <= 3:
+            if z_dim_last: ax[c].imshow(img[0,:,:,slice_num], cmap="gray")
+            else: ax[c].imshow(img[slice_num,0,:,:], cmap="gray")
+            ax[c].axis("off")
+            ax[c].set_title(f'slice {slice_num}')
+            if mask is not None:
+                if z_dim_last: mask_overlay = ax[c].imshow(masked[mask_channel,:,:,slice_num], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.4)
+                else: mask_overlay = ax[c].imshow(masked[slice_num,mask_channel,:,:], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.4)
+        else:
+            if z_dim_last: ax[r][c].imshow(img[0,:,:,slice_num], cmap="gray")
+            else: ax[r][c].imshow(img[slice_num,0,:,:], cmap="gray")
+            ax[r][c].axis("off")
+            ax[r][c].set_title(f'slice {slice_num}')
+            if mask is not None:
+                if z_dim_last: mask_overlay = ax[r][c].imshow(masked[mask_channel,:,:,slice_num], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.4)
+                else: mask_overlay = ax[r][c].imshow(masked[slice_num,mask_channel,:,:], cmap='jet', vmin=1, vmax=4, interpolation='none', alpha=0.4)
+
+    plt.suptitle(title, size=14)
+    #if mask is not None:
+    #  cbar = fig.colorbar(mask_overlay, extend='both')
+    plt.tight_layout()
+    plt.show()
+    
+    fig, ax = plt.subplots(math.ceil(n_slices/3), 3, figsize=(14, 5*math.ceil(n_slices/3)))
+    if z_dim_last: inc = img.shape[-1] // n_slices
+    else: inc = img.shape[0] // n_slices
+
+    for i in range(n_slices):
+        r, c = divmod(i, 3)
+        slice_num = i*inc if slices is None else slices[i]
+        if n_slices <= 3:
+            if z_dim_last: ax[c].imshow(img[0,:,:,slice_num], cmap="gray")
+            else: ax[c].imshow(img[slice_num,0,:,:], cmap="gray")
+            ax[c].axis("off")
+            ax[c].set_title(f'slice {slice_num}')
+        else:
+            if z_dim_last: ax[r][c].imshow(img[0,:,:,slice_num], cmap="gray")
+            else: ax[r][c].imshow(img[slice_num,0,:,:], cmap="gray")
+            ax[r][c].axis("off")
+            ax[r][c].set_title(f'slice {slice_num}')
+
+    plt.suptitle(title, size=14)
+
+    plt.tight_layout()
+    plt.show()